From 8b477de71d863953001c0b4da75654f1c3e0951d Mon Sep 17 00:00:00 2001 From: MicDiane <114597733+Micdiane@users.noreply.github.com> Date: Tue, 15 Sep 2026 10:49:03 +0800 Subject: [PATCH] [Performance][MiniMax-M3] Bound sparse decode work and split pages on Blackwell Co-authored-by: OpenAI Codex Signed-off-by: MicDiane <114597733+Micdiane@users.noreply.github.com> --- .../benchmark_minimax_m3_sparse_decode.py | 241 ++++++++++++++++++ .../kernels/check_minimax_m3_decode_model.py | 162 ++++++++++++ .../test_minimax_m3_sparse_attn_fp8_scale.py | 161 ++++++++++++ .../minimax_m3/common/ops/sparse_attn.py | 119 ++++++++- .../common/ops/sparse_decode_paged.py | 240 +++++++++++++++++ 5 files changed, 921 insertions(+), 2 deletions(-) create mode 100644 benchmarks/kernels/benchmark_minimax_m3_sparse_decode.py create mode 100644 benchmarks/kernels/check_minimax_m3_decode_model.py create mode 100644 vllm/models/minimax_m3/common/ops/sparse_decode_paged.py diff --git a/benchmarks/kernels/benchmark_minimax_m3_sparse_decode.py b/benchmarks/kernels/benchmark_minimax_m3_sparse_decode.py new file mode 100644 index 000000000000..59fd063f5202 --- /dev/null +++ b/benchmarks/kernels/benchmark_minimax_m3_sparse_decode.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare complete paged sparse decode calls against an explicitly pinned baseline. + +Extract the base revision's sparse_attn.py with git show and pass its path using +--baseline. Run from a working vLLM CUDA installation; no platform shims are used. +The benchmark checks both implementations against a dense FP32 reference before +alternating CUDA Graph timing rounds. Output is JSON Lines. This measures the +attention operator, not index selection, a model, or serving throughput. +""" + +import argparse +import hashlib +import importlib.util +import json +import statistics +import sys +from functools import partial +from pathlib import Path + +import torch + +from vllm.models.minimax_m3.common.ops import sparse_attn as candidate +from vllm.platforms import current_platform +from vllm.triton_utils import triton + + +def case(batch, heads, length, dql, mode, capacity=None, seed=781): + torch.manual_seed(seed) + pages = triton.cdiv(length, 128) + capacity = capacity or pages + q = torch.randn(batch * dql, heads * 16, 128, device="cuda", dtype=torch.bfloat16) + kv = torch.randn( + batch * pages, heads, 128, 256, device="cuda", dtype=torch.bfloat16 + ) + table = torch.zeros(batch, capacity, device="cuda", dtype=torch.int32) + table[:, :pages] = torch.randperm( + batch * pages, device="cuda", dtype=torch.int32 + ).reshape(batch, pages) + lens = torch.full((batch,), length, device="cuda", dtype=torch.int32) + top = torch.full( + (heads, batch * dql, 16), 2147483647, device="cuda", dtype=torch.int32 + ) + for t in range(batch * dql): + visible = length - dql + t % dql + 1 + n = triton.cdiv(visible, 128) + for h in range(heads): + top[h, t, : min(16, n)] = torch.randperm( + n, device="cuda", dtype=torch.int32 + )[:16] + ks = vs = None + if mode != "bf16": + kv = kv.to(torch.float8_e4m3fn) + if mode == "scalar": + ks = torch.tensor(0.7, device="cuda") + vs = torch.tensor(1.3, device="cuda") + else: + ks = 0.5 + torch.rand(heads, batch * pages * 128, device="cuda") + vs = 0.5 + torch.rand_like(ks) + return dict( + q=q, + kv_cache=kv, + topk_idx=top, + block_table=table, + seq_lens=lens, + num_kv_heads=heads, + sm_scale=128**-0.5, + decode_query_len=dql, + k_scale=ks, + v_scale=vs, + ) + + +def dense(c): + q, kv = c["q"], c["kv_cache"] + out = torch.empty_like(q) + for t in range(q.shape[0]): + req, local = divmod(t, c["decode_query_len"]) + length = int(c["seq_lens"][req]) - c["decode_query_len"] + local + 1 + for h in range(c["num_kv_heads"]): + kk = [] + vv = [] + for logical in c["topk_idx"][ + h, t, : min(16, triton.cdiv(length, 128)) + ].tolist(): + p = int(c["block_table"][req, logical]) + n = min(128, length - logical * 128) + k = kv[p, h, :n, :128].bfloat16() + v = kv[p, h, :n, 128:].bfloat16() + if c["k_scale"] is not None: + ks = c["k_scale"] + vs = c["v_scale"] + if ks.numel() > 1: + ks = ks[h, p * 128 : p * 128 + n, None] + vs = vs[h, p * 128 : p * 128 + n, None] + k = (k.float() * ks).bfloat16() + v = (v.float() * vs).bfloat16() + kk.append(k) + vv.append(v) + k = torch.cat(kk).float() + v = torch.cat(vv).float() + out[t, h * 16 : (h + 1) * 16] = ( + torch.softmax( + q[t, h * 16 : (h + 1) * 16].float() @ k.T * c["sm_scale"], -1 + ) + @ v + ).bfloat16() + return out + + +@torch.inference_mode() +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--baseline", + type=Path, + required=True, + help="Unmodified sparse_attn.py extracted from the base commit", + ) + parser.add_argument( + "--matrix", choices=("development", "heldout"), default="development" + ) + parser.add_argument("--seeds", type=int, nargs="+", default=[781, 997]) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--rep-ms", type=int, default=60) + parser.add_argument("--pad-capacity", type=int, default=0) + args = parser.parse_args() + if args.rounds < 1 or args.rep_ms < 1: + parser.error("rounds and rep-ms must be positive") + spec = importlib.util.spec_from_file_location( + "msa_benchmark_baseline", args.baseline + ) + assert spec is not None and spec.loader is not None + baseline = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = baseline + spec.loader.exec_module(baseline) + print( + json.dumps( + dict( + torch=torch.__version__, + triton=triton.__version__, + gpu=torch.cuda.get_device_name(), + baseline_sha256=hashlib.sha256(args.baseline.read_bytes()).hexdigest(), + matrix=args.matrix, + seeds=args.seeds, + pdl=current_platform.is_arch_support_pdl(), + rounds=args.rounds, + rep_ms=args.rep_ms, + pad_capacity=args.pad_capacity, + ) + ), + flush=True, + ) + shapes = ( + [ + (1, 1, 73, 1), + (4, 1, 127, 1), + (2, 2, 113, 3), + (16, 2, 127, 1), + (33, 1, 127, 1), + (1, 1, 383, 1), + (4, 1, 1025, 1), + (2, 4, 2047, 3), + (8, 4, 8192, 1), + (1, 1, 16385, 1), + ] + if args.matrix == "development" + else [ + (2, 1, 73, 1), + (3, 1, 127, 2), + (2, 4, 113, 1), + (5, 4, 127, 1), + (2, 1, 193, 1), + (3, 1, 449, 1), + (5, 1, 129, 1), + (2, 4, 1025, 2), + (3, 1, 2047, 4), + (7, 4, 4095, 1), + (9, 1, 16385, 2), + (5, 4, 16385, 1), + ] + ) + for seed in args.seeds: + for shape in shapes: + for mode in ["bf16", "scalar", "per_token_head"]: + capacity = max(triton.cdiv(shape[2], 128), args.pad_capacity) + inputs = case(*shape, mode, capacity=capacity, seed=seed) + base_output = torch.empty_like(inputs["q"]) + candidate_output = torch.empty_like(base_output) + + run_base = partial( + baseline.minimax_m3_sparse_attn_decode, **inputs, output=base_output + ) + run_candidate = partial( + candidate.minimax_m3_sparse_attn_decode, + **inputs, + output=candidate_output, + ) + run_base() + run_candidate() + reference = dense(inputs) + torch.testing.assert_close(base_output, reference, atol=0.02, rtol=0.02) + torch.testing.assert_close( + candidate_output, reference, atol=0.02, rtol=0.02 + ) + rel_rms = float( + (candidate_output.float() - reference.float()).norm() + / reference.float().norm().clamp_min(1e-12) + ) + samples = [[], []] + for repeat in range(args.rounds): + for index in [0, 1] if repeat % 2 == 0 else [1, 0]: + samples[index].append( + triton.testing.do_bench_cudagraph( + [run_base, run_candidate][index], rep=args.rep_ms + ) + * 1000 + ) + base_us, candidate_us = [ + statistics.median(values) for values in samples + ] + print( + json.dumps( + dict( + shape=shape, + seed=seed, + mode=mode, + capacity=capacity, + rel_rms=rel_rms, + baseline_us=base_us, + candidate_us=candidate_us, + speedup=base_us / candidate_us, + samples_us=samples, + ) + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/check_minimax_m3_decode_model.py b/benchmarks/kernels/check_minimax_m3_decode_model.py new file mode 100644 index 000000000000..4d4595ce0d86 --- /dev/null +++ b/benchmarks/kernels/check_minimax_m3_decode_model.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check a small dummy-weight MiniMax-M3 execution against a pinned baseline. + +Run once with --baseline-source and once without it, supplying the first JSON +output as --reference-output. This checks model execution and generated-output +parity, not pretrained model quality. Both runs use eager execution and the +Torch sampler; the attention implementation is the only changed backend. +""" + +import argparse +import importlib.util +import json +import os +import sys +import tempfile +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline-source", type=Path) + parser.add_argument("--reference-output", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + os.environ["VLLM_USE_FLASHINFER_SAMPLER"] = "0" + + from vllm import LLM, SamplingParams + from vllm.models.minimax_m3.common import sparse_attention as common_impl + from vllm.models.minimax_m3.common.ops import sparse_attn + from vllm.models.minimax_m3.nvidia import sparse_attention_msa as msa_impl + from vllm.transformers_utils.configs.minimax_m3 import ( + MiniMaxM3Config, + MiniMaxM3TextConfig, + ) + + decode = sparse_attn.minimax_m3_sparse_attn_decode + if args.baseline_source: + spec = importlib.util.spec_from_file_location( + "msa_model_baseline", args.baseline_source + ) + assert spec is not None and spec.loader is not None + baseline = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = baseline + spec.loader.exec_module(baseline) + decode = baseline.minimax_m3_sparse_attn_decode + + calls = [] + + def instrumented(q, kv_cache, topk_idx, block_table, *rest, **kwargs): + calls.append( + dict(query_shape=list(q.shape), page_table_shape=list(block_table.shape)) + ) + return decode(q, kv_cache, topk_idx, block_table, *rest, **kwargs) + + common_impl.minimax_m3_sparse_attn_decode = instrumented + msa_impl.minimax_m3_sparse_attn_decode = instrumented + config = MiniMaxM3TextConfig( + vocab_size=256, + hidden_size=512, + intermediate_size=256, + dense_intermediate_size=1024, + shared_intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=16, + num_key_value_heads=1, + head_dim=128, + max_position_embeddings=1024, + num_local_experts=4, + num_experts_per_tok=2, + num_mtp_modules=0, + moe_layer_freq=[0, 0], + sparse_attention_config=dict( + use_sparse_attention=True, + sparse_index_dim=128, + sparse_num_index_heads=1, + sparse_topk_blocks=16, + sparse_block_size=128, + sparse_disable_index_value=[1, 1], + sparse_score_type="max", + sparse_init_block=0, + sparse_local_block=1, + sparse_attention_freq=[1, 1], + ), + architectures=["MiniMaxM3SparseForCausalLM"], + torch_dtype="bfloat16", + ) + with tempfile.TemporaryDirectory(prefix="minimax-m3-decode-") as model_dir: + MiniMaxM3Config( + text_config=config, architectures=["MiniMaxM3SparseForCausalLM"] + ).save_pretrained(model_dir) + llm = LLM( + model=model_dir, + skip_tokenizer_init=True, + load_format="dummy", + dtype="bfloat16", + max_model_len=1024, + max_num_seqs=4, + max_num_batched_tokens=1024, + gpu_memory_utilization=0.15, + enforce_eager=True, + seed=731, + distributed_executor_backend="uni", + disable_log_stats=True, + block_size=128, + enable_prefix_caching=False, + kv_cache_dtype="auto", + ) + calls.clear() + try: + outputs = llm.generate( + [ + dict(prompt_token_ids=[1 + i % 200 for i in range(length)]) + for length in [31, 129, 385] + ], + SamplingParams( + temperature=0, + max_tokens=8, + ignore_eos=True, + detokenize=False, + logprobs=5, + ), + use_tqdm=False, + ) + finally: + llm.llm_engine.engine_core.shutdown() + + assert len(calls) == 14, "Expected two sparse layers over seven decode steps" + rows = [] + for output in outputs: + sequence = output.outputs[0] + assert sequence.logprobs is not None + rows.append( + dict( + token_ids=list(sequence.token_ids), + logprobs=[ + {str(token): value.logprob for token, value in step.items()} + for step in sequence.logprobs + ], + ) + ) + result = dict(dummy_weights=True, calls=calls, outputs=rows) + if args.reference_output: + reference = json.loads(args.reference_output.read_text()) + assert reference["calls"] == calls + differences = [] + for ref, actual in zip(reference["outputs"], rows, strict=True): + assert ref["token_ids"] == actual["token_ids"] + for left, right in zip(ref["logprobs"], actual["logprobs"], strict=True): + assert left.keys() == right.keys() + differences.extend(abs(left[token] - right[token]) for token in left) + maximum = max(differences) + assert maximum <= 1e-3, maximum + result["max_logprob_difference"] = maximum + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print("Model execution check passed: 24 output tokens, 14 decode calls") + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py b/tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py index 9cd481b5e336..6a35100f0802 100644 --- a/tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py +++ b/tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py @@ -187,3 +187,164 @@ def test_minimax_m3_sparse_decode_fp8_kv_scales(scale_mode: str): torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1) + + +@pytest.mark.parametrize("page_capacity", [1, 4, 17]) +@pytest.mark.parametrize("mode", ["bf16", "scalar", "per_token_head"]) +@pytest.mark.parametrize("decode_query_len", [1, 3]) +@pytest.mark.parametrize("num_reqs,num_kv_heads", [(1, 1), (2, 2), (4, 1)]) +@torch.inference_mode() +def test_page_split_decode_matches_dense( + mode, decode_query_len, num_reqs, num_kv_heads, page_capacity +): + """Check causal tails, permuted physical pages and both FP8 scale layouts.""" + if not current_platform.is_cuda() or torch.cuda.get_device_capability() not in ( + (10, 3), + (12, 0), + ): + pytest.skip("Blackwell single-page path") + if page_capacity > 1 and ( + torch.cuda.get_device_capability() != (12, 0) + or (mode == "bf16" and num_reqs * decode_query_len * num_kv_heads > 16) + ): + pytest.skip("Outside page-split dispatch") + torch.manual_seed(73) + total_q = num_reqs * decode_query_len + heads = num_kv_heads * 16 + q = torch.randn(total_q, heads, 128, device=DEVICE, dtype=DTYPE) + kv = torch.randn( + (num_reqs + 1) * page_capacity, + num_kv_heads, + 128, + 256, + device=DEVICE, + dtype=DTYPE, + ) + if mode != "bf16": + kv = kv.to(torch.float8_e4m3fn) + ks = vs = None + if mode == "scalar": + ks = torch.tensor([0.7], device=DEVICE) + vs = torch.tensor([1.3], device=DEVICE) + elif mode == "per_token_head": + ks = 0.5 + torch.rand( + num_kv_heads, (num_reqs + 1) * page_capacity * 128, device=DEVICE + ) + vs = 0.5 + torch.rand_like(ks) + table = torch.randperm( + num_reqs * page_capacity, device=DEVICE, dtype=torch.int32 + ).reshape(num_reqs, page_capacity) + lens = torch.tensor( + ([1, 63, 127, 128][:num_reqs]), device=DEVICE, dtype=torch.int32 + ) + if page_capacity > 1: + lens += (page_capacity - 1) * 128 + topk = torch.full( + (num_kv_heads, total_q, 16), 2147483647, device=DEVICE, dtype=torch.int32 + ) + for token in range(total_q): + req, local = divmod(token, decode_query_len) + visible = max(int(lens[req]) - decode_query_len + local + 1, 0) + pages = (visible + 127) // 128 + for head in range(num_kv_heads): + topk[head, token, : min(16, pages)] = torch.randperm( + pages, device=DEVICE, dtype=torch.int32 + )[:16] + output = torch.full_like(q, float("nan")) + minimax_m3_sparse_attn_decode( + q, + kv, + topk, + table, + lens, + num_kv_heads, + 128**-0.5, + output, + decode_query_len, + ks, + vs, + ) + reference = torch.zeros_like(q) + for token in range(total_q): + req, local = divmod(token, decode_query_len) + visible = max(int(lens[req]) - decode_query_len + local + 1, 0) + if visible == 0: + continue + for head in range(num_kv_heads): + keys, values = [], [] + for logical in topk[ + head, token, : min(16, (visible + 127) // 128) + ].tolist(): + page = int(table[req, logical]) + count = min(128, visible - logical * 128) + k = kv[page, head, :count, :128].to(DTYPE) + v = kv[page, head, :count, 128:].to(DTYPE) + if ks is not None: + assert vs is not None + k_scale = ( + ks + if ks.numel() == 1 + else ks[head, page * 128 : page * 128 + count, None] + ) + v_scale = ( + vs + if vs.numel() == 1 + else vs[head, page * 128 : page * 128 + count, None] + ) + k = (k.float() * k_scale).to(DTYPE) + v = (v.float() * v_scale).to(DTYPE) + keys.append(k) + values.append(v) + k, v = torch.cat(keys), torch.cat(values) + qq = q[token, head * 16 : (head + 1) * 16].float() + reference[token, head * 16 : (head + 1) * 16] = ( + torch.softmax(qq @ k.float().T * 128**-0.5, dim=-1) @ v.float() + ).to(DTYPE) + assert torch.isfinite(output).all() + torch.testing.assert_close(output, reference, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("page_capacity", [1, 4, 17]) +@pytest.mark.parametrize("fp8", [False, True]) +@torch.inference_mode() +def test_page_split_decode_graph_length_changes(fp8, page_capacity): + """A reused output must become zero for empty rows, including after replay.""" + if not current_platform.is_cuda() or torch.cuda.get_device_capability() not in ( + (10, 3), + (12, 0), + ): + pytest.skip("Blackwell single-page path") + if page_capacity > 1 and torch.cuda.get_device_capability() != (12, 0): + pytest.skip("SM120 page-split path") + q = torch.randn(1, 16, 128, device=DEVICE, dtype=DTYPE) + kv = torch.randn(page_capacity, 1, 128, 256, device=DEVICE, dtype=DTYPE) + if fp8: + kv = kv.to(torch.float8_e4m3fn) + topk = torch.arange(16, device=DEVICE, dtype=torch.int32).reshape(1, 1, 16) + table = torch.arange(page_capacity, device=DEVICE, dtype=torch.int32).reshape( + 1, page_capacity + ) + lens = torch.tensor([128 * page_capacity], device=DEVICE, dtype=torch.int32) + output = torch.empty_like(q) + + def run(): + minimax_m3_sparse_attn_decode(q, kv, topk, table, lens, 1, 128**-0.5, output, 1) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for length in [0, 1, min(129, page_capacity * 128), page_capacity * 128, 0]: + lens.fill_(length) + output.fill_(float("nan")) + graph.replay() + assert torch.isfinite(output).all() + if length == 0: + assert torch.count_nonzero(output) == 0 + elif length == 1: + expected = kv[0, 0, 0, 128:].to(DTYPE).expand_as(output) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + torch.cuda.current_stream().wait_stream(stream) diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index af703c086226..8c900da6466f 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -44,8 +44,9 @@ { "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]), - "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"] - * triton.next_power_of_2(args["gqa_group_size"]), + "BLOCK_SIZE_QH": lambda args: ( + args["BLOCK_SIZE_Q"] * triton.next_power_of_2(args["gqa_group_size"]) + ), } ) @triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"]) @@ -644,6 +645,120 @@ def minimax_m3_sparse_attn_decode( # launch_pdl was specified but unrecognised"). Only pass it when PDL is # actually supported -- on ROCm use_pdl is always False, so it's omitted. pdl_launch = {"launch_pdl": True} if use_pdl else {} + # A single physical page needs no partial-output workspace or merge launch. + if ( + current_platform.is_cuda() + and ( + current_platform.is_device_capability((10, 3)) + or current_platform.is_device_capability((12, 0)) + ) + and block_table.shape[1] == 1 + and 0 < total_q * num_kv_heads <= 32 + and gqa_group_size == 16 + and head_dim == 128 + and max_topk >= 1 + and q.dtype == torch.bfloat16 + and kv_cache.dtype in (torch.bfloat16, torch.float8_e4m3fn) + and q.is_contiguous() + and kv_cache.is_contiguous() + and output.is_contiguous() + and kv_cache.shape[1:] == (num_kv_heads, 128, 256) + ): + from .sparse_decode_paged import _single_page_decode + + _single_page_decode[(total_q * num_kv_heads, 2)]( + q, + kv_cache, + topk_idx, + block_table, + seq_lens, + k_scale_arg, + v_scale_arg, + output, + num_kv_heads, + decode_query_len, + block_table.stride(0), + *topk_idx.stride(), + stride_ks_h, + stride_ks_t, + stride_vs_h, + stride_vs_t, + 1, + 128, + 64, + sm_scale, + kv_scale_mode, + use_fp8, + use_pdl, + num_warps=8 if use_fp8 else 4, + num_stages=1, + **pdl_launch, + ) + return + # Bound work by page-table capacity without reading graph inputs on the CPU. + if ( + current_platform.is_cuda() + and current_platform.is_device_capability((12, 0)) + and gqa_group_size == 16 + and head_dim == 128 + and max_topk == 16 + and block_table.shape[1] > 0 + and total_q > 0 + and q.dtype == torch.bfloat16 + and kv_cache.dtype in (torch.bfloat16, torch.float8_e4m3fn) + and (use_fp8 or total_q * num_kv_heads <= 16) + and q.is_contiguous() + and kv_cache.is_contiguous() + and output.is_contiguous() + and kv_cache.shape[1:] == (num_kv_heads, 128, 256) + ): + from .sparse_decode_paged import _merge_page_splits, _page_split_decode + + groups = total_q * num_kv_heads + tile = 64 if groups <= (32 if use_fp8 else 4) else 128 + splits = min(max_topk, block_table.shape[1]) * (128 // tile) + partial = torch.empty((groups, splits, 16, 128), dtype=q.dtype, device=q.device) + lse = torch.empty((groups, splits, 16), dtype=torch.float32, device=q.device) + _page_split_decode[(groups, splits)]( + q, + kv_cache, + topk_idx, + block_table, + seq_lens, + k_scale_arg, + v_scale_arg, + partial, + lse, + num_kv_heads, + decode_query_len, + block_table.stride(0), + *topk_idx.stride(), + stride_ks_h, + stride_ks_t, + stride_vs_h, + stride_vs_t, + tile, + splits, + sm_scale, + kv_scale_mode, + use_fp8, + use_pdl, + num_warps=4, + num_stages=1, + **pdl_launch, + ) + _merge_page_splits[(groups, 16)]( + partial, + lse, + output, + splits, + triton.next_power_of_2(splits), + 128, + use_pdl, + num_warps=4, + **pdl_launch, + ) + return # split-K over the selected blocks; chunk count is shape-constant (cuda graph). TARGET_GRID = 256 target = max(1, min(max_topk, TARGET_GRID // max(1, total_q * num_kv_heads))) diff --git a/vllm/models/minimax_m3/common/ops/sparse_decode_paged.py b/vllm/models/minimax_m3/common/ops/sparse_decode_paged.py new file mode 100644 index 000000000000..fcd805df920e --- /dev/null +++ b/vllm/models/minimax_m3/common/ops/sparse_decode_paged.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Page-split and direct single-page kernels for MiniMax-M3 decode.""" + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _single_page_decode( + q_ptr, + kv_ptr, + topk_ptr, + block_table_ptr, + seq_lens_ptr, + k_scale_ptr, + v_scale_ptr, + output_ptr, + NUM_KV_HEADS: tl.constexpr, + DECODE_QUERY_LEN: tl.constexpr, + STRIDE_BT_REQ: tl.constexpr, + STRIDE_TOPK_HEAD: tl.constexpr, + STRIDE_TOPK_TOKEN: tl.constexpr, + STRIDE_TOPK_SLOT: tl.constexpr, + STRIDE_K_SCALE_HEAD: tl.constexpr, + STRIDE_K_SCALE_TOKEN: tl.constexpr, + STRIDE_V_SCALE_HEAD: tl.constexpr, + STRIDE_V_SCALE_TOKEN: tl.constexpr, + NUM_SLOTS: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, + BLOCK_OUTPUT_DIM: tl.constexpr, + SM_SCALE: tl.constexpr, + SCALE_MODE: tl.constexpr, + USE_FP8: tl.constexpr, + USE_PDL: tl.constexpr, +): + if USE_PDL: + tl.extra.cuda.gdc_wait() + sh, output_part = tl.program_id(0), tl.program_id(1) + t, h = sh // NUM_KV_HEADS, sh % NUM_KV_HEADS + req, local = t // DECODE_QUERY_LEN, t % DECODE_QUERY_LEN + visible = tl.maximum(tl.load(seq_lens_ptr + req) - DECODE_QUERY_LEN + local + 1, 0) + n, d, g = tl.arange(0, BLOCK_TOKENS), tl.arange(0, 128), tl.arange(0, 16) + od = output_part * BLOCK_OUTPUT_DIM + tl.arange(0, BLOCK_OUTPUT_DIM) + slots, offset = n // 128, n % 128 + valid = (slots < NUM_SLOTS) & (slots < tl.minimum(16, tl.cdiv(visible, 128))) + logical = tl.load( + topk_ptr + + h * STRIDE_TOPK_HEAD + + t * STRIDE_TOPK_TOKEN + + slots * STRIDE_TOPK_SLOT, + valid, + other=0, + ) + physical = tl.load( + block_table_ptr + req * STRIDE_BT_REQ + logical, valid, other=0 + ).to(tl.int64) + mask = valid & (logical * 128 + offset < visible) + base = (physical * NUM_KV_HEADS + h) * 128 * 256 + offset * 256 + q = tl.load(q_ptr + (sh * 16 + g[:, None]) * 128 + d[None, :]) + k = tl.load(kv_ptr + base[None, :] + d[:, None], mask[None, :], other=0.0).to( + q.dtype + ) + if USE_FP8: + if SCALE_MODE == 1: + k = (k * tl.load(k_scale_ptr)).to(q.dtype) + elif SCALE_MODE == 2: + ks = tl.load( + k_scale_ptr + + h * STRIDE_K_SCALE_HEAD + + (physical * 128 + offset) * STRIDE_K_SCALE_TOKEN, + mask, + other=1.0, + ) + k = (k * ks[None, :]).to(q.dtype) + scores = tl.dot(q, k) * (SM_SCALE * 1.4426950408889634) + scores = tl.where(mask[None, :], scores, float("-inf")) + maximum = tl.max(scores, 1) + maximum = tl.where(maximum == float("-inf"), 0.0, maximum) + prob = tl.exp2(scores - maximum[:, None]) + denom = tl.sum(prob, 1) + # Load V after QK so the compiler may reuse K's shared-memory lifetime. + v = tl.load( + kv_ptr + base[:, None] + 128 + od[None, :], mask[:, None], other=0.0 + ).to(q.dtype) + if USE_FP8: + if SCALE_MODE == 1: + v = (v * tl.load(v_scale_ptr)).to(q.dtype) + elif SCALE_MODE == 2: + vs = tl.load( + v_scale_ptr + + h * STRIDE_V_SCALE_HEAD + + (physical * 128 + offset) * STRIDE_V_SCALE_TOKEN, + mask, + other=1.0, + ) + v = (v * vs[:, None]).to(q.dtype) + accum = tl.dot(prob.to(q.dtype), v) + out = accum / tl.where(denom > 0, denom, 1.0)[:, None] + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + tl.store(output_ptr + (sh * 16 + g[:, None]) * 128 + od[None, :], out) + + +@triton.jit +def _page_split_decode( + q_ptr, + kv_ptr, + topk_ptr, + block_table_ptr, + seq_lens_ptr, + k_scale_ptr, + v_scale_ptr, + partial_ptr, + lse_ptr, + NUM_KV_HEADS: tl.constexpr, + DECODE_QUERY_LEN: tl.constexpr, + STRIDE_BT_REQ: tl.constexpr, + STRIDE_TOPK_HEAD: tl.constexpr, + STRIDE_TOPK_TOKEN: tl.constexpr, + STRIDE_TOPK_SLOT: tl.constexpr, + STRIDE_K_SCALE_HEAD: tl.constexpr, + STRIDE_K_SCALE_TOKEN: tl.constexpr, + STRIDE_V_SCALE_HEAD: tl.constexpr, + STRIDE_V_SCALE_TOKEN: tl.constexpr, + TOKENS_PER_SPLIT: tl.constexpr, + NUM_SPLITS: tl.constexpr, + SM_SCALE: tl.constexpr, + SCALE_MODE: tl.constexpr, + USE_FP8: tl.constexpr, + USE_PDL: tl.constexpr, +): + if USE_PDL: + tl.extra.cuda.gdc_wait() + sh = tl.program_id(0) + t, h = sh // NUM_KV_HEADS, sh % NUM_KV_HEADS + piece = tl.program_id(1) + slot, sub = piece // (128 // TOKENS_PER_SPLIT), piece % (128 // TOKENS_PER_SPLIT) + req, local = t // DECODE_QUERY_LEN, t % DECODE_QUERY_LEN + visible = tl.maximum(tl.load(seq_lens_ptr + req) - DECODE_QUERY_LEN + local + 1, 0) + valid = slot < tl.minimum(16, tl.cdiv(visible, 128)) + logical = tl.load( + topk_ptr + + h * STRIDE_TOPK_HEAD + + t * STRIDE_TOPK_TOKEN + + slot * STRIDE_TOPK_SLOT, + valid, + other=0, + ) + active = valid & (logical * 128 + sub * TOKENS_PER_SPLIT < visible) + d, g = tl.arange(0, 128), tl.arange(0, 16) + po = ((sh * NUM_SPLITS + piece) * 16 + g[:, None]) * 128 + d[None, :] + lo = (sh * NUM_SPLITS + piece) * 16 + g + if active: + page = tl.load( + block_table_ptr + req * STRIDE_BT_REQ + logical, valid, other=0 + ).to(tl.int64) + n = sub * TOKENS_PER_SPLIT + tl.arange(0, TOKENS_PER_SPLIT) + mask = valid & (logical * 128 + n < visible) + q = tl.load( + q_ptr + ((t * NUM_KV_HEADS + h) * 16 + g[:, None]) * 128 + d[None, :] + ) + base = (page * NUM_KV_HEADS + h) * 128 * 256 + k = tl.load( + kv_ptr + base + n[None, :] * 256 + d[:, None], mask[None, :], other=0.0 + ).to(q.dtype) + v = tl.load( + kv_ptr + base + n[:, None] * 256 + 128 + d[None, :], + mask[:, None], + other=0.0, + ).to(q.dtype) + if USE_FP8: + if SCALE_MODE == 1: + k = (k * tl.load(k_scale_ptr)).to(q.dtype) + v = (v * tl.load(v_scale_ptr)).to(q.dtype) + elif SCALE_MODE == 2: + ks = tl.load( + k_scale_ptr + + h * STRIDE_K_SCALE_HEAD + + (page * 128 + n) * STRIDE_K_SCALE_TOKEN, + mask, + other=1.0, + ) + vs = tl.load( + v_scale_ptr + + h * STRIDE_V_SCALE_HEAD + + (page * 128 + n) * STRIDE_V_SCALE_TOKEN, + mask, + other=1.0, + ) + k = (k * ks[None, :]).to(q.dtype) + v = (v * vs[:, None]).to(q.dtype) + score = tl.dot(q, k) * (SM_SCALE * 1.4426950408889634) + score = tl.where(mask[None, :], score, float("-inf")) + m = tl.max(score, 1) + safe_m = tl.where(m == float("-inf"), 0.0, m) + prob = tl.exp2(score - safe_m[:, None]) + z = tl.sum(prob, 1) + accum = tl.dot(prob.to(q.dtype), v) + out = accum / tl.where(z > 0, z, 1.0)[:, None] + lse = tl.where(z > 0, safe_m + tl.log2(z), float("-inf")) + tl.store(partial_ptr + po, out) + tl.store(lse_ptr + lo, lse) + else: + # Every slot is overwritten on every replay: no stale scratch on shrink. + tl.store(partial_ptr + po, tl.full((16, 128), 0.0, tl.float32)) + tl.store(lse_ptr + lo, tl.full((16,), float("-inf"), tl.float32)) + + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _merge_page_splits( + partial_ptr, + lse_ptr, + output_ptr, + NUM_SPLITS: tl.constexpr, + PADDED_SPLITS: tl.constexpr, + BLOCK_D: tl.constexpr, + USE_PDL: tl.constexpr, +): + if USE_PDL: + tl.extra.cuda.gdc_wait() + sh, gd = tl.program_id(0), tl.program_id(1) + g, part = gd // (128 // BLOCK_D), gd % (128 // BLOCK_D) + ss = tl.arange(0, PADDED_SPLITS) + d = part * BLOCK_D + tl.arange(0, BLOCK_D) + lo = (sh * NUM_SPLITS + ss) * 16 + g + po = ((sh * NUM_SPLITS + ss[:, None]) * 16 + g) * 128 + d[None, :] + lse_values = tl.load(lse_ptr + lo, ss < NUM_SPLITS, other=float("-inf")) + m = tl.max(lse_values, 0) + m = tl.where(m == float("-inf"), 0.0, m) + w = tl.exp2(lse_values - m) + z = tl.sum(w, 0) + w = w / tl.where(z > 0, z, 1.0) + p = tl.load(partial_ptr + po, (ss < NUM_SPLITS)[:, None], other=0.0).to(tl.float32) + out = tl.sum(p * w[:, None], 0) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + tl.store(output_ptr + (sh * 16 + g) * 128 + d, out)