diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index e1443f49..bb1656ac 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -176,13 +176,22 @@ def bench_fixed(configs): cu_seqlens=cu_seqlens, lower_bound=lower_bound, ) + common_cula = dict(common) + if init_state is not None: + common_cula["init_state"] = init_state.transpose(-1, -2).contiguous() # Accuracy - o_fla, _ = run_fla(**common) - o_cula, _ = run_cula(**common) + o_fla, ht_fla = run_fla(**common) + o_cula, ht_cula_vk = run_cula(**common_cula) torch.cuda.synchronize() relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) + if ht_fla is not None and ht_cula_vk is not None: + ht_cula = ht_cula_vk.transpose(-1, -2) + state_rms, state_max, state_mean = relative_rms_error_rel_max_mean_abs(ht_fla, ht_cula) + relative_rms_error = max(relative_rms_error, state_rms) + rel_max = max(rel_max, state_max) + mean_diff = max(mean_diff, state_mean) # Performance ms_fla = benchmark_cuda_mode_fn( @@ -193,7 +202,7 @@ def bench_fixed(configs): sanitizer_mode=SANITIZER_MODE, ) ms_cula = benchmark_cuda_mode_fn( - lambda: run_cula(**common), + lambda: run_cula(**common_cula), default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE, @@ -216,7 +225,7 @@ def bench_fixed(configs): } ) - del o_fla, o_cula, q, k, v, g, beta, A_log, dt_bias, inputs + del o_fla, o_cula, ht_fla, ht_cula_vk, q, k, v, g, beta, A_log, dt_bias, inputs torch.cuda.empty_cache() return results @@ -267,13 +276,22 @@ def bench_varlen(configs): cu_seqlens=cu_seqlens, lower_bound=lower_bound, ) + common_cula = dict(common) + if init_state is not None: + common_cula["init_state"] = init_state.transpose(-1, -2).contiguous() # Accuracy - o_fla, _ = run_fla(**common) - o_cula, _ = run_cula(**common) + o_fla, ht_fla = run_fla(**common) + o_cula, ht_cula_vk = run_cula(**common_cula) torch.cuda.synchronize() relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) + if ht_fla is not None and ht_cula_vk is not None: + ht_cula = ht_cula_vk.transpose(-1, -2) + state_rms, state_max, state_mean = relative_rms_error_rel_max_mean_abs(ht_fla, ht_cula) + relative_rms_error = max(relative_rms_error, state_rms) + rel_max = max(rel_max, state_max) + mean_diff = max(mean_diff, state_mean) # Performance ms_fla = benchmark_cuda_mode_fn( @@ -284,7 +302,7 @@ def bench_varlen(configs): sanitizer_mode=SANITIZER_MODE, ) ms_cula = benchmark_cuda_mode_fn( - lambda: run_cula(**common), + lambda: run_cula(**common_cula), default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE, @@ -314,7 +332,7 @@ def bench_varlen(configs): } ) - del o_fla, o_cula, q, k, v, g, beta, A_log, dt_bias, inputs + del o_fla, o_cula, ht_fla, ht_cula_vk, q, k, v, g, beta, A_log, dt_bias, inputs torch.cuda.empty_cache() return results diff --git a/benchmarks/bench_qwen35_decode.py b/benchmarks/bench_qwen35_decode.py new file mode 100755 index 00000000..bfbdcd07 --- /dev/null +++ b/benchmarks/bench_qwen35_decode.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +"""Benchmark actual cuLA Qwen GDN decode against SGLang's packed inference path. + +Only config.json is read. State reset is outside both CUDA event windows. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import pathlib +import statistics +import sys +from collections.abc import Callable + +import torch + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +import cula.cudac as cula_cuda + + +def load_shape(config_path: pathlib.Path, tp_size: int) -> dict[str, int | str]: + with config_path.open(encoding="utf-8") as f: + root = json.load(f) + config = root.get("text_config", root) + global_h = int(config["linear_num_key_heads"]) + global_hv = int(config["linear_num_value_heads"]) + if global_h % tp_size or global_hv % tp_size: + raise ValueError(f"TP={tp_size} must divide H={global_h} and HV={global_hv}") + h, hv = global_h // tp_size, global_hv // tp_size + k = int(config["linear_key_head_dim"]) + v = int(config["linear_value_head_dim"]) + if hv % h or k != 128 or v != 128: + raise ValueError(f"unsupported local GVA shape H={h} HV={hv} K={k} V={v}") + return { + "model": config_path.parent.name, + "global_h": global_h, + "global_hv": global_hv, + "h": h, + "hv": hv, + "k": k, + "v": v, + } + + +def load_sglang(sglang_path: pathlib.Path): + for candidate in (sglang_path, sglang_path / "python"): + if candidate.exists(): + sys.path.insert(0, str(candidate)) + from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel + + kernel = TritonGDNKernel() + if not kernel.supports_packed_decode: + raise RuntimeError("SGLang Triton packed GDN decode is unavailable") + return kernel + + +def benchmark_cuda( + fn: Callable[[], object], + *, + setup: Callable[[], None], + warmup: int, + rep: int, +) -> float: + for _ in range(warmup): + setup() + fn() + torch.cuda.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + for start, end in zip(starts, ends, strict=True): + setup() + start.record() + fn() + end.record() + torch.cuda.synchronize() + samples = sorted(start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)) + if len(samples) < 4: + return statistics.mean(samples) + return statistics.mean(samples[len(samples) // 4 : 3 * len(samples) // 4]) + + +def relative_rms(reference: torch.Tensor, actual: torch.Tensor) -> float: + ref = reference.float() + diff = ref - actual.float() + return (diff.square().mean().sqrt() / ref.square().mean().sqrt().clamp_min(1e-8)).item() + + +def make_inputs(tokens: int, shape: dict[str, int | str], seed: int) -> dict[str, torch.Tensor]: + torch.manual_seed(seed) + device = torch.device("cuda") + h, hv, k, v = (int(shape[name]) for name in ("h", "hv", "k", "v")) + conv_dim = 2 * h * k + hv * v + state_kv = torch.randn(tokens, hv, k, v, device=device, dtype=torch.float32) * 0.01 + return { + "mixed_qkv": torch.randn(tokens, conv_dim, device=device, dtype=torch.bfloat16), + "a": torch.randn(tokens, hv, device=device, dtype=torch.bfloat16), + "b": torch.randn(tokens, hv, device=device, dtype=torch.bfloat16), + "A_log": -torch.rand(hv, device=device, dtype=torch.float32), + "dt_bias": torch.randn(hv, device=device, dtype=torch.float32) * 0.1, + "state_kv": state_kv, + "state_vk": state_kv.transpose(-1, -2).contiguous(), + "indices": torch.arange(tokens, device=device, dtype=torch.int32), + } + + +@torch.inference_mode() +def run_case(tokens: int, shape, sglang_kernel, args) -> dict[str, float | int]: + x = make_inputs(tokens, shape, args.seed) + hv, k, v = (int(shape[name]) for name in ("hv", "k", "v")) + state_cula = torch.empty_like(x["state_kv"]) + state_sglang = torch.empty_like(x["state_vk"]) + out_cula = torch.empty(tokens, hv, v, device="cuda", dtype=torch.bfloat16) + + def setup_cula(): + state_cula.copy_(x["state_kv"]) + + def setup_sglang(): + state_sglang.copy_(x["state_vk"]) + + def run_cula(): + cula_cuda.qwen35_layout_scalar_kda_decode( + x["mixed_qkv"], x["a"], x["b"], x["A_log"], x["dt_bias"], + state_cula, x["indices"], out_cula, + ) + + def run_sglang(): + return sglang_kernel.packed_decode( + mixed_qkv=x["mixed_qkv"], a=x["a"], b=x["b"], + A_log=x["A_log"], dt_bias=x["dt_bias"], scale=k**-0.5, + ssm_states=state_sglang, cache_indices=x["indices"], + num_v_heads=hv, head_v_dim=v, + ) + + setup_cula() + run_cula() + setup_sglang() + out_sglang = run_sglang().squeeze(0) + torch.cuda.synchronize() + out_rrms = relative_rms(out_sglang, out_cula) + state_rrms = relative_rms(state_sglang, state_cula.transpose(-1, -2)) + + sglang_ms = benchmark_cuda(run_sglang, setup=setup_sglang, warmup=args.warmup, rep=args.rep) + cula_ms = benchmark_cuda(run_cula, setup=setup_cula, warmup=args.warmup, rep=args.rep) + return { + "tokens": tokens, + "sglang_packed_ms": sglang_ms, + "cula_fused_ms": cula_ms, + "speedup": sglang_ms / cula_ms, + "out_rel_rms": out_rrms, + "state_rel_rms": state_rrms, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--sglang-path", type=pathlib.Path, default=pathlib.Path("/sgl-workspace/sglang")) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--tokens", type=int, nargs="+", default=(1, 2, 4, 8, 16, 32, 64, 128)) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--csv", type=pathlib.Path) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + shape = load_shape(args.config_json, args.tp_size) + sglang_kernel = load_sglang(args.sglang_path) + print( + f"model={shape['model']} device={torch.cuda.get_device_name(0)} TP={args.tp_size} " + f"global_H/HV={shape['global_h']}/{shape['global_hv']} " + f"local_H/HV={shape['h']}/{shape['hv']} K/V={shape['k']}/{shape['v']}" + ) + print("state reset is outside timing; SGLang packed decode vs cuLA fused packed decode") + print("| tokens | sglang_packed_ms | cula_fused_ms | speedup | out_rrms | state_rrms |") + print("|---:|---:|---:|---:|---:|---:|") + rows = [] + for tokens in args.tokens: + row = run_case(tokens, shape, sglang_kernel, args) + rows.append(row) + print( + f"| {tokens} | {row['sglang_packed_ms']:.4f} | {row['cula_fused_ms']:.4f} | " + f"{row['speedup']:.3f}x | {row['out_rel_rms']:.3e} | {row['state_rel_rms']:.3e} |" + ) + if args.csv: + args.csv.parent.mkdir(parents=True, exist_ok=True) + with args.csv.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_qwen35_prefill.py b/benchmarks/bench_qwen35_prefill.py new file mode 100644 index 00000000..34cb989f --- /dev/null +++ b/benchmarks/bench_qwen35_prefill.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +"""Benchmark cuLA native-GVA Qwen3.5 prefill against SGLang's inference path. + +Only config.json is read. Model weights are not loaded: tensors are generated +from the Qwen3.5 linear-attention shapes and dtype declared by the config. +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import statistics +import sys +from collections.abc import Callable + +import torch + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +from cula.ops.qwen35_fused_kda_prefill import qwen35_fused_kda_prefill + + +def load_qwen35_shape(config_path: pathlib.Path, tp_size: int) -> dict[str, object]: + with config_path.open(encoding="utf-8") as f: + root = json.load(f) + config = root.get("text_config", root) + + required = ( + "linear_num_key_heads", + "linear_num_value_heads", + "linear_key_head_dim", + "linear_value_head_dim", + ) + missing = [key for key in required if key not in config] + if missing: + raise ValueError(f"{config_path} is missing Qwen3.5 fields: {missing}") + + global_h = int(config["linear_num_key_heads"]) + global_hv = int(config["linear_num_value_heads"]) + if global_h % tp_size or global_hv % tp_size: + raise ValueError(f"TP={tp_size} must divide H={global_h} and HV={global_hv}") + + h = global_h // tp_size + hv = global_hv // tp_size + k = int(config["linear_key_head_dim"]) + v = int(config["linear_value_head_dim"]) + if hv % h: + raise ValueError(f"Qwen3.5 GVA requires local HV % H == 0, got H={h} HV={hv}") + if k != 128 or v != 128: + raise ValueError(f"cuLA native-GVA prefill currently requires K=V=128, got K={k} V={v}") + + dtype_name = str( + config.get("torch_dtype", config.get("dtype", root.get("torch_dtype", root.get("dtype", "bfloat16")))) + ).lower() + if dtype_name not in ("bfloat16", "bf16", "torch.bfloat16"): + raise ValueError(f"This benchmark expects Qwen3.5 bf16 activations, got torch_dtype={dtype_name}") + + return { + "model_type": config.get("model_type", root.get("model_type", "unknown")), + "global_h": global_h, + "global_hv": global_hv, + "h": h, + "hv": hv, + "k": k, + "v": v, + "dtype": torch.bfloat16, + } + + +def load_sglang(sglang_path: pathlib.Path | None): + if sglang_path is not None: + for candidate in (sglang_path, sglang_path / "python"): + if candidate.exists(): + sys.path.insert(0, str(candidate)) + + from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating + from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel + + return fused_gdn_gating, TritonGDNKernel() + + +def benchmark_cuda( + fn: Callable[[], object], + warmup: int, + rep: int, + setup: Callable[[], None] | None = None, +) -> float: + for _ in range(warmup): + if setup is not None: + setup() + fn() + torch.cuda.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + for start, end in zip(starts, ends, strict=True): + if setup is not None: + setup() + start.record() + fn() + end.record() + torch.cuda.synchronize() + + samples = sorted(start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)) + if len(samples) < 4: + return statistics.mean(samples) + return statistics.mean(samples[len(samples) // 4 : 3 * len(samples) // 4]) + + +def relative_rms(ref: torch.Tensor, out: torch.Tensor) -> float: + ref_f = ref.float() + diff_rms = (ref_f - out.float()).square().mean().sqrt() + return (diff_rms / ref_f.square().mean().sqrt().clamp_min(1.0e-8)).item() + + +def make_inputs( + batch: int, + seq_len: int, + shape: dict[str, object], + device: torch.device, + seed: int, + random_initial_state: bool, +) -> dict[str, torch.Tensor]: + torch.manual_seed(seed) + total = batch * seq_len + h, hv, k, v = (int(shape[name]) for name in ("h", "hv", "k", "v")) + dtype = shape["dtype"] + + state = torch.zeros(batch, hv, k, v, device=device, dtype=torch.float32) + if random_initial_state: + state.normal_(mean=0.0, std=0.01) + + return { + "q": torch.randn(1, total, h, k, device=device, dtype=dtype), + "k": torch.randn(1, total, h, k, device=device, dtype=dtype), + "v": torch.randn(1, total, hv, v, device=device, dtype=dtype), + "a": torch.randn(total, hv, device=device, dtype=dtype), + "b": torch.randn(total, hv, device=device, dtype=dtype), + "A_log": -torch.rand(hv, device=device, dtype=torch.float32), + "dt_bias": torch.randn(hv, device=device, dtype=torch.float32) * 0.1, + "state_kv": state, + "state_vk": state.transpose(-1, -2).contiguous(), + "cu_seqlens": torch.arange(0, total + 1, seq_len, device=device, dtype=torch.int32), + "cache_indices": torch.arange(batch, device=device, dtype=torch.int32), + } + + +@torch.inference_mode() +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--sglang-path", type=pathlib.Path, default=pathlib.Path("/sgl-workspace/sglang")) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-lens", type=int, nargs="+", default=(128, 256, 512, 1024, 2048, 4096)) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rep", type=int, default=30) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--random-initial-state", action="store_true") + parser.add_argument("--skip-accuracy", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + + shape = load_qwen35_shape(args.config_json, args.tp_size) + fused_gdn_gating, sglang_kernel = load_sglang(args.sglang_path) + device = torch.device("cuda") + + print("Qwen3.5 GDN prefill: cuLA native GVA vs SGLang Triton inference path") + print(f"config={args.config_json} model_type={shape['model_type']}") + print( + f"device={torch.cuda.get_device_name(device)} TP={args.tp_size} " + f"global H/HV={shape['global_h']}/{shape['global_hv']} " + f"local H/HV={shape['h']}/{shape['hv']} K/V={shape['k']}/{shape['v']}" + ) + print(f"batch={args.batch} warmup={args.warmup} rep={args.rep}") + print() + print(f"{'T/seq':>8} {'tokens':>8} {'SGLang ms':>11} {'cuLA ms':>10} {'speedup':>9} {'out rrms':>11} {'state rrms':>12}") + print("-" * 79) + + for seq_len in args.seq_lens: + x = make_inputs( + args.batch, + seq_len, + shape, + device, + args.seed, + args.random_initial_state, + ) + state_sglang = torch.empty_like(x["state_vk"]) + + def setup_sglang(): + state_sglang.copy_(x["state_vk"]) + + def run_cula(): + return qwen35_fused_kda_prefill( + x["q"], + x["k"], + x["v"], + x["a"], + x["b"], + x["A_log"], + x["dt_bias"], + initial_state=x["state_kv"], + cu_seqlens=x["cu_seqlens"], + output_final_state=True, + ) + + def run_sglang(): + g, beta = fused_gdn_gating( + x["A_log"], + x["a"], + x["b"], + x["dt_bias"], + ) + return sglang_kernel.extend( + x["q"], + x["k"], + x["v"], + g, + beta, + ssm_states=state_sglang, + cache_indices=x["cache_indices"], + query_start_loc=x["cu_seqlens"], + ) + + rrms = float("nan") + state_rrms = float("nan") + if not args.skip_accuracy: + out_cula, state_cula = run_cula() + setup_sglang() + out_sglang = run_sglang()[0] + torch.cuda.synchronize() + rrms = relative_rms(out_sglang, out_cula) + state_rrms = relative_rms(state_sglang, state_cula.transpose(-1, -2)) + + sglang_ms = benchmark_cuda(run_sglang, args.warmup, args.rep, setup=setup_sglang) + cula_ms = benchmark_cuda(run_cula, args.warmup, args.rep) + total = args.batch * seq_len + print( + f"{seq_len:8d} {total:8d} {sglang_ms:11.4f} {cula_ms:10.4f} " + f"{sglang_ms / cula_ms:8.3f}x {rrms:11.3e} {state_rrms:12.3e}" + ) + del x + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_qwen35_scalar_prefill_core.py b/benchmarks/bench_qwen35_scalar_prefill_core.py new file mode 100644 index 00000000..7b55fe27 --- /dev/null +++ b/benchmarks/bench_qwen35_scalar_prefill_core.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 +"""Benchmark the optimized legacy Qwen scalar prefill CUDA kernel. + +The reference path is the real SGLang ``TritonGDNKernel.extend`` path. Both +implementations receive compact native-GVA Q/K and full-HV V/gates. The table +is deliberately a compute-kernel comparison: SGLang's +``fused_gdn_gating`` is warmed and evaluated before timing, while cuLA's +legacy kernel includes its raw ``a/b`` gate conversion, making the comparison +conservative for cuLA. +""" + +from __future__ import annotations + +import argparse +import csv +import importlib.metadata +import json +import pathlib +import statistics +import subprocess +import sys +from typing import Callable + +import torch + +ROOT = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill, qwen35_scalar_kda_prefill_core + + +def _run_text(command: list[str], *, cwd: pathlib.Path) -> str | None: + try: + result = subprocess.run( + command, + cwd=cwd, + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() + + +def _git_head(path: pathlib.Path) -> str: + return _run_text(["git", "rev-parse", "HEAD"], cwd=path) or "unavailable" + + +def _tracked_source_state(path: pathlib.Path) -> str: + status = _run_text( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=path, + ) + if status is None: + return "unavailable" + return "dirty" if status else "clean" + + +def load_shape(path: pathlib.Path, tp: int) -> dict[str, int | torch.dtype | str]: + root = json.loads(path.read_text(encoding="utf-8")) + cfg = root.get("text_config", root) + h_global = int(cfg["linear_num_key_heads"]) + hv_global = int(cfg["linear_num_value_heads"]) + if h_global % tp or hv_global % tp: + raise ValueError(f"TP={tp} must divide global H/HV={h_global}/{hv_global}") + h, hv = h_global // tp, hv_global // tp + if int(cfg["linear_key_head_dim"]) != 128 or int(cfg["linear_value_head_dim"]) != 128: + raise ValueError("This scalar benchmark requires K=V=128") + if hv % h: + raise ValueError(f"local HV={hv} must be divisible by local H={h}") + return { + "model_type": str(cfg.get("model_type", root.get("model_type", "unknown"))), + "h_global": h_global, + "hv_global": hv_global, + "h": h, + "hv": hv, + "dtype": torch.bfloat16, + } + + +def _timed(fn: Callable[[], object], repeats: int) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) + + +def _capture_cuda_graph(fn: Callable[[], object]) -> Callable[[], object]: + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + torch.cuda.synchronize() + return graph.replay + + +def _rrms(a: torch.Tensor, b: torch.Tensor) -> float: + af, bf = a.float(), b.float() + return ((af - bf).square().mean().sqrt() / af.square().mean().sqrt().clamp_min(1.0e-8)).item() + + +@torch.inference_mode() +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--sglang-path", type=pathlib.Path, default=pathlib.Path("/sgl-workspace/sglang")) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-lens", type=int, nargs="+", default=(1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096)) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--inner", type=int, default=1, help="kernel calls per CUDA event sample") + parser.add_argument( + "--preheat-iters", + type=int, + default=0, + help="large BF16 GEMMs before timing to stabilize GPU clocks", + ) + parser.add_argument( + "--eager-timing", + action="store_true", + help="time Python launches directly instead of CUDA Graph replay", + ) + parser.add_argument("--random-initial-state", action="store_true") + parser.add_argument("--skip-accuracy", action="store_true") + parser.add_argument( + "--core-only", + action="store_true", + help="compare the preprocessed g/beta calculation core; exclude raw a/b gate conversion on both sides", + ) + parser.add_argument("--csv", type=pathlib.Path, help="write the exact per-shape medians, IQRs, and errors") + parser.add_argument( + "--min-speedup", + type=float, + help="fail unless every acceptance shape reaches this paired-median speedup; requires --core-only", + ) + parser.add_argument( + "--acceptance-seq-lens", + type=int, + nargs="+", + default=(256, 512), + help="sequence lengths checked by --min-speedup", + ) + parser.add_argument( + "--require-clean-source", + action="store_true", + help="fail when tracked source changes are present; ignored build artifacts are allowed", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if args.batch != 1: + raise ValueError("The packed SGLang comparison currently requires --batch 1") + if args.min_speedup is not None and not args.core_only: + parser.error("--min-speedup is an acceptance gate and requires the apples-to-apples --core-only scope") + missing_acceptance_shapes = sorted(set(args.acceptance_seq_lens) - set(args.seq_lens)) + if args.min_speedup is not None and missing_acceptance_shapes: + parser.error(f"acceptance sequence lengths are missing from --seq-lens: {missing_acceptance_shapes}") + source_state = _tracked_source_state(ROOT) + if args.require_clean_source and source_state != "clean": + parser.error(f"formal runs require clean tracked source, got source_state={source_state}") + shape = load_shape(args.config_json, args.tp_size) + import cula.cudac as cula_cuda + + for candidate in (args.sglang_path, args.sglang_path / "python"): + if candidate.exists(): + sys.path.insert(0, str(candidate)) + + from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating + from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel + + device = torch.device("cuda") + sg_kernel = TritonGDNKernel() + if not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core"): + raise RuntimeError("the loaded cuLA extension does not expose qwen35_scalar_kda_prefill_core") + core_op = cula_cuda.qwen35_scalar_kda_prefill_core + extension_module = sys.modules.get(getattr(core_op, "__module__", "")) + extension_path = getattr(extension_module, "__file__", "unavailable") + try: + sglang_version = importlib.metadata.version("sglang") + except importlib.metadata.PackageNotFoundError: + sglang_version = "unavailable" + repo_head = _git_head(ROOT) + sglang_head = _git_head(args.sglang_path) + h, hv = int(shape["h"]), int(shape["hv"]) + print( + f"repo_head={repo_head} tracked_source_state={source_state} " + f"extension={extension_path}" + ) + print( + f"torch={torch.__version__} cuda={torch.version.cuda} " + f"sglang={sglang_version} sglang_head={sglang_head}" + ) + print( + f"device={torch.cuda.get_device_name(device)} config={args.config_json} " + f"model_type={shape['model_type']} TP={args.tp_size} " + f"global H/HV={shape['h_global']}/{shape['hv_global']} local H/HV={h}/{hv}" + ) + print( + f"batch={args.batch} warmup={args.warmup} rep={args.rep} inner={args.inner} " + f"graph={'off' if args.eager_timing else 'on'} " + f"(scope={'preprocessed core' if args.core_only else 'raw CULA gate vs SGLang core'})" + ) + if args.preheat_iters: + heat_a = torch.randn(8192, 8192, device=device, dtype=torch.bfloat16) + heat_b = torch.randn_like(heat_a) + for _ in range(args.preheat_iters): + torch.mm(heat_a, heat_b) + torch.cuda.synchronize() + del heat_a, heat_b + print(f"{'T':>6} {'SGLang ms':>25} {'cuLA ms':>25} {'speedup':>10} {'out rrms':>12} {'state rrms':>12}") + print("-" * 110) + rows: list[dict[str, int | float | str]] = [] + + for seq_len in args.seq_lens: + torch.manual_seed(7000 + seq_len + hv * 17 + args.tp_size) + total = args.batch * seq_len + q = torch.randn(args.batch, seq_len, h, 128, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(args.batch, seq_len, hv, 128, device=device, dtype=torch.bfloat16) + a = torch.randn(args.batch, seq_len, hv, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(hv, device=device, dtype=torch.float32) + dt_bias = torch.randn(hv, device=device, dtype=torch.float32) * 0.1 + state_kv = torch.zeros(args.batch, hv, 128, 128, device=device, dtype=torch.float32) + if args.random_initial_state: + state_kv.normal_(mean=0.0, std=0.01) + state_vk = state_kv.transpose(-1, -2).contiguous() + state_sg = state_vk.clone() + cu = torch.arange(0, total + 1, seq_len, device=device, dtype=torch.int32) + cache_indices = torch.arange(args.batch, device=device, dtype=torch.int32) + + # Keep gating outside both timed calls. This is the same input format + # that SGLang's gdn_backend passes to TritonGDNKernel.extend. + g, beta = fused_gdn_gating(A_log, a.reshape(total, hv), b.reshape(total, hv), dt_bias) + g_core = g.reshape(args.batch, seq_len, hv).contiguous() + beta_core = beta.reshape(args.batch, seq_len, hv).contiguous() + + out_cula = torch.empty_like(v) + state_cula = torch.empty_like(state_kv) + empty_initial = torch.empty(0, device=device, dtype=torch.float32) + + def run_cula() -> None: + # Call the extension ABI directly: output/state allocation and + # Python wrapper overhead are not part of the compute measurement. + cula_state = state_kv + if args.core_only: + cula_cuda.qwen35_scalar_kda_prefill_core( + q, k, v, g_core, beta_core, cula_state, cu, out_cula, state_cula + ) + else: + cula_cuda.qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, cula_state, cu, out_cula, state_cula + ) + + def reset_sg() -> None: + state_sg.copy_(state_vk) + + def run_sg() -> None: + sg_kernel.extend( + q, + k, + v, + g, + beta, + ssm_states=state_sg, + cache_indices=cache_indices, + query_start_loc=cu, + ) + + # Warm up each path independently, including the first Triton compile. + for _ in range(args.warmup): + run_cula() + reset_sg() + run_sg() + torch.cuda.synchronize() + + out_rrms = float("nan") + state_rrms = float("nan") + if not args.skip_accuracy: + run_cula() + reset_sg() + run_sg() + torch.cuda.synchronize() + # SGLang mutates [V,K], while cuLA writes [K,V]. + reset_sg() + out_sg = sg_kernel.extend( + q, + k, + v, + g, + beta, + ssm_states=state_sg, + cache_indices=cache_indices, + query_start_loc=cu, + )[0] + torch.cuda.synchronize() + out_rrms = _rrms(out_sg, out_cula) + state_rrms = _rrms(state_sg, state_cula.transpose(-1, -2)) + + timed_sg: Callable[[], object] = run_sg + timed_cula: Callable[[], object] = run_cula + if not args.eager_timing: + reset_sg() + timed_sg = _capture_cuda_graph(run_sg) + timed_cula = _capture_cuda_graph(run_cula) + + sg_ms: list[float] = [] + cu_ms: list[float] = [] + # Alternate order and restore the state outside each CUDA event. This + # avoids a systematic clock/thermal bias between the two kernels. + for i in range(args.rep): + if i & 1: + reset_sg() + cu_ms.append(_timed(timed_cula, args.inner) / args.inner) + reset_sg() + sg_ms.append(_timed(timed_sg, args.inner) / args.inner) + else: + reset_sg() + sg_ms.append(_timed(timed_sg, args.inner) / args.inner) + cu_ms.append(_timed(timed_cula, args.inner) / args.inner) + + def middle(xs: list[float]) -> tuple[float, float, float]: + ys = sorted(xs) + return statistics.median(ys), ys[len(ys) // 4], ys[(3 * len(ys)) // 4] + + sg_med, sg_q1, sg_q3 = middle(sg_ms) + cu_med, cu_q1, cu_q3 = middle(cu_ms) + paired = [s / c for s, c in zip(sg_ms, cu_ms)] + speed_med, speed_q1, speed_q3 = middle(paired) + rows.append( + { + "repo_head": repo_head, + "source_state": source_state, + "extension": extension_path, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda or "unavailable", + "sglang_version": sglang_version, + "sglang_head": sglang_head, + "config_json": str(args.config_json.resolve()), + "scope": "core" if args.core_only else "raw_cula_vs_sglang_core", + "cuda_graph": not args.eager_timing, + "random_initial_state": args.random_initial_state, + "warmup": args.warmup, + "rep": args.rep, + "inner": args.inner, + "seq_len": seq_len, + "batch": args.batch, + "tp_size": args.tp_size, + "qk_heads": h, + "v_heads": hv, + "sglang_ms": sg_med, + "sglang_q1_ms": sg_q1, + "sglang_q3_ms": sg_q3, + "cula_ms": cu_med, + "cula_q1_ms": cu_q1, + "cula_q3_ms": cu_q3, + "paired_speedup": speed_med, + "paired_speedup_q1": speed_q1, + "paired_speedup_q3": speed_q3, + "out_rrms": out_rrms, + "state_rrms": state_rrms, + } + ) + print( + f"{seq_len:6d} {sg_med:8.4f} [{sg_q1:8.4f},{sg_q3:8.4f}] " + f"{cu_med:8.4f} [{cu_q1:8.4f},{cu_q3:8.4f}] " + f"{speed_med:8.3f}x [{speed_q1:6.3f},{speed_q3:6.3f}] " + f"{out_rrms:12.3e} {state_rrms:12.3e}" + ) + + if args.csv: + args.csv.parent.mkdir(parents=True, exist_ok=True) + with args.csv.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + if args.min_speedup is not None: + acceptance = {int(row["seq_len"]): float(row["paired_speedup"]) for row in rows} + failures = { + seq_len: acceptance[seq_len] + for seq_len in args.acceptance_seq_lens + if acceptance[seq_len] < args.min_speedup + } + if failures: + formatted = ", ".join(f"T{seq_len}={speedup:.3f}x" for seq_len, speedup in failures.items()) + raise SystemExit(f"speedup acceptance failed (required {args.min_speedup:.3f}x): {formatted}") + print( + "speedup acceptance passed: " + + ", ".join( + f"T{seq_len}={acceptance[seq_len]:.3f}x" for seq_len in args.acceptance_seq_lens + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/profile_qwen35_decode.py b/benchmarks/profile_qwen35_decode.py new file mode 100644 index 00000000..b1993f5f --- /dev/null +++ b/benchmarks/profile_qwen35_decode.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Small Nsight Compute target for Qwen3.5 decode kernels. + +Use with `ncu --profile-from-start off` so only the decode loop bracketed by +cudaProfilerStart/Stop is collected. +""" + +from __future__ import annotations + +import argparse +import pathlib +import sys + +import torch + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +import cula.cudac as cula_cuda +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as GLOBAL_CONFIG +from cula.qwen35.common import Qwen35LinearAttentionConfig + + +def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: + if tp_size not in (1, 2, 4, 8): + raise ValueError(f"tp_size must be one of 1, 2, 4, 8, got {tp_size}") + return Qwen35LinearAttentionConfig( + hidden_size=GLOBAL_CONFIG.hidden_size // tp_size, + conv_kernel_size=GLOBAL_CONFIG.conv_kernel_size, + num_k_heads=GLOBAL_CONFIG.num_k_heads // tp_size, + num_v_heads=GLOBAL_CONFIG.num_v_heads // tp_size, + head_k_dim=GLOBAL_CONFIG.head_k_dim, + head_v_dim=GLOBAL_CONFIG.head_v_dim, + qkv_dtype=GLOBAL_CONFIG.qkv_dtype, + state_dtype=GLOBAL_CONFIG.state_dtype, + ) + + +def make_fused_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): + torch.manual_seed(seed) + mixed_qkv_conv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state = torch.randn( + tokens, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + state_work = torch.empty_like(state) + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + out = torch.empty(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) + return mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out + + +def make_native_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): + torch.manual_seed(seed) + q = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + k = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + v = torch.randn(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state = torch.randn( + tokens, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + state_work = torch.empty_like(state) + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + out = torch.empty_like(v) + return q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out + + +def profiler_start() -> None: + torch.cuda.profiler.start() + + +def profiler_stop() -> None: + torch.cuda.profiler.stop() + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--op", choices=("fused", "native"), default="fused") + parser.add_argument("--tokens", type=int, default=128) + parser.add_argument("--tp-size", type=int, default=1) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rep", type=int, default=3) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--device", type=int, default=0) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("No CUDA device is available") + + torch.cuda.set_device(args.device) + device = torch.device("cuda", torch.cuda.current_device()) + config = local_config_from_tp_size(args.tp_size) + + if args.op == "fused": + mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_fused_inputs( + args.tokens, device, args.seed, config + ) + + def run() -> None: + cula_cuda.qwen35_layout_scalar_kda_decode( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + state_work, + state_indices, + out, + ) + + else: + q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_native_inputs( + args.tokens, device, args.seed, config + ) + + def run() -> None: + cula_cuda.qwen35_scalar_kda_decode( + q, + k, + v, + a, + b, + A_log, + dt_bias, + state_work, + state_indices, + out, + ) + + for _ in range(args.warmup): + state_work.copy_(state) + run() + torch.cuda.synchronize() + + state_work.copy_(state) + torch.cuda.synchronize() + + print( + f"profile op={args.op} tokens={args.tokens} tp={args.tp_size} " + f"warmup={args.warmup} rep={args.rep} device={torch.cuda.get_device_name(device)}" + ) + profiler_start() + for _ in range(args.rep): + run() + torch.cuda.synchronize() + profiler_stop() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/profile_qwen35_scalar_prefill.py b/benchmarks/profile_qwen35_scalar_prefill.py new file mode 100644 index 00000000..390c18f3 --- /dev/null +++ b/benchmarks/profile_qwen35_scalar_prefill.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Nsight Compute target for the legacy Qwen scalar CUDA prefill path. + +This calls only ``qwen35_scalar_kda_prefill_core``. It never resolves or +imports the experimental CuTe prefill backend. Use ``ncu +--profile-from-start off`` so only the launches bracketed by +``cudaProfilerStart/Stop`` are collected. + +Examples:: + + ncu --profile-from-start off --kernel-name-base demangled \ + --kernel-name 'regex:.*qwen35_chunk_state_output_sm100_ts_kernel.*' \ + -o /tmp/qwen35_scalar_t256_state_output \ + python benchmarks/profile_qwen35_scalar_prefill.py \ + --config-json /data/xinhaowei/qwen_configs/Qwen3.5-27B/config.json \ + --seq-len 256 --rep 1 + + ncu --profile-from-start off --kernel-name-base demangled \ + --kernel-name 'regex:.*qwen35_chunk_(preprocess|state_output).*' \ + -o /tmp/qwen35_scalar_t512_stages \ + python benchmarks/profile_qwen35_scalar_prefill.py \ + --config-json /data/xinhaowei/qwen_configs/Qwen3.5-27B/config.json \ + --seq-len 512 --rep 1 +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys + +import torch +import torch.nn.functional as F + +ROOT = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +def load_shape(config_path: pathlib.Path, tp_size: int) -> tuple[int, int]: + root = json.loads(config_path.read_text(encoding="utf-8")) + config = root.get("text_config", root) + global_h = int(config["linear_num_key_heads"]) + global_hv = int(config["linear_num_value_heads"]) + if global_h % tp_size or global_hv % tp_size: + raise ValueError(f"TP={tp_size} must divide global H/HV={global_h}/{global_hv}") + h, hv = global_h // tp_size, global_hv // tp_size + if hv % h: + raise ValueError(f"local HV={hv} must be divisible by local H={h}") + if int(config["linear_key_head_dim"]) != 128 or int(config["linear_value_head_dim"]) != 128: + raise ValueError("the scalar CUDA prefill path requires K=V=128") + return h, hv + + +def extension_path(cula_cuda) -> str: + if not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core"): + raise RuntimeError("the loaded cuLA extension does not expose qwen35_scalar_kda_prefill_core") + op = cula_cuda.qwen35_scalar_kda_prefill_core + module = sys.modules.get(getattr(op, "__module__", "")) + return str(getattr(module, "__file__", "unavailable")) + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--seq-len", type=int, default=256) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rep", type=int, default=1) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--device", type=int, default=0) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("no CUDA device is available") + if args.rep < 1 or args.warmup < 0: + parser.error("--rep must be positive and --warmup must be non-negative") + if args.seq_len < 32: + parser.error("the chunk scalar CUDA path requires --seq-len >= 32") + + import cula.cudac as cula_cuda + + torch.cuda.set_device(args.device) + device = torch.device("cuda", torch.cuda.current_device()) + h, hv = load_shape(args.config_json, args.tp_size) + torch.manual_seed(args.seed) + + q = torch.randn(1, args.seq_len, h, 128, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, args.seq_len, hv, 128, device=device, dtype=torch.bfloat16) + a = torch.randn(1, args.seq_len, hv, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(hv, device=device, dtype=torch.float32) + dt_bias = torch.randn(hv, device=device, dtype=torch.float32) * 0.1 + g = -torch.exp(A_log).view(1, 1, hv) * F.softplus(a.float() + dt_bias.view(1, 1, hv)) + beta = torch.sigmoid(b.float()) + initial_state = torch.randn(1, hv, 128, 128, device=device, dtype=torch.float32) * 0.01 + cu_seqlens = torch.tensor([0, args.seq_len], device=device, dtype=torch.int32) + out = torch.empty_like(v) + final_state = torch.empty_like(initial_state) + + def run() -> None: + cula_cuda.qwen35_scalar_kda_prefill_core( + q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + out, + final_state, + ) + + for _ in range(args.warmup): + run() + torch.cuda.synchronize() + + print( + f"profile scalar_cuda_core T={args.seq_len} TP={args.tp_size} H/HV={h}/{hv} " + f"warmup={args.warmup} rep={args.rep} device={torch.cuda.get_device_name(device)}" + ) + print(f"extension={extension_path(cula_cuda)}") + torch.cuda.profiler.start() + for _ in range(args.rep): + run() + torch.cuda.synchronize() + torch.cuda.profiler.stop() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/tune_qwen35_tp_policy.py b/benchmarks/tune_qwen35_tp_policy.py new file mode 100644 index 00000000..28653e9c --- /dev/null +++ b/benchmarks/tune_qwen35_tp_policy.py @@ -0,0 +1,362 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Tune Qwen3.5 TP-local kernel policies. + +This is a configuration-driven tuner. It benchmarks only policies that are +compiled into the current extension and records unsupported candidates in the +result file. The initial compiled policy is the decode traits currently used by +the CUDA/CuTe kernels: + + layout_vec=4, kda_threads=128, kda_tile_v=16, kda_tile_k=16, heads_per_cta=1 + +When more C++ policy specializations are added, extend `compiled_policy_key` +and the kernel dispatch path; this script can then sweep them without changing +the output format. +""" + +from __future__ import annotations + +import argparse +import csv +import itertools +import json +import pathlib +import sys +from dataclasses import asdict, dataclass +from typing import Any + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + + +@dataclass(frozen=True) +class DecodePolicy: + name: str + layout_vec: int + kda_threads: int + kda_tile_v: int + kda_tile_k: int + heads_per_cta: int = 1 + + @property + def key(self) -> tuple[int, int, int, int, int]: + return (self.layout_vec, self.kda_threads, self.kda_tile_v, self.kda_tile_k, self.heads_per_cta) + + +CURRENT_DECODE_POLICY = DecodePolicy( + name="current", + layout_vec=4, + kda_threads=128, + kda_tile_v=16, + kda_tile_k=16, + heads_per_cta=1, +) + + +def decode_benchmarks(): + from benchmarks import bench_qwen35_decode + + return bench_qwen35_decode + + +def compiled_policy_key(policy: DecodePolicy) -> str | None: + """Return the compiled backend selector for a policy, or None if absent.""" + if policy.key == CURRENT_DECODE_POLICY.key: + return "current" + return None + + +def _list_from_json(data: dict[str, Any], key: str, default: list[int]) -> list[int]: + value = data.get(key, default) + if not isinstance(value, list) or not value: + raise ValueError(f"{key} must be a non-empty list") + return [int(item) for item in value] + + +def load_decode_policies(path: pathlib.Path | None) -> list[DecodePolicy]: + if path is None: + return [CURRENT_DECODE_POLICY] + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + + if isinstance(data, list): + policies = [] + for idx, item in enumerate(data): + if not isinstance(item, dict): + raise ValueError(f"Policy entry {idx} must be an object") + policies.append( + DecodePolicy( + name=str(item.get("name", f"policy_{idx}")), + layout_vec=int(item["layout_vec"]), + kda_threads=int(item["kda_threads"]), + kda_tile_v=int(item["kda_tile_v"]), + kda_tile_k=int(item["kda_tile_k"]), + heads_per_cta=int(item.get("heads_per_cta", 1)), + ) + ) + return policies + + if not isinstance(data, dict): + raise ValueError("Policy grid must be a JSON object or list") + if data.get("mode", "decode") != "decode": + raise ValueError("Only decode policy grids are supported by this tuner") + + policies = [] + for idx, combo in enumerate( + itertools.product( + _list_from_json(data, "layout_vec", [CURRENT_DECODE_POLICY.layout_vec]), + _list_from_json(data, "kda_threads", [CURRENT_DECODE_POLICY.kda_threads]), + _list_from_json(data, "kda_tile_v", [CURRENT_DECODE_POLICY.kda_tile_v]), + _list_from_json(data, "kda_tile_k", [CURRENT_DECODE_POLICY.kda_tile_k]), + _list_from_json(data, "heads_per_cta", [CURRENT_DECODE_POLICY.heads_per_cta]), + ) + ): + layout_vec, kda_threads, kda_tile_v, kda_tile_k, heads_per_cta = combo + policies.append( + DecodePolicy( + name=f"p{idx}_lv{layout_vec}_th{kda_threads}_tv{kda_tile_v}_tk{kda_tile_k}_h{heads_per_cta}", + layout_vec=layout_vec, + kda_threads=kda_threads, + kda_tile_v=kda_tile_v, + kda_tile_k=kda_tile_k, + heads_per_cta=heads_per_cta, + ) + ) + return policies + + +def write_example_grid(path: pathlib.Path) -> None: + example = { + "mode": "decode", + "layout_vec": [4, 8], + "kda_threads": [64, 128, 256], + "kda_tile_v": [8, 16, 32], + "kda_tile_k": [8, 16, 32], + "heads_per_cta": [1, 2, 4], + } + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(example, f, indent=2, sort_keys=True) + + +def bucket_name(tokens: int) -> str: + if tokens <= 4: + return "tokens<=4" + if tokens <= 16: + return "tokens<=16" + if tokens <= 64: + return "tokens<=64" + return "tokens>64" + + +def run_decode_policy( + *, + scope: str, + tokens: int, + tp_size: int, + warmup: int, + rep: int, + seed: int, + policy: DecodePolicy, +) -> dict[str, Any]: + decode_bench = decode_benchmarks() + config = decode_bench.local_config_from_tp_size(tp_size) + compiled_key = compiled_policy_key(policy) + row: dict[str, Any] = { + "mode": "decode", + "scope": scope, + "tokens": tokens, + "token_bucket": bucket_name(tokens), + "tp_size": tp_size, + "local_k_heads": config.num_k_heads, + "local_v_heads": config.num_v_heads, + "conv_dim": config.conv_dim, + "policy": policy.name, + "compiled_policy": compiled_key, + **asdict(policy), + } + if compiled_key is None: + row.update({"status": "unsupported", "ms": None, "us_per_token": None}) + return row + + device = decode_bench.accelerator_device() + if scope == "core": + ms = decode_bench.bench_native_core(tokens, device, warmup, rep, seed, config) + elif scope == "fused": + ms = decode_bench.bench_fused_layout_kda(tokens, device, warmup, rep, seed, config) + elif scope == "full": + ms = decode_bench.bench_full(tokens, device, warmup, rep, seed, config) + else: + raise ValueError(f"Unsupported decode scope={scope}") + + row.update({"status": "ok", "ms": ms, "us_per_token": ms * 1000.0 / tokens}) + return row + + +def choose_best(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} + for row in rows: + if row["status"] != "ok": + continue + key = (row["mode"], row["scope"], row["tp_size"], row["local_v_heads"], row["token_bucket"]) + groups.setdefault(key, []).append(row) + + best_rows = [] + for key, candidates in sorted(groups.items()): + best = min(candidates, key=lambda row: float(row["ms"])) + mode, scope, tp_size, local_v_heads, token_bucket = key + best_rows.append( + { + "mode": mode, + "scope": scope, + "tp_size": tp_size, + "local_v_heads": local_v_heads, + "token_bucket": token_bucket, + "policy": best["policy"], + "compiled_policy": best["compiled_policy"], + "ms": best["ms"], + "us_per_token": best["us_per_token"], + "layout_vec": best["layout_vec"], + "kda_threads": best["kda_threads"], + "kda_tile_v": best["kda_tile_v"], + "kda_tile_k": best["kda_tile_k"], + "heads_per_cta": best["heads_per_cta"], + } + ) + return best_rows + + +def write_csv(path: pathlib.Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fieldnames = [ + "mode", + "scope", + "tokens", + "token_bucket", + "tp_size", + "local_k_heads", + "local_v_heads", + "conv_dim", + "policy", + "compiled_policy", + "status", + "ms", + "us_per_token", + "layout_vec", + "kda_threads", + "kda_tile_v", + "kda_tile_k", + "heads_per_cta", + ] + with path.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + writer.writerows(rows) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Tune Qwen3.5 TP-local kernel policies.") + parser.add_argument("--mode", choices=["decode"], default="decode") + parser.add_argument("--scope", choices=["core", "fused", "full", "all"], default="fused") + parser.add_argument("--tp-sizes", nargs="+", type=int, choices=[1, 2, 4, 8], default=[1, 2, 4, 8]) + parser.add_argument("--tokens", nargs="+", type=int, default=[1, 2, 4, 8, 16, 32, 64, 128]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--policy-grid", type=pathlib.Path, default=None) + parser.add_argument("--write-example-grid", type=pathlib.Path, default=None) + parser.add_argument("--output-json", type=pathlib.Path, default=pathlib.Path("tmp/qwen35_tp_policy_tune.json")) + parser.add_argument("--csv", type=pathlib.Path, default=None) + parser.add_argument("--fail-on-unsupported", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.write_example_grid is not None: + write_example_grid(args.write_example_grid) + print(f"wrote example policy grid: {args.write_example_grid}") + return 0 + + policies = load_decode_policies(args.policy_grid) + scopes = ["core", "fused", "full"] if args.scope == "all" else [args.scope] + decode_bench = decode_benchmarks() + device = decode_bench.accelerator_device() + device_name = decode_bench.accelerator_name(device) + + print(f"Qwen3.5 TP policy tuner: mode={args.mode} device={device_name}") + print(f"tp_sizes={args.tp_sizes} tokens={args.tokens} scopes={scopes}") + print(f"policies={len(policies)} compiled={sum(compiled_policy_key(p) is not None for p in policies)}") + + rows: list[dict[str, Any]] = [] + for policy in policies: + compiled_key = compiled_policy_key(policy) + if compiled_key is None: + print(f"skip unsupported policy={policy.name} {asdict(policy)}") + for scope in scopes: + for tp_size in args.tp_sizes: + for tokens in args.tokens: + row = run_decode_policy( + scope=scope, + tokens=tokens, + tp_size=tp_size, + warmup=args.warmup, + rep=args.rep, + seed=args.seed, + policy=policy, + ) + rows.append(row) + if row["status"] == "ok": + print( + f"{scope:>5} tp={tp_size} hv={row['local_v_heads']:>2} tokens={tokens:>4} " + f"policy={policy.name} ms={row['ms']:.4f} us/tok={row['us_per_token']:.2f}" + ) + + unsupported = [row for row in rows if row["status"] == "unsupported"] + if unsupported and args.fail_on_unsupported: + raise RuntimeError(f"{len(unsupported)} policy/shape rows are unsupported by the compiled extension") + + best_rows = choose_best(rows) + result = { + "device": device_name, + "mode": args.mode, + "warmup": args.warmup, + "rep": args.rep, + "seed": args.seed, + "rows": rows, + "best": best_rows, + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + with args.output_json.open("w", encoding="utf-8") as f: + json.dump(result, f, indent=2, sort_keys=True) + print(f"wrote {args.output_json}") + + if args.csv is not None: + write_csv(args.csv, rows) + print(f"wrote {args.csv}") + + if best_rows: + print("best policies:") + for row in best_rows: + print( + f" {row['scope']:>5} tp={row['tp_size']} hv={row['local_v_heads']:>2} " + f"{row['token_bucket']}: {row['policy']} {row['ms']:.4f} ms" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index 5a0f6299..c2611f94 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -17,6 +17,9 @@ #include #include +#include "qwen35/decode/qwen35_decode_common.cuh" +#include "qwen35/prefill/qwen35_prefill_common.cuh" + #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) void ChunkKDAFwdIntra( @@ -70,6 +73,176 @@ kda_fwd_prefill( std::optional raw_cu_seqlens_); #endif +void +qwen35_conv1d_decode( + at::Tensor mixed_qkv, + at::Tensor conv_state, + at::Tensor conv_weight, + at::Tensor out) { + cula::qwen35::decode::ConvDecodeParams params{ + mixed_qkv, + conv_state, + conv_weight, + out, + }; + cula::qwen35::decode::run_qwen35_conv1d_decode(params); +} + +void +qwen35_layout_decode( + at::Tensor mixed_qkv_conv, + at::Tensor a, + at::Tensor b, + at::Tensor q_rep, + at::Tensor k_rep, + at::Tensor v, + at::Tensor a_kernel, + at::Tensor b_kernel) { + cula::qwen35::decode::LayoutDecodeParams params{ + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + }; + cula::qwen35::decode::run_qwen35_layout_decode(params); +} + +void +qwen35_scalar_kda_decode( + at::Tensor q_rep, + at::Tensor k_rep, + at::Tensor v, + at::Tensor a_kernel, + at::Tensor b_kernel, + at::Tensor A_log, + at::Tensor dt_bias, + at::Tensor recurrent_state, + at::Tensor pool_idx, + at::Tensor out) { + cula::qwen35::decode::ScalarKdaDecodeParams params{ + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + }; + cula::qwen35::decode::run_qwen35_scalar_kda_decode(params); +} + +void +qwen35_layout_scalar_kda_decode( + at::Tensor mixed_qkv_conv, + at::Tensor a, + at::Tensor b, + at::Tensor A_log, + at::Tensor dt_bias, + at::Tensor recurrent_state, + at::Tensor pool_idx, + at::Tensor out) { + cula::qwen35::decode::LayoutScalarKdaDecodeParams params{ + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + }; + cula::qwen35::decode::run_qwen35_layout_scalar_kda_decode(params); +} + +void +qwen35_scalar_kda_prefill( + at::Tensor q, + at::Tensor k, + at::Tensor v, + at::Tensor a, + at::Tensor b, + at::Tensor A_log, + at::Tensor dt_bias, + at::Tensor initial_state, + at::Tensor cu_seqlens, + at::Tensor out, + at::Tensor final_state) { + cula::qwen35::prefill::ScalarKdaPrefillParams params{ + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + }; + cula::qwen35::prefill::run_qwen35_scalar_kda_prefill(params); +} + +void +qwen35_scalar_kda_prefill_core( + at::Tensor q, + at::Tensor k, + at::Tensor v, + at::Tensor g, + at::Tensor beta, + at::Tensor initial_state, + at::Tensor cu_seqlens, + at::Tensor out, + at::Tensor final_state) { + cula::qwen35::prefill::ScalarKdaPrefillCoreParams params{ + q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + out, + final_state, + }; + cula::qwen35::prefill::run_qwen35_scalar_kda_prefill_core(params); +} + +void +qwen35_layout_prefill( + at::Tensor mixed_qkv_conv, + at::Tensor a, + at::Tensor b, + at::Tensor q_rep, + at::Tensor k_rep, + at::Tensor v, + at::Tensor a_kernel, + at::Tensor b_kernel) { + cula::qwen35::prefill::LayoutPrefillParams params{ + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + }; + cula::qwen35::prefill::run_qwen35_layout_prefill(params); +} + +void +qwen35_chunk_qk_prefill_sm90(at::Tensor q, at::Tensor k, at::Tensor out) { + cula::qwen35::prefill::sm90::qwen35_chunk_qk_prefill_sm90(q, k, out); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "cuLA"; #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) @@ -96,4 +269,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { pybind11::arg("cp_seq_map_") = std::nullopt, pybind11::arg("raw_cu_seqlens_") = std::nullopt); #endif + m.def("qwen35_conv1d_decode", &qwen35_conv1d_decode); + m.def("qwen35_layout_decode", &qwen35_layout_decode); + m.def("qwen35_scalar_kda_decode", &qwen35_scalar_kda_decode); + m.def("qwen35_layout_scalar_kda_decode", &qwen35_layout_scalar_kda_decode); + m.def("qwen35_layout_prefill", &qwen35_layout_prefill); + m.def("qwen35_scalar_kda_prefill", &qwen35_scalar_kda_prefill); + m.def("qwen35_scalar_kda_prefill_core", &qwen35_scalar_kda_prefill_core); + m.def("qwen35_chunk_qk_prefill_sm90", &qwen35_chunk_qk_prefill_sm90); } diff --git a/csrc/kda/sm100/fwd_helpers.hpp b/csrc/kda/sm100/fwd_helpers.hpp index da9b5ae6..336387df 100644 --- a/csrc/kda/sm100/fwd_helpers.hpp +++ b/csrc/kda/sm100/fwd_helpers.hpp @@ -16,7 +16,10 @@ #include +#include + #include "kerutils/kerutils.cuh" +#include "kda/sm100/kda_fwd_common.cuh" namespace kda::sm100 { @@ -26,6 +29,21 @@ using ku::nvbf16x4; using ku::store_128b; using namespace cute; +template +CUTE_DEVICE void +gate_exp2_float4(float2& s1, float2& s2) { + if constexpr (std::is_same_v, ScalarGateView>) { + const float scale = exp2f(s1.x); + s1 = make_float2(scale, scale); + s2 = make_float2(scale, scale); + } else { + s1.x = exp2f(s1.x); + s1.y = exp2f(s1.y); + s2.x = exp2f(s2.x); + s2.y = exp2f(s2.y); + } +} + // ============================================================ // Forward Prologue: B-matrix (SMEM) helper functions // ============================================================ @@ -104,10 +122,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_0_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_0_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -117,10 +132,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_1_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_1_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -130,10 +142,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_2_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_2_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -143,10 +152,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -193,10 +199,7 @@ fwd_setup_kg_col1_3out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_1_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_1_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -206,10 +209,7 @@ fwd_setup_kg_col1_3out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_2_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_2_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -219,10 +219,7 @@ fwd_setup_kg_col1_3out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -266,10 +263,7 @@ fwd_setup_kg_col2_2out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_2_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_2_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -279,10 +273,7 @@ fwd_setup_kg_col2_2out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -317,10 +308,7 @@ fwd_setup_kg_col3_1out(G_TENSOR& sG, K_TENSOR& sK, KG_TENSOR& sKG_intra, int idx // intra(3,3): exp2(g_first_3 - g[x]) * K[x] float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], reinterpret_cast(&g)[0]); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], reinterpret_cast(&g)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, __bfloat1622float2(k.a)); reinterpret_cast(&res)[1] = float2_mul(s2, __bfloat1622float2(k.b)); @@ -365,10 +353,7 @@ fwd_setup_A_inter_intra_all( float4 g_ref = *reinterpret_cast(&sG(g_first_row, y)); float2 s1 = float2_sub(g_a, reinterpret_cast(&g_ref)[0]); float2 s2 = float2_sub(g_b, reinterpret_cast(&g_ref)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); reinterpret_cast(&res_inter[i * 4])[0] = float2_mul(s1, va); reinterpret_cast(&res_inter[i * 4])[1] = float2_mul(s2, vb); } @@ -377,10 +362,7 @@ fwd_setup_A_inter_intra_all( float4 g_ref = *reinterpret_cast(&sG(g_half_row, y)); float2 s1 = float2_sub(g_a, reinterpret_cast(&g_ref)[0]); float2 s2 = float2_sub(g_b, reinterpret_cast(&g_ref)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); reinterpret_cast(&res_intra[i * 4])[0] = float2_mul(s1, va); reinterpret_cast(&res_intra[i * 4])[1] = float2_mul(s2, vb); } @@ -417,10 +399,7 @@ fwd_setup_A_inter_all( float4 g_ref = *reinterpret_cast(&sG(g_first_row, y)); float2 s1 = float2_sub(reinterpret_cast(&g)[0], reinterpret_cast(&g_ref)[0]); float2 s2 = float2_sub(reinterpret_cast(&g)[1], reinterpret_cast(&g_ref)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); reinterpret_cast(&res_inter[i * 4])[0] = float2_mul(s1, va); reinterpret_cast(&res_inter[i * 4])[1] = float2_mul(s2, vb); } diff --git a/csrc/kda/sm100/kda_fwd_common.cuh b/csrc/kda/sm100/kda_fwd_common.cuh index 6bc227a0..ad5a8f46 100644 --- a/csrc/kda/sm100/kda_fwd_common.cuh +++ b/csrc/kda/sm100/kda_fwd_common.cuh @@ -18,14 +18,32 @@ namespace kda::sm100 { +// Presents a compact per-row scalar gate as the float4-addressable 2-D view +// used by the existing gated Q/K helpers. Each row stores four identical +// values; all logical K columns intentionally alias that four-float record. +struct ScalarGateView { + float* ptr; + + __device__ __forceinline__ float& + operator()(int row, int) const { + return ptr[row * 4]; + } +}; + // KDA forward kernels // KDA forward intra-chunk kernel void run_kda_fwd_intra_sm100(KDA_fwd_intra_params& params, cudaStream_t stream); +void +run_kda_fwd_intra_sm100_qwen_scalar_g(KDA_fwd_intra_params& params, cudaStream_t stream); + // KDA forward recompute W & U kernel void run_kda_fwd_recomp_w_u_sm100(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream); -} // namespace kda::sm100 \ No newline at end of file +void +run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream); + +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp index 3689bee3..2679772e 100644 --- a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp @@ -129,7 +129,9 @@ struct KdaChunkFwdIntraKernelSm100 { if (warp_idx == 0 && lane_predicate) { cute::prefetch_tma_descriptor(tma_params.tma_q.get_tma_descriptor()); cute::prefetch_tma_descriptor(tma_params.tma_k.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + if constexpr (!Mainloop::ScalarG) { + cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + } } // Allocate TMEM (warp 0 only) @@ -144,9 +146,10 @@ struct KdaChunkFwdIntraKernelSm100 { // === Unified TMA load pipeline: Q + K + G === typename PipelineQKG::Params qkg_load_pipe_params; - qkg_load_pipe_params.transaction_bytes = sizeof(ku::bf16) * cosize_v + // Q - sizeof(ku::bf16) * cosize_v + // K - sizeof(float) * cosize_v; // G + qkg_load_pipe_params.transaction_bytes = + sizeof(ku::bf16) * cosize_v + // Q + sizeof(ku::bf16) * cosize_v + // K + (Mainloop::ScalarG ? 0 : sizeof(float) * cosize_v); qkg_load_pipe_params.is_leader = lane_predicate && (role == WarpRole::Load); qkg_load_pipe_params.num_consumers = NumCudaCoreThreads; @@ -337,10 +340,16 @@ run_kda_fwd_intra_sm100_impl_dispatch(KDA_fwd_intra_params& params, cudaStream_t make_tensor(make_gmem_ptr((ku::bf16*)params.k_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); - auto tma_G = cute::make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), - typename Kernel::SmemLayoutInputFP32{}); + auto tma_G = [&]() { + if constexpr (Kernel::Mainloop::ScalarG) { + return 0; + } else { + return cute::make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), + typename Kernel::SmemLayoutInputFP32{}); + } + }(); // --- Pack TMA params --- typename Kernel:: @@ -374,11 +383,30 @@ run_kda_fwd_intra_sm100_impl(KDA_fwd_intra_params& params, cudaStream_t stream) BOOL_SWITCH(params.unified_gref, kUnifiedGRef, [&] { // Currently we hardcode RoundingTF32=false to align with FLA implementation, the precision is enough using Kernel = KdaChunkFwdIntraKernelSm100< - KdaChunkFwdIntraMainloopSm100>; + KdaChunkFwdIntraMainloopSm100< + kUseTF32Inverse, + /*RoundingTF32=*/false, + kUnifiedGRef, + /*ScalarG=*/false, + BetaType>>; run_kda_fwd_intra_sm100_impl_dispatch(params, stream); }); }); }); } -} // namespace kda::sm100 \ No newline at end of file +// Qwen GDN uses one scalar gate per token/value-head. Keep this entrypoint +// separate from the public vector-G dispatcher so the generic ABI and its +// template combinations remain unchanged. +inline void +run_kda_fwd_intra_sm100_qwen_scalar_g_impl(KDA_fwd_intra_params& params, cudaStream_t stream) { + using Kernel = KdaChunkFwdIntraKernelSm100>; + run_kda_fwd_intra_sm100_impl_dispatch(params, stream); +} + +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp index 55cdc686..1b6bb999 100644 --- a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp @@ -48,6 +48,7 @@ template < bool UseTF32Inverse_ = true, bool RoundingTF32_ = false, bool UnifiedGRef_ = false, + bool ScalarG_ = false, typename ElementBeta_ = float> struct KdaChunkFwdIntraMainloopSm100 { // ===================== Tile / Buffer Constants ===================== @@ -76,6 +77,7 @@ struct KdaChunkFwdIntraMainloopSm100 { // This makes inter and intra A-matrices identical, allowing the intra A-matrix to be skipped entirely. // Saves 50% of A-matrix exp2f computation and one TMEM store per k-iteration. static constexpr bool UnifiedGRef = UnifiedGRef_; + static constexpr bool ScalarG = ScalarG_; using ElementBeta = ElementBeta_; // double buffer in TMEM, overlap prologue A matrix with MMA @@ -194,12 +196,25 @@ struct KdaChunkFwdIntraMainloopSm100 { GmemLayoutAtom{}, Layout>>{})); // Val layout, 8 or 16 vals per store + struct VectorGateStorage { + array_aligned> g[StagesLoad]; + }; + + struct ScalarGateStorage { + // Four identical floats per row preserve the float4 load contract of + // the existing helpers while all logical K columns alias this record. + array_aligned g[StagesAcc]; + }; + + using GateStorage = std::conditional_t; + // ===================== Shared Memory Plan ===================== struct SharedMemoryPlan { - // Q, K, G double buffer + // Q/K use the TMA pipeline. Generic vector-G keeps a matching staged + // matrix; Qwen scalar-G keeps four duplicated floats per token row. array_aligned> q[StagesLoad]; // 12KB array_aligned> k[StagesLoad]; // 12KB - array_aligned> g[StagesLoad]; // 24KB + GateStorage gate; // Gated MMA K^T, double buffer struct { @@ -331,6 +346,13 @@ struct KdaChunkFwdIntraMainloopSm100 { int seq_len = cu_seqlens_ptr[batch_idx + 1] - cu_seqlens_ptr[batch_idx]; int sub_seq_len = min(TileT, seq_len - tile_idx * TileT); + // The Qwen specialization publishes compact scalar-g together + // with beta. Keep this stage until the tile epilogue so all four + // K slices can consume the expanded shared views. + if constexpr (ScalarG) { + beta_pipeline.consumer_wait(beta_pipe_state_read); + } + constexpr int kg_offset = SubTileT * TileK; // stride between sub_tile buffers CUTE_NO_UNROLL @@ -345,7 +367,16 @@ struct KdaChunkFwdIntraMainloopSm100 { // Step 2: Create SMEM tensor views for this buffer slot // ============================================================ Tensor sK = make_tensor(make_smem_ptr(shared_plan->k[buf_load_idx].data()), SmemLayoutInputBF16{}); - Tensor sG = make_tensor(make_smem_ptr(shared_plan->g[buf_load_idx].data()), SmemLayoutInputFP32{}); + auto sG = [&]() { + if constexpr (ScalarG) { + return ScalarGateView{ + shared_plan->gate.g[beta_pipe_state_read.index()].data()}; + } else { + return make_tensor( + make_smem_ptr(shared_plan->gate.g[buf_load_idx].data()), + SmemLayoutInputFP32{}); + } + }(); qkg_inter_pipeline.producer_acquire(qkg_inter_pipe_state_write); int buf_idx = qkg_inter_pipe_state_write.index(); @@ -484,7 +515,9 @@ struct KdaChunkFwdIntraMainloopSm100 { // and beta data before waiting for MMA, overlapping independent waits. kk_inv_pipeline.producer_acquire(kk_inv_pipe_state_write); - beta_pipeline.consumer_wait(beta_pipe_state_read); + if constexpr (!ScalarG) { + beta_pipeline.consumer_wait(beta_pipe_state_read); + } qk_done_pipeline.consumer_wait(qk_done_pipe_state_read); int buf_acc_idx = qk_done_pipe_state_read.index(); @@ -729,8 +762,15 @@ struct KdaChunkFwdIntraMainloopSm100 { make_coord(token_offset, _0{}, _0{}), tma_params.tma_q.get_tma_tensor(tma_params.shape_qk)); Tensor mK = domain_offset( make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_qk)); - Tensor mG = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + auto mG = [&]() { + if constexpr (ScalarG) { + return 0; + } else { + return domain_offset( + make_coord(token_offset, _0{}, _0{}), + tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + } + }(); // TMA load body (Q, K, G — unified pipeline, single barrier per stage) CUTE_NO_UNROLL @@ -738,20 +778,28 @@ struct KdaChunkFwdIntraMainloopSm100 { int buf_idx = qkg_load_pipe_state_write.index(); Tensor sQ = make_tensor(make_smem_ptr(shared_plan->q[buf_idx].data()), SmemLayoutInputBF16{}); Tensor sK = make_tensor(make_smem_ptr(shared_plan->k[buf_idx].data()), SmemLayoutInputBF16{}); - Tensor sG = make_tensor(make_smem_ptr(shared_plan->g[buf_idx].data()), SmemLayoutInputFP32{}); - // GVA: K and Q are sliced by qk_head_idx; G is sliced by head_idx (v-head). + // GVA: K and Q are sliced by qk_head_idx. Generic vector-G + // is sliced by v-head; Qwen scalar-G arrives via the aux + // pipeline and therefore has no TMA transfer here. Tensor gK = local_tile( mK(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); - Tensor gG = local_tile( - mG(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); Tensor gQ = local_tile( mQ(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); // Single acquire for all three TMA copies qkg_load_pipeline.producer_acquire(qkg_load_pipe_state_write); auto& barrier = *qkg_load_pipeline.producer_get_barrier(qkg_load_pipe_state_write); - ku::launch_tma_copy(tma_params.tma_g, gG, sG, barrier); + if constexpr (!ScalarG) { + Tensor sG = make_tensor( + make_smem_ptr(shared_plan->gate.g[buf_idx].data()), + SmemLayoutInputFP32{}); + Tensor gG = local_tile( + mG(_, _, head_idx), + make_shape(Int{}, Int{}), + make_coord(tile_idx, k_idx)); + ku::launch_tma_copy(tma_params.tma_g, gG, sG, barrier); + } ku::launch_tma_copy(tma_params.tma_k, gK, sK, barrier); ku::launch_tma_copy(tma_params.tma_q, gQ, sQ, barrier); ++qkg_load_pipe_state_write; @@ -911,11 +959,22 @@ struct KdaChunkFwdIntraMainloopSm100 { // Beta loading body beta_pipeline.producer_acquire(beta_pipe_state_write); if (thread_idx < TileT) { + const int token = token_offset + tile_idx * TileT + thread_idx; shared_plan->beta_smem[beta_pipe_state_write.index()][thread_idx] = (thread_idx < sub_seq_len) - ? float(reinterpret_cast( - params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h_v + head_idx]) + ? float(reinterpret_cast(params.beta_ptr)[token * params.h_v + head_idx]) : float(0); + if constexpr (ScalarG) { + const float gate = + (thread_idx < sub_seq_len) + ? reinterpret_cast(params.g_ptr)[token * params.h_v + head_idx] + : 0.0f; + float* gate4 = shared_plan->gate.g[beta_pipe_state_write.index()].data() + thread_idx * 4; + gate4[0] = gate; + gate4[1] = gate; + gate4[2] = gate; + gate4[3] = gate; + } } fence_view_async_shared(); beta_pipeline.producer_commit(beta_pipe_state_write); @@ -924,4 +983,4 @@ struct KdaChunkFwdIntraMainloopSm100 { } }; -} // namespace kda::sm100 \ No newline at end of file +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp index 2cae6c04..482290dd 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp @@ -152,7 +152,9 @@ struct KdaChunkFwdRecompWUKernelSm100 { cute::prefetch_tma_descriptor(tma_params.tma_akk.get_tma_descriptor()); cute::prefetch_tma_descriptor(tma_params.tma_k.get_tma_descriptor()); cute::prefetch_tma_descriptor(tma_params.tma_v.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + if constexpr (!Mainloop::ScalarG) { + cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + } if constexpr (StoreQG) { cute::prefetch_tma_descriptor(tma_params.tma_q.get_tma_descriptor()); } @@ -453,10 +455,16 @@ run_kda_fwd_recomp_w_u_sm100_impl_dispatch(KDA_fwd_recomp_w_u_params& params, cu make_tensor(make_gmem_ptr((bf16*)params.k_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); - auto tma_G = cute::make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), - typename Kernel::SmemLayoutInputFP32{}); + auto tma_G = [&]() { + if constexpr (Kernel::Mainloop::ScalarG) { + return 0; + } else { + return cute::make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), + typename Kernel::SmemLayoutInputFP32{}); + } + }(); auto tma_Akk = cute::make_tma_copy( SM90_TMA_LOAD{}, @@ -502,10 +510,18 @@ inline void run_kda_fwd_recomp_w_u_sm100_impl(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { BETA_TYPE_SWITCH(params.is_beta_bf16, BetaType, [&] { BOOL_SWITCH(params.store_qg, kStoreQG, [&] { - using Kernel = KdaChunkFwdRecompWUKernelSm100>; + using Kernel = KdaChunkFwdRecompWUKernelSm100< + KdaChunkFwdRecompWUMainloopSm100>; run_kda_fwd_recomp_w_u_sm100_impl_dispatch(params, stream); }); }); } -} // namespace kda::sm100 \ No newline at end of file +inline void +run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g_impl(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { + using Kernel = KdaChunkFwdRecompWUKernelSm100< + KdaChunkFwdRecompWUMainloopSm100>; + run_kda_fwd_recomp_w_u_sm100_impl_dispatch(params, stream); +} + +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp index d702e46f..30568877 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp @@ -35,7 +35,7 @@ struct KdaChunkFwdRecompWUSm100NamedBarriers { // constants, and the persistent loop bodies for each warp role. // The Kernel struct is templated on this Mainloop. // =================================================================== -template +template struct KdaChunkFwdRecompWUMainloopSm100 { // ===================== Tile / Buffer Constants ===================== static constexpr int TileT = 64; @@ -49,6 +49,7 @@ struct KdaChunkFwdRecompWUMainloopSm100 { static constexpr int StagesQ = 1; static constexpr bool StoreQG = StoreQG_; + static constexpr bool ScalarG = ScalarG_; using ElementBeta = ElementBeta_; // TODO: try optimization with tcgen05.mma.ws @@ -160,13 +161,22 @@ struct KdaChunkFwdRecompWUMainloopSm100 { struct QSmemBufferDisabled {}; // empty, zero-cost using QSmemBuffer = cute::conditional_t; + struct VectorGateStorage { + array_aligned> g[StagesLoadStore]; + }; + struct ScalarGateStorage { + array_aligned g[StagesA]; + }; + using GateStorage = cute::conditional_t; + struct SharedMemoryPlan { // Akk, single buffer array_aligned> akk[StagesA]; // 16KB - // K, V, G double buffer + // K/V are double buffered. Generic vector-G keeps the full staged + // matrix; Qwen scalar-G keeps four duplicated floats per token row. array_aligned> k[StagesLoadStore]; // 32KB array_aligned> v[StagesLoadStore]; // 32KB - array_aligned> g[StagesLoadStore]; // 64KB + GateStorage gate; // Q double buffer (only present when StoreQG=true) QSmemBuffer q_buf; // MMA B-operand staging: K_proc/V_proc after prologue, [N=TileK, K=TileT] MN-major, double buffer @@ -326,9 +336,19 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } - g_pipeline.consumer_wait(g_pipe_state_read); - Tensor sG = - make_tensor(make_smem_ptr(shared_plan->g[g_pipe_state_read.index()].data()), SmemLayoutInputFP32{}); + if constexpr (!ScalarG) { + g_pipeline.consumer_wait(g_pipe_state_read); + } + auto sG = [&]() { + if constexpr (ScalarG) { + return ScalarGateView{ + shared_plan->gate.g[beta_pipe_state_read.index()].data()}; + } else { + return make_tensor( + make_smem_ptr(shared_plan->gate.g[g_pipe_state_read.index()].data()), + SmemLayoutInputFP32{}); + } + }(); // Load G with same 16x64 column mapping as K (two float4 per iteration) #pragma unroll for (int ti = 0; ti < TileT / 16; ++ti) { @@ -368,15 +388,27 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // lo half (cols y..y+3): k_reg a01, a23 + g_reg lo float2 kf_01 = __bfloat1622float2(k_reg[ti][k_yi].a01); float2 kf_23 = __bfloat1622float2(k_reg[ti][k_yi].a23); - float2 g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; - float2 g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + float2 g_01; + float2 g_23; + float2 g_45; + float2 g_67; + if constexpr (ScalarG) { + const float scale = exp2f(g_reg[ti][k_yi][0].x); + g_01 = make_float2(scale, scale); + g_23 = make_float2(scale, scale); + g_45 = make_float2(scale, scale); + g_67 = make_float2(scale, scale); + } else { + g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; + g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; + g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; + } float2 res_01 = float2_mul(float2_mul(kf_01, beta2), g_01); float2 res_23 = float2_mul(float2_mul(kf_23, beta2), g_23); // hi half (cols y+4..y+7): k_reg a45, a67 + g_reg hi float2 kf_45 = __bfloat1622float2(k_reg[ti][k_yi].a45); float2 kf_67 = __bfloat1622float2(k_reg[ti][k_yi].a67); - float2 g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; - float2 g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; float2 res_45 = float2_mul(float2_mul(kf_45, beta2), g_45); float2 res_67 = float2_mul(float2_mul(kf_67, beta2), g_67); // Single 128-bit store @@ -431,23 +463,36 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // lo half (cols y..y+3): k_reg a01, a23 float2 kf_01 = __bfloat1622float2(k_reg[ti][k_yi].a01); float2 kf_23 = __bfloat1622float2(k_reg[ti][k_yi].a23); - float2 gd_01 = { - exp2f(g_last_reg[k_yi][0].x - g_reg[ti][k_yi][0].x), - exp2f(g_last_reg[k_yi][0].y - g_reg[ti][k_yi][0].y)}; - float2 gd_23 = { - exp2f(g_last_reg[k_yi][0].z - g_reg[ti][k_yi][0].z), - exp2f(g_last_reg[k_yi][0].w - g_reg[ti][k_yi][0].w)}; + float2 gd_01; + float2 gd_23; + float2 gd_45; + float2 gd_67; + if constexpr (ScalarG) { + const float scale = exp2f( + g_last_reg[k_yi][0].x - g_reg[ti][k_yi][0].x); + gd_01 = make_float2(scale, scale); + gd_23 = make_float2(scale, scale); + gd_45 = make_float2(scale, scale); + gd_67 = make_float2(scale, scale); + } else { + gd_01 = { + exp2f(g_last_reg[k_yi][0].x - g_reg[ti][k_yi][0].x), + exp2f(g_last_reg[k_yi][0].y - g_reg[ti][k_yi][0].y)}; + gd_23 = { + exp2f(g_last_reg[k_yi][0].z - g_reg[ti][k_yi][0].z), + exp2f(g_last_reg[k_yi][0].w - g_reg[ti][k_yi][0].w)}; + gd_45 = { + exp2f(g_last_reg[k_yi][1].x - g_reg[ti][k_yi][1].x), + exp2f(g_last_reg[k_yi][1].y - g_reg[ti][k_yi][1].y)}; + gd_67 = { + exp2f(g_last_reg[k_yi][1].z - g_reg[ti][k_yi][1].z), + exp2f(g_last_reg[k_yi][1].w - g_reg[ti][k_yi][1].w)}; + } float2 res_01 = float2_mul(kf_01, gd_01); float2 res_23 = float2_mul(kf_23, gd_23); // hi half (cols y+4..y+7): k_reg a45, a67 float2 kf_45 = __bfloat1622float2(k_reg[ti][k_yi].a45); float2 kf_67 = __bfloat1622float2(k_reg[ti][k_yi].a67); - float2 gd_45 = { - exp2f(g_last_reg[k_yi][1].x - g_reg[ti][k_yi][1].x), - exp2f(g_last_reg[k_yi][1].y - g_reg[ti][k_yi][1].y)}; - float2 gd_67 = { - exp2f(g_last_reg[k_yi][1].z - g_reg[ti][k_yi][1].z), - exp2f(g_last_reg[k_yi][1].w - g_reg[ti][k_yi][1].w)}; float2 res_45 = float2_mul(kf_45, gd_45); float2 res_67 = float2_mul(kf_67, gd_67); // Single 128-bit store @@ -468,8 +513,10 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } - g_pipeline.consumer_release(g_pipe_state_read); - ++g_pipe_state_read; + if constexpr (!ScalarG) { + g_pipeline.consumer_release(g_pipe_state_read); + ++g_pipe_state_read; + } // Ensure all 128 prologue threads have finished writing sKG_out cutlass::arch::NamedBarrier::arrive_and_wait( @@ -534,15 +581,27 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // lo half (cols y..y+3) float2 qf_01 = __bfloat1622float2(q_reg[ti][k_yi].a01); float2 qf_23 = __bfloat1622float2(q_reg[ti][k_yi].a23); - float2 g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; - float2 g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + float2 g_01; + float2 g_23; + float2 g_45; + float2 g_67; + if constexpr (ScalarG) { + const float scale = exp2f(g_reg[ti][k_yi][0].x); + g_01 = make_float2(scale, scale); + g_23 = make_float2(scale, scale); + g_45 = make_float2(scale, scale); + g_67 = make_float2(scale, scale); + } else { + g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; + g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; + g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; + } float2 res_01 = float2_mul(qf_01, g_01); float2 res_23 = float2_mul(qf_23, g_23); // hi half (cols y+4..y+7) float2 qf_45 = __bfloat1622float2(q_reg[ti][k_yi].a45); float2 qf_67 = __bfloat1622float2(q_reg[ti][k_yi].a67); - float2 g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; - float2 g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; float2 res_45 = float2_mul(qf_45, g_45); float2 res_67 = float2_mul(qf_67, g_67); // Single 128-bit store @@ -919,8 +978,15 @@ struct KdaChunkFwdRecompWUMainloopSm100 { make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_qk)); Tensor mV = domain_offset( make_coord(token_offset, _0{}, _0{}), tma_params.tma_v.get_tma_tensor(tma_params.shape_vg)); - Tensor mG = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + auto mG = [&]() { + if constexpr (ScalarG) { + return 0; + } else { + return domain_offset( + make_coord(token_offset, _0{}, _0{}), + tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + } + }(); Tensor mA = domain_offset( make_coord(token_offset, _0{}, _0{}), tma_params.tma_akk.get_tma_tensor(tma_params.shape_Akk)); @@ -958,26 +1024,37 @@ struct KdaChunkFwdRecompWUMainloopSm100 { make_smem_ptr(shared_plan->k[k_pipe_state_write.index()].data()), SmemLayoutInputBF16{}); Tensor sV = make_tensor( make_smem_ptr(shared_plan->v[v_pipe_state_write.index()].data()), SmemLayoutInputBF16{}); - Tensor sG = make_tensor( - make_smem_ptr(shared_plan->g[g_pipe_state_write.index()].data()), SmemLayoutInputFP32{}); - // GVA slicing: K uses qk_head_idx; V and G use the v-head index. + // GVA slicing: K uses qk_head_idx and V uses v-head. + // Generic vector-G is loaded by TMA; Qwen scalar-G is + // published by the aux pipeline instead. Tensor gK = local_tile( mK(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); Tensor gV = local_tile( mV(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); - Tensor gG = local_tile( - mG(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); // K: Load → Compute k_pipeline.producer_acquire(k_pipe_state_write); ku::launch_tma_copy(tma_params.tma_k, gK, sK, *k_pipeline.producer_get_barrier(k_pipe_state_write)); ++k_pipe_state_write; - // G: Load → Compute - g_pipeline.producer_acquire(g_pipe_state_write); - ku::launch_tma_copy(tma_params.tma_g, gG, sG, *g_pipeline.producer_get_barrier(g_pipe_state_write)); - ++g_pipe_state_write; + // G: Load → Compute (generic vector-G only) + if constexpr (!ScalarG) { + Tensor sG = make_tensor( + make_smem_ptr(shared_plan->gate.g[g_pipe_state_write.index()].data()), + SmemLayoutInputFP32{}); + Tensor gG = local_tile( + mG(_, _, head_idx), + make_shape(Int{}, Int{}), + make_coord(tile_idx, i_k)); + g_pipeline.producer_acquire(g_pipe_state_write); + ku::launch_tma_copy( + tma_params.tma_g, + gG, + sG, + *g_pipeline.producer_get_barrier(g_pipe_state_write)); + ++g_pipe_state_write; + } // V: Load → Compute v_pipeline.producer_acquire(v_pipe_state_write); @@ -1039,12 +1116,24 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // ============================================================ beta_pipeline.producer_acquire(beta_pipe_state_write); if (thread_idx < TileT) { + const int token = token_offset + tile_idx * TileT + thread_idx; float beta_val = (thread_idx < sub_seq_len) - ? float(reinterpret_cast( - params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h_v + head_idx]) + ? float(reinterpret_cast(params.beta_ptr)[token * params.h_v + head_idx]) : float(0); shared_plan->beta_smem[beta_pipe_state_write.index()][thread_idx] = beta_val; + if constexpr (ScalarG) { + const float gate = + (thread_idx < sub_seq_len) + ? reinterpret_cast(params.g_ptr)[token * params.h_v + head_idx] + : 0.0f; + float* gate4 = + shared_plan->gate.g[beta_pipe_state_write.index()].data() + thread_idx * 4; + gate4[0] = gate; + gate4[1] = gate; + gate4[2] = gate; + gate4[3] = gate; + } } fence_view_async_shared(); beta_pipeline.producer_commit(beta_pipe_state_write); @@ -1053,4 +1142,4 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } }; -} // namespace kda::sm100 \ No newline at end of file +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_sm100.cu b/csrc/kda/sm100/kda_fwd_sm100.cu index edaaf0d9..1c10d7f2 100644 --- a/csrc/kda/sm100/kda_fwd_sm100.cu +++ b/csrc/kda/sm100/kda_fwd_sm100.cu @@ -23,9 +23,19 @@ run_kda_fwd_intra_sm100(KDA_fwd_intra_params& params, cudaStream_t stream) { kda::sm100::run_kda_fwd_intra_sm100_impl(params, stream); } +void +run_kda_fwd_intra_sm100_qwen_scalar_g(KDA_fwd_intra_params& params, cudaStream_t stream) { + kda::sm100::run_kda_fwd_intra_sm100_qwen_scalar_g_impl(params, stream); +} + void run_kda_fwd_recomp_w_u_sm100(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { kda::sm100::run_kda_fwd_recomp_w_u_sm100_impl(params, stream); } +void +run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { + kda::sm100::run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g_impl(params, stream); +} + } // namespace kda::sm100 diff --git a/csrc/kda/sm90/collective/load_tma.hpp b/csrc/kda/sm90/collective/load_tma.hpp index 1d427b0c..f05dcbcd 100644 --- a/csrc/kda/sm90/collective/load_tma.hpp +++ b/csrc/kda/sm90/collective/load_tma.hpp @@ -106,22 +106,28 @@ struct CollectiveLoadTma { return g_full; } else if constexpr (kind == LoadKind::kAlpha) { // Alpha (gate) is per V/O head under GVA. + constexpr int AlphaWidth = decltype(size<1>(SmemLayout{}))::value; DPRINTF0_W( "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", to_string(kind), work_desc.seq_idx, work_desc.o_head_idx(), work_desc.tok_offset); + constexpr bool ScalarAlpha = AlphaWidth == 4; + const int alpha_head_groups = + ScalarAlpha ? problem_size.num_v_heads / AlphaWidth : problem_size.num_v_heads; + const int alpha_head_group = + ScalarAlpha ? work_desc.o_head_idx() / AlphaWidth : work_desc.o_head_idx(); Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( problem_size.total_seqlen, - problem_size.head_size, - problem_size.num_v_heads)); // global view to the packed varlen sequence - Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx + Int{}, + alpha_head_groups)); // global view to packed tokens x 4-head groups + Tensor m_varlen = m_varlen_head(_, _, alpha_head_group); // slice group containing current V head Tensor m_offset = domain_offset( make_coord(work_desc.tok_offset, _0{}), m_varlen); // offset to start of the current sequence Tensor g_full = - local_tile(m_offset, make_tile(BlkSeqQ, HeadSize), make_coord(_, _0{})); // (blk, d, iter_blk) + local_tile(m_offset, make_tile(BlkSeqQ, Int{}), make_coord(_, _0{})); // (blk, d, iter_blk) return g_full; } else { // K lives in the QK head space; V lives in the V head space. @@ -140,8 +146,10 @@ struct CollectiveLoadTma { problem_size.total_seqlen, num_kv_heads)); // global view to the packed varlen sequence Tensor m_varlen = m_varlen_head(_, _, head_idx); // slice into current head_idx + const int feature_offset = + kIsK ? 0 : work_desc.value_tile_idx * HeadSize; Tensor m_offset = domain_offset( - make_coord(_0{}, work_desc.tok_offset), + make_coord(feature_offset, work_desc.tok_offset), m_varlen); // offset to start of the current sequence Tensor g_full = local_tile(m_offset, make_tile(HeadSize, BlkSeqKV), make_coord(_0{}, _)); // (d, blk, iter_blk) diff --git a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp index 66563e5a..4c7af669 100644 --- a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp +++ b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp @@ -90,10 +90,12 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr bool kIsPersistent = find_option_t::value; static constexpr bool kInitStateFromInput = find_option_t::value; + static constexpr bool SplitValueDim = find_option_t::value; static constexpr int NumLoadWarpGroups = 1; - static constexpr int NumStateMmaWarpGroups = 2; + static constexpr int NumStateMmaWarpGroups = SplitValueDim ? 1 : 2; static constexpr int NumAuxMmaWarpGroups = 1; + static constexpr int NumValueTiles = SplitValueDim ? 2 : 1; static constexpr int StageCountQ = find_option_t, Options>::value; static constexpr int StageCountK = find_option_t, Options>::value; @@ -101,6 +103,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr int NeedsAlpha = find_option_t::value; static constexpr int NeedsBeta = find_option_t::value; + static constexpr bool ScalarAlpha = find_option_t::value; + static constexpr bool StateKVLayout = find_option_t::value; static_assert(NeedsAlpha && NeedsBeta, "Alpha and Beta are both used in KDA."); static constexpr int SafeGate = true; // only support safe_gate=true @@ -136,7 +140,11 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr auto BlkSeqKV = get<1>(TileShape{}); // Blk_K/V static constexpr auto HeadSize = get<2>(TileShape{}); // D (Dq, Dk, Dv all equal) static constexpr auto HeadSizeQK = HeadSize; - static constexpr auto HeadSizeV = HeadSize; + using HeadSizeVType = std::conditional_t< + SplitValueDim, + _64, + std::remove_cv_t>; + static constexpr auto HeadSizeV = HeadSizeVType{}; using HeadSizeHalf = _64; using HeadSizeQuar = _32; @@ -151,6 +159,11 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { using TileShapeO2 = decltype(make_shape(HeadSizeV, BlkSeqQ, BlkSeqKV)); using TileShapeO1 = decltype(make_shape(HeadSizeV, BlkSeqQ, HeadSizeQK)); + using StateMmaSchedule = std::conditional_t< + NumStateMmaWarpGroups == 1, + cutlass::gemm::KernelTmaWarpSpecialized, + cutlass::gemm::KernelTmaWarpSpecializedCooperative>; + static_assert(BlkSeqQ % 64 == 0); static_assert(BlkSeqQ == 64 || BlkSeqQ == 128); static_assert(BlkSeqQ == BlkSeqKV); @@ -200,24 +213,37 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeKV, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using SmemLayoutAlphaAtom = GMMA::Layout_K_SW128_Atom; - using SmemLayoutAlpha_SD = decltype(tile_to_shape( + using SmemLayoutAlphaVector_SD = decltype(tile_to_shape( SmemLayoutAlphaAtom{}, make_shape( shape<1>(TileShapeQK{}), shape<2>(TileShapeQK{}), Int{}))); // (blk_kv, head_size), (64, 128) - using GmemShapeAlpha = Shape; // (seqlen_k, d, h) + // TMA requires its first global mode to be contiguous. A single scalar + // head is strided by HV across tokens, so scalar mode loads four adjacent + // heads as one 16-byte unit and each CTA selects its hv % 4 lane. + using AlphaWidth = std::conditional_t(TileShapeQK{}))>; + using SmemLayoutAlphaScalar_SD = decltype(make_layout( + make_shape(shape<1>(TileShapeQK{}), _4{}, Int{}), + make_stride(_4{}, _1{}, Int<4 * size<1>(TileShapeQK{})>{}))); + using SmemLayoutAlphaLoad_SD = + std::conditional_t; + // Vector-alpha compute layouts stay intact for the generic path. Scalar + // specializations bypass them and fill MMA fragments from the compact + // [token, stage] tensor using token coordinates. + using SmemLayoutAlpha_SD = SmemLayoutAlphaVector_SD; + using GmemShapeAlpha = Shape; // (seqlen_k, d-or-1, h) using GmemStrideAlpha = Stride; using GmemLayoutAlpha = Layout; using GmemTiledCopyAlpha = cute::SM90_TMA_LOAD; using TMA_Alpha = decltype(make_tma_copy( GmemTiledCopyAlpha{}, make_tensor(make_gmem_ptr(static_cast(nullptr)), GmemLayoutAlpha{}), - take<0, 2>(SmemLayoutAlpha_SD{}), - select<1, 2>(TileShapeQK{}), + take<0, 2>(SmemLayoutAlphaLoad_SD{}), + make_shape(shape<1>(TileShapeQK{}), AlphaWidth{}), size<0>(ClusterShape{}))); // raw layout for copy @@ -247,7 +273,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeKV, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using RefLayoutKV = decltype(make_layout(select<0, 1>(TileShapeKV{}), LayoutRight{})); // (dv, dk) using CollectiveMmaO1 = typename cutlass::gemm::collective::CollectiveBuilder< @@ -263,7 +289,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeO1, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; // (blk_q,blk_k) to align with O2 mma, LayoutRight to align with QK mma output using DesiredLayoutQK = decltype(make_layout(select<0, 1>(TileShapeQK{}), LayoutRight{})); @@ -280,7 +306,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeO2, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using TiledMmaQK = typename CollectiveMmaQK::TiledMma; // Q@K^t using TiledMmaKV = decltype(convert_to_gmma_rs(typename CollectiveMmaKV::TiledMma{})); @@ -346,7 +372,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeSK, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using ElementAccumulatorNewV = float; using TileShapeNewV = decltype(make_shape(HeadSizeV, BlkSeqKV, BlkSeqKV)); @@ -365,7 +391,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeNewV, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; // FIXME: K@K^t are not exactly the same as Q@K^t, but similar enough (what does this mean??) using TiledMmaKK = typename CollectiveMmaQK::TiledMma; // T = inv(I + strict_lower_triangular(K@K^t)) @@ -379,6 +405,9 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // only store the last row in Alpha using SmemLayoutAlphaLast = decltype(make_layout(make_shape(HeadSize, Int{}))); + using SmemLayoutAlphaLastScalar = decltype(make_layout(make_shape(_1{}, Int{}))); + using SmemLayoutAlphaLastStorage = + std::conditional_t; using SmemLayoutBeta = decltype(make_layout(make_shape(BlkSeqQ, Int{}))); using MainloopQPipeline = cutlass::PipelineTmaAsync; @@ -416,7 +445,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr int LoadQBytes = size(QKSmemLayoutQ{}(_, _, _0{})) * sizeof(Element); static constexpr int LoadKBytes = size(KVSmemLayoutK{}(_, _, _0{})) * sizeof(Element); static constexpr int LoadVBytes = size(KVSmemLayoutV{}(_, _, _0{})) * sizeof(Element); - static constexpr int LoadAlphaBytes = size(QKQSmemLayoutAlpha{}(_, _, _0{})) * sizeof(ElementAlpha); + static constexpr int LoadAlphaBytes = size(SmemLayoutAlphaLoad_SD{}(_, _, _0{})) * sizeof(ElementAlpha); static constexpr int StoreOBytes = CollectiveStoreO::TmaTransactionBytes; using SharedStorageO = typename CollectiveStoreO::SharedStorage; @@ -429,7 +458,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { alignas( alignment_for_swizzle(KVSmemLayoutV{})) cute::array_aligned> smem_v; alignas(alignment_for_swizzle( - QKQSmemLayoutAlpha{})) cute::array_aligned> smem_alpha; + SmemLayoutAlphaLoad_SD{})) cute::array_aligned> smem_alpha; alignas( alignment_for_swizzle(SmemLayoutQK{})) cute::array_aligned> smem_qk; alignas(alignment_for_swizzle( @@ -442,7 +471,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { cute::array_aligned> smem_beta; // store last row in Alpha separately, used for S'=K^T NewV's epilogue and S+=decay(S') (one fused epilogue) - cute::array_aligned> smem_alpha_last; + cute::array_aligned> smem_alpha_last; }; using TMA_Q = typename CollectiveMmaQK::Params::TMA_A; @@ -454,7 +483,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { using LoadK = CollectiveLoadTma; using LoadV = CollectiveLoadTma; using LoadAlpha = - CollectiveLoadTma; + CollectiveLoadTma; using LoadBeta = CollectiveLoadVector< LoadKindVector::kBeta, @@ -493,6 +522,34 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { GmemLayoutBeta beta_layout; }; + template + CUTE_DEVICE static auto + make_state_tensor(Ptr ptr, ProblemShape const& problem_size, NumSeqs num_seqs) { + // The global state always remains a full [K=128,V=128] matrix. Split + // V64 CTAs select disjoint column tiles after constructing this view. + auto state_shape = make_shape( + Int{}, problem_size.head_size, problem_size.num_v_heads, num_seqs); + if constexpr (StateKVLayout) { + // Qwen exposes state as contiguous [N, HV, K, V]. Express that + // physical layout in the kernel's logical (K,V,HV,N) coordinates + // instead of paying for pre/post transpose kernels. + auto state_stride = make_stride( + int64_t(problem_size.head_size), + _1{}, + int64_t(HeadSizeQK) * problem_size.head_size, + int64_t(problem_size.num_v_heads) * int(HeadSizeQK) * problem_size.head_size); + return make_tensor(make_gmem_ptr(ptr), make_layout(state_shape, state_stride)); + } else { + return make_tensor(make_gmem_ptr(ptr), make_layout(state_shape, LayoutLeft{})); + } + } + + template + CUTE_DEVICE static auto + make_state_tensor(Ptr ptr, ProblemShape const& problem_size) { + return make_state_tensor(ptr, problem_size, problem_size.num_seqs); + } + template static bool can_implement(ProblemShape const& problem_size, Arguments const& args) { @@ -525,7 +582,9 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }, /*workspace=*/nullptr); - auto alpha_shape = make_shape(s, d, problem_size.num_v_heads); + const int32_t alpha_head_groups = + ScalarAlpha ? problem_size.num_v_heads / int(AlphaWidth{}) : problem_size.num_v_heads; + auto alpha_shape = make_shape(s, AlphaWidth{}, alpha_head_groups); auto alpha_stride = make_stride( get<0>(args.dAlpha), // seqlen stride get<1>(args.dAlpha), // head_dim stride @@ -535,8 +594,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TMA_Alpha tma_load_alpha = make_tma_copy( GmemTiledCopyAlpha{}, mAlpha, - take<0, 2>(SmemLayoutAlpha_SD{}), - select<1, 2>(TileShapeQK{}), + take<0, 2>(SmemLayoutAlphaLoad_SD{}), + make_shape(shape<1>(TileShapeQK{}), AlphaWidth{}), size<0>(ClusterShape{})); auto params_kv_v = CollectiveMmaKV_G2S::to_underlying_arguments( @@ -617,10 +676,11 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto k_collective_load = LoadK(params.tma_load_k, k_pipeline, storage.smem_k); auto v_collective_load = LoadV(params.tma_load_v, v_pipeline, storage.smem_v); auto alpha_collective_load = LoadAlpha{params.tma_load_alpha, alpha_pipeline, storage.smem_alpha}; + auto v_load_tile_shape = make_shape(BlkSeqQ, BlkSeqKV, HeadSizeV); auto q_src_dst = q_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); auto k_src_dst = k_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); - auto v_src_dst = v_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); + auto v_src_dst = v_collective_load.partition_SD(problem_size, v_load_tile_shape, work_desc); auto alpha_src_dst = alpha_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); CUTE_NO_UNROLL @@ -673,7 +733,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarp; Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); - Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); + Tensor sAlphaLoad = make_tensor(make_smem_ptr(storage.smem_alpha.data()), SmemLayoutAlphaLoad_SD{}); + Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLastStorage{}); auto extract_loop_body = [&](int blk, auto is_final_block_) INLINE_LAMBDA { constexpr bool is_final_block = decltype(is_final_block_)::value; @@ -681,15 +742,22 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { int B = is_final_block ? valid_seq_len(work_desc, blk) : BlkSeqKV; auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); + auto sAlphaLoadCurr = sAlphaLoad(_, _, alpha_smem_pipe_read.index()); Tensor sAlast_out = sAlast(_, alpha_last_smem_pipe_write.index()); alpha_pipeline.consumer_wait(alpha_smem_pipe_read); alpha_last_pipeline.producer_acquire(alpha_last_smem_pipe_write); // each thread copy 4 elements, total 128 elements with one warp - CUTE_UNROLL - for (int t = thread_idx; t < HeadSize; t += 32) { - sAlast_out(t) = sAqkq_curr(B - 1, t); + if constexpr (ScalarAlpha) { + if (thread_idx == 0) { + sAlast_out(_0{}) = sAlphaLoadCurr(B - 1, work_desc.o_head_idx() & 3); + } + } else { + CUTE_UNROLL + for (int t = thread_idx; t < HeadSize; t += 32) { + sAlast_out(t) = sAqkq_curr(B - 1, t); + } } cutlass::arch::fence_view_async_shared(); @@ -775,10 +843,13 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { Tensor Beta = make_tensor(make_smem_ptr(storage.smem_beta.data()), SmemLayoutBeta{}); Tensor AlphaLast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); + Tensor AlphaLastStorage = + make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLastStorage{}); Tensor sQqk = make_tensor(make_smem_ptr(storage.smem_q.data()), QKSmemLayoutQ{}); Tensor sKqk = make_tensor(make_smem_ptr(storage.smem_k.data()), QKSmemLayoutK{}); Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); + Tensor sAlphaLoad = make_tensor(make_smem_ptr(storage.smem_alpha.data()), SmemLayoutAlphaLoad_SD{}); Tensor sVkv = make_tensor(make_smem_ptr(storage.smem_v.data()), KVSmemLayoutV{}); Tensor sQK = make_tensor(make_smem_ptr(storage.smem_qk.data()), SmemLayoutQK{}); Tensor sO = make_tensor(make_smem_ptr(storage.smem_o.data()), SmemLayoutO{}); @@ -883,7 +954,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto thr_copy_o = tiled_copy_o.get_thread_slice(thread_idx); auto tOsO = thr_copy_o.partition_D(sO); - auto const cO = make_identity_tensor(Shape, Int>{}); + auto const cO = make_identity_tensor(Shape, Int>{}); Tensor tOcO = o1_thr_mma.partition_C(cO); auto const seq_idx = work_desc.seq_idx; @@ -902,12 +973,13 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto kv_load = [&](auto& tKVrKV) INLINE_LAMBDA { DPRINTF0_WG("[%d,%d,%d,%d]>> load tKVgKV -> tKVrKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); // GVA: state is stored per V/O head. - int num_state_heads = problem_size.num_v_heads; int state_head_idx = work_desc.o_head_idx(); - auto gKV = make_tensor( - make_gmem_ptr(params.ptr_input_state), - make_layout(make_shape(Int{}, Int{}, num_state_heads, problem_size.num_seqs)))( - _, _, state_head_idx, seq_idx); // (KDim, VDim), K-contiguous + auto gKV_full = make_state_tensor(params.ptr_input_state, problem_size)( + _, _, state_head_idx, seq_idx); // full (KDim, VDim) + auto gKV = local_tile( + gKV_full, + make_tile(Int{}, HeadSizeVType{}), + make_coord(_0{}, work_desc.value_tile_idx)); auto tiled_copy_kv = make_tiled_copy_C(Copy_Atom{}, kv_tiled_mma); auto thr_copy_kv = tiled_copy_kv.get_thread_slice(thread_idx); @@ -944,12 +1016,13 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { } DPRINTF0_WG("[%d,%d,%d,%d]>> save tKVrKV -> tKVgKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); // GVA: state is stored per V/O head. - int num_state_heads = problem_size.num_v_heads; int state_head_idx = work_desc.o_head_idx(); - auto gKV = make_tensor( - make_gmem_ptr(params.ptr_output_state), - make_layout(make_shape(Int{}, Int{}, num_state_heads, out_num_seqs)))( - _, _, state_head_idx, out_seq_idx); // (KDim, VDim), K-contiguous + auto gKV_full = make_state_tensor(params.ptr_output_state, problem_size, out_num_seqs)( + _, _, state_head_idx, out_seq_idx); // full (KDim, VDim) + auto gKV = local_tile( + gKV_full, + make_tile(Int{}, HeadSizeVType{}), + make_coord(_0{}, work_desc.value_tile_idx)); auto tiled_copy_kv = make_tiled_copy_C(Copy_Atom{}, kv_tiled_mma); auto thr_copy_kv = tiled_copy_kv.get_thread_slice(thread_idx); @@ -959,12 +1032,17 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }; auto s_decay = [&](auto& tKVrKV, auto const& alpha_last_smem_pipe_read) INLINE_LAMBDA { - Tensor alpha_last_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tKVcS(i); - auto [s, t] = coord; // (head_size_v, head_size_k) - tKVrKV(i) *= exp2f(alpha_last_curr(t)); - }); + if constexpr (ScalarAlpha) { + const float decay = exp2f(AlphaLastStorage(_0{}, alpha_last_smem_pipe_read.index())); + for_each(make_int_sequence{}, [&](auto i) { tKVrKV(i) *= decay; }); + } else { + Tensor alpha_last_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); + for_each(make_int_sequence{}, [&](auto i) { + auto coord = tKVcS(i); + auto [s, t] = coord; // (head_size_v, head_size_k) + tKVrKV(i) *= exp2f(alpha_last_curr(t)); + }); + } }; auto o1_epi = [&](auto& tOrO1) INLINE_LAMBDA { @@ -1034,6 +1112,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto sK_scaled_curr = sQ_K_scaled(_, _, _1{}); auto sAlast_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); + auto sAlphaLoadCurr = sAlphaLoad(_, _, alpha_smem_pipe_read.index()); auto sQqk_slice = flat_divide(sQqk_curr, tiler_qk); auto sKqk_slice = flat_divide(sKqk_curr, tiler_qk); auto sQ_scaled_slice = flat_divide(sQ_scaled_curr, tiler_qk); @@ -1054,12 +1133,12 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { if constexpr (!is_first_block) { // make sure sQ_K_scaled is already consumed for previous K^@V cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // Each WG iterates over 2 slices of 32 elements each. - // WG0 (thread_idx < 128): wg_idx=0, processes alpha indices {0,1}, Q/K dim1=0 - // WG1 (thread_idx >= 128): wg_idx=1, processes alpha indices {2,3}, Q/K dim1=1 + // Divide the four 32-d Q/K quarters over the active state WGs. + // V128 uses two quarters per WG; split V64 uses one WG for all four. { - int wg_idx = thread_idx / 128; // 0 or 1 - int alpha_base = wg_idx * 2; // 0 or 2 + constexpr int kQuarterCount = int(HeadSizeQK) / int(HeadSizeQuar{}); + constexpr int kQuartersPerWG = kQuarterCount / NumStateMmaWarpGroups; + int wg_idx = thread_idx / 128; // Allocate Q/K register fragments once (reused across slices) // Only shape/layout matters for partition_fragment_A, use compile-time indices @@ -1069,17 +1148,26 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { qk_thr_mma_rs_quar.partition_fragment_A(sKqk_slice(_, _, _0{}, make_coord(_0{}, _0{}))); auto tArA = make_fragment_like(tQKrQ_wg); - for (int s = 0; s < 2; ++s) { - // S2R Alpha: alpha_col = wg_idx * 2 + s - int alpha_col = alpha_base + s; - auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); - auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); - copy(CopyAlphaAtom{}, tAsA_cur, tArA); + for (int local_quarter = 0; local_quarter < kQuartersPerWG; ++local_quarter) { + int quarter = wg_idx * kQuartersPerWG + local_quarter; + int quarter_col = quarter & 1; + int quarter_group = quarter >> 1; + int alpha_col = quarter; + if constexpr (ScalarAlpha) { + for_each(make_int_sequence{}, [&](auto i) { + auto [seq, _] = tQcMq_quar(i); + tArA(i) = sAlphaLoadCurr(seq, work_desc.o_head_idx() & 3); + }); + } else { + auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); + auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); + copy(CopyAlphaAtom{}, tAsA_cur, tArA); + } cute::transform(tArA, [](auto g) { return exp2f(g); }); // S2R Q - auto sQqk_cur = sQqk_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sQqk_cur = sQqk_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsQ_cur = thr_load_qk_quar.partition_S(sQqk_cur); auto tQKrQ_cv = thr_load_qk_quar.retile_D(tQKrQ_wg); copy(tiled_load_qk_quar, tQKsQ_cur, tQKrQ_cv); @@ -1091,13 +1179,14 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }); // R2S Q -> stage 0 - auto sQ_scaled_cur = sQ_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sQ_scaled_cur = + sQ_scaled_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsQ_out = thr_store_qk_quar.partition_D(sQ_scaled_cur); auto tQKrQ_out_cv = thr_store_qk_quar.retile_S(tQKrQ_wg); copy(tiled_store_qk_quar, tQKrQ_out_cv, tQKsQ_out); // S2R K - auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_cur = thr_load_qk_quar.partition_S(sKqk_cur); auto tQKrK_cv = thr_load_qk_quar.retile_D(tQKrK_wg); copy(tiled_load_qk_quar, tQKsK_cur, tQKrK_cv); @@ -1109,7 +1198,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }); // R2S K -> stage 1 - auto sK_scaled_cur = sK_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sK_scaled_cur = + sK_scaled_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_out = thr_store_qk_quar.partition_D(sK_scaled_cur); auto tQKrK_out_cv = thr_store_qk_quar.retile_S(tQKrK_wg); copy(tiled_store_qk_quar, tQKrK_out_cv, tQKsK_out); @@ -1286,40 +1376,58 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // synchronize 2 WGs before rewriting sQ_K_scaled cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // exp(alpha_last - alpha) * K - // Each WG iterates over 2 slices of 32 elements each. - // WG0 (thread_idx < 128): wg_idx=0, alpha_last indices {0,1}, K/output dim1=0 - // WG1 (thread_idx >= 128): wg_idx=1, alpha_last indices {2,3}, K/output dim1=1 + // exp(alpha_last - alpha) * K over all four 32-d quarters. { - int wg_idx = thread_idx / 128; // 0 or 1 - int alpha_base = wg_idx * 2; // 0 or 2 + constexpr int kQuarterCount = int(HeadSizeQK) / int(HeadSizeQuar{}); + constexpr int kQuartersPerWG = kQuarterCount / NumStateMmaWarpGroups; + int wg_idx = thread_idx / 128; // Allocate K/Alpha register fragments once (reused across slices) auto tQKrK_wg = qk_thr_mma_rs_quar.partition_fragment_A(sKqk_slice(_, _, _0{}, make_coord(_0{}, _0{}))); auto tArA_wg = make_fragment_like(tQKrK_wg); + float scalar_alpha_last = 0.0f; + if constexpr (ScalarAlpha) { + scalar_alpha_last = AlphaLastStorage(_0{}, alpha_last_smem_pipe_read.index()); + } - for (int s = 0; s < 2; ++s) { + for (int local_quarter = 0; local_quarter < kQuartersPerWG; ++local_quarter) { + int quarter = wg_idx * kQuartersPerWG + local_quarter; + int quarter_col = quarter & 1; + int quarter_group = quarter >> 1; // S2R Alpha - int alpha_col = alpha_base + s; - auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); - auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); - copy(CopyAlphaAtom{}, tAsA_cur, tArA_wg); + int alpha_col = quarter; + if constexpr (ScalarAlpha) { + for_each(make_int_sequence{}, [&](auto i) { + auto [seq, _] = tQcMq_quar(i); + tArA_wg(i) = sAlphaLoadCurr(seq, work_desc.o_head_idx() & 3); + }); + } else { + auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); + auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); + copy(CopyAlphaAtom{}, tAsA_cur, tArA_wg); + } // S2R K - auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_cur = thr_load_qk_quar.partition_S(sKqk_cur); auto tQKrK_cv = thr_load_qk_quar.retile_D(tQKrK_wg); copy(tiled_load_qk_quar, tQKsK_cur, tQKrK_cv); // element-wise: exp(alpha_last - alpha) * K - int alast_idx = alpha_base + s; + int alast_idx = quarter; auto alpha_last_cur = sAlast_slice(_, alast_idx); for_each(make_int_sequence{}, [&](auto i) { auto coord = tQcMq_quar(i); auto [seq, t] = coord; auto alpha = tArA_wg(i); auto k = tQKrK_wg(i); - auto alpha_last = alpha_last_cur(t); + auto alpha_last = [&]() { + if constexpr (ScalarAlpha) { + return scalar_alpha_last; + } else { + return alpha_last_cur(t); + } + }(); auto k_scaled = Element(exp2f(alpha_last - alpha) * float(k)); tQKrK_wg(i) = k_scaled; if constexpr (is_final_block) { @@ -1330,7 +1438,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }); // R2S K -> stage 0 (reuse for KV update) - auto sQ_scaled_cur = sQ_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sQ_scaled_cur = + sQ_scaled_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_out = thr_store_qk_quar.partition_D(sQ_scaled_cur); auto tQKrK_out_cv = thr_store_qk_quar.retile_S(tQKrK_wg); copy(tiled_store_qk_quar, tQKrK_out_cv, tQKsK_out); @@ -1444,6 +1553,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); Tensor sAqkk = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKKSmemLayoutAlpha{}); + Tensor sAlphaLoad = make_tensor(make_smem_ptr(storage.smem_alpha.data()), SmemLayoutAlphaLoad_SD{}); Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); Tensor sKkv = make_tensor(make_smem_ptr(storage.smem_k.data()), KVSmemLayoutK{}); @@ -1517,6 +1627,10 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // index tensor auto cMqk_subchunk = make_identity_tensor(select<0, 1>(TileShape_SubChunk{})); auto tQKcMqk_subchunk = thr_mma_subchunk.partition_C(cMqk_subchunk); + auto cMq_subchunk = make_identity_tensor(select<0, 2>(TileShape_SubChunk{})); + auto tQcMq_bf16_subchunk = thr_mma_bf16_subchunk.partition_A(cMq_subchunk); + auto cNk_subchunk = make_identity_tensor(select<1, 2>(TileShape_SubChunk{})); + auto tKcNk_bf16_subchunk = thr_mma_bf16_subchunk.partition_B(cNk_subchunk); // do MMA at the granularity of 16x16x64 with two warps constexpr auto tiler_subchunk_alpha = Shape<_16, Shape<_32, _1>>{}; @@ -1525,6 +1639,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto sQqk_curr = sQqk(_, _, q_smem_pipe_read.index()); auto sKqk_curr = sKqk(_, _, k_smem_pipe_read.index()); auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); + auto sAlphaLoadCurr = sAlphaLoad(_, _, alpha_smem_pipe_read.index()); Tensor sBeta_curr = Beta(_, beta_smem_pipe_read.index()); // (_16,(_32,_1),_4,(_2,_2)):(_64,(_1,_0),_1024,(_32,_4096)) @@ -1565,11 +1680,20 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // layout) auto s2r_compute_subchunk_operandA = [&](auto r_, int j, int j0, int j1) INLINE_LAMBDA { // S2R g_r_j in BF16 MMA operand A layout (single load) - Tensor sAqkq_r_j = sAqkq_slice(_, _, r_, make_coord(_0{}, j)); - Tensor tAsA_r_j = alpha_Q_bf16_thr_copy.partition_S(sAqkq_r_j); Tensor tArA_r_j = make_fragment_like(tv_layout_bf16_mma_A); - Tensor tArA_r_j_cv = alpha_Q_bf16_thr_copy.retile_D(tArA_r_j); - copy(alpha_Q_bf16_tiled_copy, tAsA_r_j, tArA_r_j_cv); + if constexpr (ScalarAlpha) { + constexpr int kSubchunkRows = size<0>(TileShape_SubChunk{}); + for_each(make_int_sequence{}, [&](auto i) { + auto [row, _] = tQcMq_bf16_subchunk(i); + tArA_r_j(i) = + sAlphaLoadCurr(int(r_) * kSubchunkRows + int(row), work_desc.o_head_idx() & 3); + }); + } else { + Tensor sAqkq_r_j = sAqkq_slice(_, _, r_, make_coord(_0{}, j)); + Tensor tAsA_r_j = alpha_Q_bf16_thr_copy.partition_S(sAqkq_r_j); + Tensor tArA_r_j_cv = alpha_Q_bf16_thr_copy.retile_D(tArA_r_j); + copy(alpha_Q_bf16_tiled_copy, tAsA_r_j, tArA_r_j_cv); + } // Derive g_first (alpha[row=0, :]) from tArA_r_j via warp shuffle, // directly into operand B layout (8 values instead of 16). @@ -1577,7 +1701,14 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // v1=0 subset of operand A. We shuffle v1=0 values from t1=0 thread and // output directly as operand B fragment, saving 8 float registers. Tensor tArAfirst_r_j_kt = make_fragment_like(tv_layout_bf16_mma_B); - broadcast_row0_operandA_to_operandB_bf16_layout(tArA_r_j, tArAfirst_r_j_kt, local_thread_idx); + if constexpr (ScalarAlpha) { + const float g_first = sAlphaLoadCurr( + int(r_) * size<0>(TileShape_SubChunk{}), work_desc.o_head_idx() & 3); + fill(tArAfirst_r_j_kt, g_first); + } else { + broadcast_row0_operandA_to_operandB_bf16_layout( + tArA_r_j, tArAfirst_r_j_kt, local_thread_idx); + } // gqn_r_j = exp2(g_r_j - g_r_j_first[None, :]) in BF16 MMA A layout. // g_first per k-iter is in tArAfirst_r_j_kt: frag_B(2j)=K_lo, frag_B(2j+1)=K_hi. @@ -1635,11 +1766,20 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto s2r_compute_subchunk_operandB = [&](auto c_, int j, int j0, int j1, auto const& tArAfirst_kt) INLINE_LAMBDA { // S2R g_c_j in BF16 MMA operand B layout - Tensor sAqkq_c_j = sAqkq_slice(_, _, c_, make_coord(_0{}, j)); - Tensor tAsA_c_j = alpha_Kt_bf16_thr_copy.partition_S(sAqkq_c_j); Tensor tArA_c_j = make_fragment_like(tv_layout_bf16_mma_B); - Tensor tArA_c_j_cv = alpha_Kt_bf16_thr_copy.retile_D(tArA_c_j); - copy(alpha_Kt_bf16_tiled_copy, tAsA_c_j, tArA_c_j_cv); + if constexpr (ScalarAlpha) { + constexpr int kSubchunkCols = size<1>(TileShape_SubChunk{}); + for_each(make_int_sequence{}, [&](auto i) { + auto [col, _] = tKcNk_bf16_subchunk(i); + tArA_c_j(i) = + sAlphaLoadCurr(int(c_) * kSubchunkCols + int(col), work_desc.o_head_idx() & 3); + }); + } else { + Tensor sAqkq_c_j = sAqkq_slice(_, _, c_, make_coord(_0{}, j)); + Tensor tAsA_c_j = alpha_Kt_bf16_thr_copy.partition_S(sAqkq_c_j); + Tensor tArA_c_j_cv = alpha_Kt_bf16_thr_copy.retile_D(tArA_c_j); + copy(alpha_Kt_bf16_tiled_copy, tAsA_c_j, tArA_c_j_cv); + } // compute gktn_c_j = exp2(g_first - g_c_j) in BF16 MMA B layout cute::transform( diff --git a/csrc/kda/sm90/collective/store_tma.hpp b/csrc/kda/sm90/collective/store_tma.hpp index 0f7f7c1a..e2349f6a 100644 --- a/csrc/kda/sm90/collective/store_tma.hpp +++ b/csrc/kda/sm90/collective/store_tma.hpp @@ -184,7 +184,6 @@ struct CollectiveStoreTma { CUTE_DEVICE auto partition_SD(ProblemSize const& problem_size, TileShape const& tile_shape, WorkDesc const& work_desc) { constexpr auto BlkSeqQ = decltype(get<0>(tile_shape))::value; - constexpr auto HeadSize = decltype(get<2>(tile_shape))::value; Tensor g = [&] { DPRINTF0_W( @@ -198,10 +197,10 @@ struct CollectiveStoreTma { problem_size.num_v_heads)); // O lives in the V/O head space under GVA Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx Tensor m_offset = domain_offset( - make_coord(_0{}, work_desc.tok_offset), + make_coord(work_desc.value_tile_idx * int(SizeM{}), work_desc.tok_offset), m_varlen); // offset to start of the current sequence Tensor g_full = - local_tile(m_offset, make_tile(HeadSize, BlkSeqQ), make_coord(_0{}, _)); // (d, blk, iter_blk) + local_tile(m_offset, make_tile(SizeM{}, BlkSeqQ), make_coord(_0{}, _)); // (d, blk, iter_blk) return g_full; }(); Tensor s = make_tensor(make_smem_ptr(storage_.data()), SmemLayoutO{}); diff --git a/csrc/kda/sm90/kda_fwd_sm90.cu b/csrc/kda/sm90/kda_fwd_sm90.cu index ed855db6..bfa7ff0a 100644 --- a/csrc/kda/sm90/kda_fwd_sm90.cu +++ b/csrc/kda/sm90/kda_fwd_sm90.cu @@ -22,6 +22,7 @@ namespace kda::sm90 { using namespace cute; +using bf16 = cute::bfloat16_t; // Forward declaration of the per-variant launcher (defined in .cuh, instantiated in separate TUs) template < @@ -33,7 +34,10 @@ template < typename TO, typename TQKV, typename TState, - typename TBeta = float> + typename TBeta = float, + bool ScalarAlpha = false, + bool StateKVLayout = false, + bool SplitValueDim = false> void launch_kda_fwd_prefill_kernel_gbai( cudaStream_t stream, @@ -58,6 +62,78 @@ launch_kda_fwd_prefill_kernel_gbai( int32_t const* raw_cu_seqlens, int32_t raw_num_seqs); +void +launch_qwen35_scalar_kda_fwd_prefill_kernel( + cudaStream_t stream, + void* output, + float* output_state, + void const* q, + void const* k, + void const* v, + float const* input_state, + float const* alpha, + float const* beta, + int32_t const* cu_seqlens, + uint8_t* workspace_buffer, + int32_t num_seqs, + int32_t num_qk_heads, + int32_t num_v_heads, + int32_t head_size, + int64_t total_seqlen, + float scale, + bool has_initial_state, + int32_t sm_count) { + if (has_initial_state) { + launch_kda_fwd_prefill_kernel_gbai< + true, true, true, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + stream, + static_cast(output), + output_state, + static_cast(q), + static_cast(k), + static_cast(v), + input_state, + alpha, + beta, + cu_seqlens, + workspace_buffer, + num_seqs, + num_qk_heads, + num_v_heads, + head_size, + total_seqlen, + scale, + sm_count, + nullptr, + nullptr, + num_seqs); + } else { + launch_kda_fwd_prefill_kernel_gbai< + true, true, false, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + stream, + static_cast(output), + output_state, + static_cast(q), + static_cast(k), + static_cast(v), + nullptr, + alpha, + beta, + cu_seqlens, + workspace_buffer, + num_seqs, + num_qk_heads, + num_v_heads, + head_size, + total_seqlen, + scale, + sm_count, + nullptr, + nullptr, + num_seqs); + } +} + template < typename ArchTag, // TODO: hide this typename TO, @@ -132,8 +208,6 @@ launch_kda_fwd_prefill_kernel( #undef LAUNCH } -using bf16 = cute::bfloat16_t; - // TBeta=float (default) template void launch_kda_fwd_prefill_kernel( diff --git a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu index 0da2986e..693cc406 100644 --- a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu +++ b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu @@ -63,4 +63,57 @@ INSTANTIATE_GBAI(true, true, true, true, bf16); #undef INSTANTIATE_GBAI +// Qwen scalar-G specialization: compact alpha [T, HV] and external state +// buffers in contiguous [K, V] layout. Keep this separate from the generic +// vector-alpha instantiations above. +template void +launch_kda_fwd_prefill_kernel_gbai< + true, true, false, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + cudaStream_t, + bf16*, + float*, + bf16 const*, + bf16 const*, + bf16 const*, + float const*, + float const*, + float const*, + int32_t const*, + uint8_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int64_t, + float, + int32_t, + int32_t const*, + int32_t const*, + int32_t); + +template void +launch_kda_fwd_prefill_kernel_gbai< + true, true, true, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + cudaStream_t, + bf16*, + float*, + bf16 const*, + bf16 const*, + bf16 const*, + float const*, + float const*, + float const*, + int32_t const*, + uint8_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int64_t, + float, + int32_t, + int32_t const*, + int32_t const*, + int32_t); + } // namespace kda::sm90 diff --git a/csrc/kda/sm90/kernel/builder_kda_fwd.hpp b/csrc/kda/sm90/kernel/builder_kda_fwd.hpp index 74e9c43a..cc9a606f 100644 --- a/csrc/kda/sm90/kernel/builder_kda_fwd.hpp +++ b/csrc/kda/sm90/kernel/builder_kda_fwd.hpp @@ -70,7 +70,7 @@ struct FlatBuilderKdaFwd< static constexpr bool kIsPersistent = find_option_t::value; static_assert(!kIsPersistent, "not implemented"); - using TileScheduler = kda::sm90::kernel::IndividualTileScheduler; + using TileScheduler = kda::sm90::kernel::IndividualTileScheduler; // using TileScheduler = std::conditional_t; diff --git a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp b/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp index ac597f7b..ef326859 100644 --- a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp +++ b/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp @@ -198,7 +198,11 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { get_register_requirements(MaxThreadsPerBlock, MinBlocksPerMultiprocessor, NumStateMmaWarpGroups); static constexpr uint32_t LdStRegisterRequirement = get<0>(RegisterRequirements); static constexpr uint32_t StateMmaRegisterRequirement = get<1>(RegisterRequirements); - static constexpr uint32_t AuxMmaRegisterRequirement = get<2>(RegisterRequirements); + // The V64 specialization statically uses 168 registers/thread. setmaxnreg + // `.inc 152` is illegal when the requested aux ceiling is below that + // initial allocation, so keep the aux WG slightly above the static value. + static constexpr uint32_t AuxMmaRegisterRequirement = + NumStateMmaWarpGroups == 1 ? 176 : get<2>(RegisterRequirements); static size_t get_workspace_size(Arguments const& args) { @@ -236,13 +240,6 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { CUTE_DEVICE void operator()(const Params& params, char* smem) { - enum class WarpGroupRole { - LdSt = 0, - Math0 = 1, - Math1 = 2, - MathA = 3, // auxiliary math WG - }; - // NOTE: CollectiveInverse will have more utilization on warp 0&1 // so we put beta and alpha preprocessing on warp 2&3 enum class LdStWarpRole { @@ -261,7 +258,10 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { int warp_idx = cutlass::canonical_warp_idx_sync(); int warp_idx_in_wg = warp_idx % cutlass::NumWarpsPerWarpGroup; int warp_group_idx = cutlass::canonical_warp_group_idx(); - auto warp_group_role = WarpGroupRole(warp_group_idx); + bool is_load_wg = warp_group_idx == 0; + bool is_state_wg = + warp_group_idx >= 1 && warp_group_idx < 1 + NumStateMmaWarpGroups; + bool is_aux_wg = warp_group_idx == 1 + NumStateMmaWarpGroups; auto ldst_warp_role = LdStWarpRole(warp_idx_in_wg); int lane_predicate = cute::elect_one_sync(); @@ -323,7 +323,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { OrderedMathBarriers math_barriers; - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadQKV) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::LoadQKV) { DPRINTF0_W("ldst_warp_role: LoadQKV Alpha\n"); q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Producer; k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Producer; @@ -332,23 +332,23 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Producer; } } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::StoreO) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::StoreO) { DPRINTF0_W("ldst_warp_role: StoreO\n"); o_pipeline_params.role = MainloopOPipeline::ThreadCategory::Consumer; } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadBeta) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::LoadBeta) { if constexpr (NeedsBeta) { beta_pipeline_params.role = MainloopBetaPipeline::ThreadCategory::Producer; } } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadAlpha) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::LoadAlpha) { // LoadAlpha warp consumes alpha_pipeline (reads last row) and produces alpha_last_pipeline if constexpr (NeedsAlpha) { alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Consumer; } alpha_last_pipeline_params.role = MainloopAlphaLastPipeline::ThreadCategory::Producer; } - if (warp_group_role == WarpGroupRole::Math0 || warp_group_role == WarpGroupRole::Math1) { + if (is_state_wg) { DPRINTF0_WG("warp_group_role: MathX\n"); q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Consumer; k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Consumer; @@ -368,7 +368,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { math_barriers.init(warp_group_idx - 1); } - if (warp_group_role == WarpGroupRole::MathA) { + if (is_aux_wg) { DPRINTF0_WG("warp_group_role: MathA\n"); q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Consumer; k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Consumer; @@ -453,7 +453,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { CollectiveMainloop collective_mainloop; - if (warp_group_role == WarpGroupRole::LdSt) { + if (is_load_wg) { DPRINTF0_WG("LsSt warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, LdStRegisterRequirement); cutlass::arch::warpgroup_reg_dealloc(); if (ldst_warp_role == LdStWarpRole::LoadQKV) { @@ -549,7 +549,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { o_smem_pipe_read, storage.tensors.mainloop.smem_o); } - } else if (warp_group_role == WarpGroupRole::Math0 || warp_group_role == WarpGroupRole::Math1) { + } else if (is_state_wg) { DPRINTF0_WG( "Compute[state]: warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, @@ -592,7 +592,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { math_barriers, storage.tensors.mainloop); } - } else if (warp_group_role == WarpGroupRole::MathA) { + } else if (is_aux_wg) { DPRINTF0_WG( "Compute[aux]: warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, AuxMmaRegisterRequirement); cutlass::arch::warpgroup_reg_alloc(); diff --git a/csrc/kda/sm90/kernel/options.hpp b/csrc/kda/sm90/kernel/options.hpp index e25fe9d4..cbd93311 100644 --- a/csrc/kda/sm90/kernel/options.hpp +++ b/csrc/kda/sm90/kernel/options.hpp @@ -82,6 +82,9 @@ enum class Tag { kInitStateFromInput, // if true, initialize state by reading global memory instead of zero initialization. kSafeGate, // KDA kElementBetaGmem, // GMEM element type for beta (default float, can be bf16) + kScalarAlpha, // Qwen GDN: one gate value per token/V head, broadcast over D + kStateKVLayout, // Qwen adapter: external state is contiguous [K, V] + kSplitValueDim, // Qwen SM90: split the V=128 feature dimension into two V64 CTAs }; } // namespace kda::sm90::kernel diff --git a/csrc/kda/sm90/kernel/tile_scheduler.hpp b/csrc/kda/sm90/kernel/tile_scheduler.hpp index c70c7a63..4d690d4a 100644 --- a/csrc/kda/sm90/kernel/tile_scheduler.hpp +++ b/csrc/kda/sm90/kernel/tile_scheduler.hpp @@ -27,6 +27,7 @@ struct WorkDesc { int32_t seq_idx; // which sequence to process int32_t qk_head_idx; // head idx for Q/K (the representative of the GVA group) int32_t head_idx; // head idx for V/O/g/beta + int32_t value_tile_idx; // V-feature tile within one value head int64_t tok_offset; // start offset of this sequence in the packed tensor // shape @@ -65,10 +66,12 @@ struct WorkDesc { } }; -// Each block handles a single (seq, v_head) work item; CTAs do not cooperate. +// Each block handles one (seq, v_head, value_tile) work item; CTAs do not cooperate. // GVA optimization: heads_per_group is precomputed on the host and stored in // Params, so the device side does not redo the integer division per CTA. +template struct IndividualTileScheduler { + static_assert(NumValueTiles >= 1); struct Params { dim3 grid; int32_t num_seqs; @@ -93,7 +96,7 @@ struct IndividualTileScheduler { // the integer division. int32_t const heads_per_group = problem_size.num_v_heads / problem_size.num_qk_heads; dim3 grid(0, 1, 1); - grid.x = problem_size.num_seqs * problem_size.num_v_heads; + grid.x = problem_size.num_seqs * problem_size.num_v_heads * NumValueTiles; DPRINTF( "to_underlying_arguments: grid:{.x:%d, .y:%d, .z:%d}, num_seqs:%d, num_qk_heads:%d, num_v_heads:%d, " "heads_per_group:%d\n", @@ -120,8 +123,10 @@ struct IndividualTileScheduler { template CUTE_DEVICE WorkDesc get_next_work(Params params, ProblemSize const& problem_size) { - int32_t seq_idx = blockIdx.x / params.num_v_heads; - int32_t head_idx = blockIdx.x % params.num_v_heads; + int32_t value_tile_idx = blockIdx.x % NumValueTiles; + int32_t work_idx = blockIdx.x / NumValueTiles; + int32_t seq_idx = work_idx / params.num_v_heads; + int32_t head_idx = work_idx % params.num_v_heads; // GVA: use the host-precomputed heads_per_group to avoid device-side division. int32_t qk_head_idx = head_idx / params.heads_per_group; @@ -134,10 +139,12 @@ struct IndividualTileScheduler { } else { scheduled = true; DPRINTF0_W( - "get_next_work: this_work={seq_idx:%d qk_head_idx:%d head_idx:%d tok_offset:%lld seq_len:%lld}\n", + "get_next_work: this_work={seq_idx:%d qk_head_idx:%d head_idx:%d value_tile_idx:%d " + "tok_offset:%lld seq_len:%lld}\n", seq_idx, qk_head_idx, head_idx, + value_tile_idx, s, seq_len); } @@ -146,6 +153,7 @@ struct IndividualTileScheduler { .seq_idx = seq_idx, .qk_head_idx = qk_head_idx, .head_idx = head_idx, + .value_tile_idx = value_tile_idx, .tok_offset = s, .seq_len = seq_len, }; diff --git a/csrc/kda/sm90/prefill_kernel.hpp b/csrc/kda/sm90/prefill_kernel.hpp index 6e54c3a3..94341ee6 100644 --- a/csrc/kda/sm90/prefill_kernel.hpp +++ b/csrc/kda/sm90/prefill_kernel.hpp @@ -51,4 +51,29 @@ launch_kda_fwd_prefill_kernel( int32_t const* raw_cu_seqlens = nullptr, int32_t raw_num_seqs = 0); +// Qwen GDN-only specialization. Alpha is compact chunk-prefix log2 gate +// [packed_tokens, num_v_heads], and state buffers use contiguous [K, V]. +// The generic vector-alpha public API above remains unchanged. +void +launch_qwen35_scalar_kda_fwd_prefill_kernel( + cudaStream_t stream, + void* output, + float* output_state, + void const* q, + void const* k, + void const* v, + float const* input_state, + float const* alpha, + float const* beta, + int32_t const* cu_seqlens, + uint8_t* workspace_buffer, + int32_t num_seqs, + int32_t num_qk_heads, + int32_t num_v_heads, + int32_t head_size, + int64_t total_seqlen, + float scale, + bool has_initial_state, + int32_t sm_count); + } // namespace kda::sm90 diff --git a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh index c53f2ae3..8a482372 100644 --- a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh +++ b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh @@ -38,7 +38,10 @@ template < typename TO, typename TQKV, typename TState, - typename TBeta = float> + typename TBeta = float, + bool ScalarAlpha = false, + bool StateKVLayout = false, + bool SplitValueDim = false> void launch_kda_fwd_prefill_kernel_gbai( cudaStream_t stream, @@ -81,17 +84,21 @@ launch_kda_fwd_prefill_kernel_gbai( using NeedsBetaType = std::conditional_t; using NeedsAlphaType = std::conditional_t; using InitStateType = std::conditional_t; + using Options0 = decltype(add_option(Option{}, DefaultOptions{})); + using Options1 = decltype(add_option(Option{}, Options0{})); + using Options2 = decltype(add_option(Option{}, Options1{})); + using Options3 = decltype(add_option(Option{}, Options2{})); + using Options4 = decltype(add_option(Option{}, Options3{})); + using Options5 = decltype(add_option(Option{}, Options4{})); + using Options6 = decltype(add_option( + Option>{}, + Options5{})); + using Options7 = decltype(add_option( + Option>{}, + Options6{})); using Options = decltype(add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option(Option{}, DefaultOptions{}))))))); + Option>{}, + Options7{})); using TileShape = Shape<_64, _64, _128>; using Scheduler = cutlass::gemm::KernelTmaWarpSpecializedCooperative; @@ -137,7 +144,11 @@ launch_kda_fwd_prefill_kernel_gbai( .ptr_K = (T*)k, .dK = {qk_tok_stride, _1{}, head_stride}, .ptr_V = (T*)v, .dV = {v_tok_stride, _1{}, head_stride}, .ptr_O = (T*)output, .dO = {v_tok_stride, _1{}, head_stride}, - .ptr_Alpha = alpha, .dAlpha = {v_tok_stride, _1{}, head_stride}, + .ptr_Alpha = alpha, + .dAlpha = { + ScalarAlpha ? int64_t(num_v_heads) : int64_t(v_tok_stride), + _1{}, + ScalarAlpha ? int32_t(4) : head_stride}, .ptr_output_state = (float*)output_state, .ptr_input_state = (float*)input_state, .scale = scale, diff --git a/csrc/qwen35/decode/qwen35_conv1d_decode.cu b/csrc/qwen35/decode/qwen35_conv1d_decode.cu new file mode 100644 index 00000000..5a3c3470 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_conv1d_decode.cu @@ -0,0 +1,192 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#include "qwen35_decode_common.cuh" + +#include +#include +#include +#include +#include +#include + +namespace { + +template +__device__ inline float to_float(T x) { + return static_cast(x); +} + +template <> +__device__ inline float to_float(c10::Half x) { + return __half2float(static_cast<__half>(x)); +} + +template <> +__device__ inline float to_float(c10::BFloat16 x) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return __bfloat162float(static_cast<__nv_bfloat16>(x)); +#else + return static_cast(x); +#endif +} + +template +__device__ inline T from_float(float x) { + return static_cast(x); +} + +template <> +__device__ inline c10::Half from_float(float x) { + return c10::Half(__float2half_rn(x)); +} + +template <> +__device__ inline c10::BFloat16 from_float(float x) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return c10::BFloat16(__float2bfloat16(x)); +#else + return c10::BFloat16(x); +#endif +} + +template +__global__ void qwen35_conv1d_decode_kernel( + const scalar_t* __restrict__ mixed_qkv, + scalar_t* __restrict__ conv_state, + const scalar_t* __restrict__ conv_weight, + scalar_t* __restrict__ out, + int batch_size, + int conv_dim) { + constexpr int kThreads = 256; + const int64_t linear_idx = static_cast(blockIdx.x) * kThreads + threadIdx.x; + const int64_t total = static_cast(batch_size) * conv_dim; + if (linear_idx >= total) { + return; + } + + const int64_t b = linear_idx / conv_dim; + const int64_t c = linear_idx % conv_dim; + + const int64_t x_idx = b * conv_dim + c; + const int64_t state_base = + (b * conv_dim + c) * cula::qwen35::decode::kConvKernelSize; + const int64_t weight_base = c * cula::qwen35::decode::kConvKernelSize; + + const float s0 = to_float(conv_state[state_base + 1]); + const float s1 = to_float(conv_state[state_base + 2]); + const float s2 = to_float(conv_state[state_base + 3]); + const float s3 = to_float(mixed_qkv[x_idx]); + + const float w0 = to_float(conv_weight[weight_base + 0]); + const float w1 = to_float(conv_weight[weight_base + 1]); + const float w2 = to_float(conv_weight[weight_base + 2]); + const float w3 = to_float(conv_weight[weight_base + 3]); + + const float conv = s0 * w0 + s1 * w1 + s2 * w2 + s3 * w3; + const float silu = conv / (1.f + expf(-conv)); + + conv_state[state_base + 0] = from_float(s0); + conv_state[state_base + 1] = from_float(s1); + conv_state[state_base + 2] = from_float(s2); + conv_state[state_base + 3] = from_float(s3); + out[x_idx] = from_float(silu); +} + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +} // namespace + +namespace cula::qwen35::decode { + +void run_qwen35_conv1d_decode(ConvDecodeParams& params) { + const at::Tensor& mixed_qkv = params.mixed_qkv; + const at::Tensor& conv_state = params.conv_state; + const at::Tensor& conv_weight = params.conv_weight; + const at::Tensor& out = params.out; + + TORCH_CHECK(mixed_qkv.is_cuda(), "mixed_qkv must be a CUDA tensor."); + const at::Device device = mixed_qkv.device(); + + check_tensor_device(conv_state, "conv_state", device); + check_tensor_device(conv_weight, "conv_weight", device); + check_tensor_device(out, "out", device); + + TORCH_CHECK(mixed_qkv.is_contiguous(), "mixed_qkv must be contiguous."); + TORCH_CHECK(conv_state.is_contiguous(), "conv_state must be contiguous."); + TORCH_CHECK(conv_weight.is_contiguous(), "conv_weight must be contiguous."); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous."); + + TORCH_CHECK( + mixed_qkv.scalar_type() == conv_state.scalar_type() && + mixed_qkv.scalar_type() == conv_weight.scalar_type() && + mixed_qkv.scalar_type() == out.scalar_type(), + "mixed_qkv/conv_state/conv_weight/out must share the same dtype."); + + TORCH_CHECK( + mixed_qkv.scalar_type() == at::kHalf || mixed_qkv.scalar_type() == at::kBFloat16, + "conv decode only supports half/bfloat16."); + + const int64_t batch_size = mixed_qkv.size(0); + const int64_t conv_dim = mixed_qkv.size(2); + TORCH_CHECK(conv_dim > 0, "conv_dim must be positive."); + TORCH_CHECK( + mixed_qkv.dim() == 3 && mixed_qkv.sizes() == at::IntArrayRef({batch_size, 1, conv_dim}), + "mixed_qkv must have shape [B, 1, local_conv_dim]."); + TORCH_CHECK( + conv_state.dim() == 3 && + conv_state.sizes() == at::IntArrayRef({batch_size, conv_dim, kConvKernelSize}), + "conv_state must have shape [B, local_conv_dim, 4]."); + TORCH_CHECK( + (conv_weight.dim() == 2 && conv_weight.sizes() == at::IntArrayRef({conv_dim, kConvKernelSize})) || + (conv_weight.dim() == 3 && + conv_weight.sizes() == at::IntArrayRef({conv_dim, 1, kConvKernelSize})), + "conv_weight must have shape [local_conv_dim, 4] or [local_conv_dim, 1, 4]."); + TORCH_CHECK( + out.dim() == 3 && out.sizes() == at::IntArrayRef({batch_size, 1, conv_dim}), + "out must have shape [B, 1, local_conv_dim]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + const at::Tensor mixed_qkv_2d = mixed_qkv.view({batch_size, conv_dim}); + const at::Tensor out_2d = out.view({batch_size, conv_dim}); + const at::Tensor weight_2d = + conv_weight.dim() == 3 ? conv_weight.view({conv_dim, kConvKernelSize}) : conv_weight; + + constexpr int kThreads = 256; + const int64_t total = batch_size * conv_dim; + const dim3 block(kThreads, 1, 1); + const dim3 grid(static_cast((total + kThreads - 1) / kThreads), 1, 1); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + mixed_qkv.scalar_type(), + "qwen35_conv1d_decode_kernel", + [&] { + qwen35_conv1d_decode_kernel<<>>( + mixed_qkv_2d.data_ptr(), + conv_state.data_ptr(), + weight_2d.data_ptr(), + out_2d.data_ptr(), + static_cast(batch_size), + static_cast(conv_dim)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_decode_common.cuh b/csrc/qwen35/decode/qwen35_decode_common.cuh new file mode 100644 index 00000000..51ce1872 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_decode_common.cuh @@ -0,0 +1,132 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include +#include +#include + +namespace cula::qwen35::decode { + +inline constexpr int kNumQKHeads = 16; +inline constexpr int kNumVHeads = 48; +inline constexpr int kHeadDimQK = 128; +inline constexpr int kHeadDimV = 128; +inline constexpr int kConvKernelSize = 4; +inline constexpr int kQDim = kNumQKHeads * kHeadDimQK; +inline constexpr int kKDim = kNumQKHeads * kHeadDimQK; +inline constexpr int kVDim = kNumVHeads * kHeadDimV; +inline constexpr int kMixedQKVDim = kQDim + kKDim + kVDim; + +inline constexpr int local_qk_heads_from_v_heads(int local_v_heads) { + return local_v_heads / (kNumVHeads / kNumQKHeads); +} + +inline constexpr int local_q_dim(int local_qk_heads) { + return local_qk_heads * kHeadDimQK; +} + +inline constexpr int local_v_dim(int local_v_heads) { + return local_v_heads * kHeadDimV; +} + +inline constexpr int local_mixed_qkv_dim(int local_qk_heads, int local_v_heads) { + return 2 * local_q_dim(local_qk_heads) + local_v_dim(local_v_heads); +} + +inline constexpr bool is_supported_local_v_heads(int local_v_heads) { + return local_v_heads == 48 || local_v_heads == 24 || local_v_heads == 12 || local_v_heads == 6; +} + +template +struct Qwen35DecodeLocalShape { + static_assert(is_supported_local_v_heads(kLocalVHeads_), "Unsupported Qwen3.5 local V-head count."); + static constexpr int kLocalVHeads = kLocalVHeads_; + static constexpr int kLocalQKHeads = local_qk_heads_from_v_heads(kLocalVHeads); + static constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; + static constexpr int kLocalQDim = local_q_dim(kLocalQKHeads); + static constexpr int kLocalKDim = local_q_dim(kLocalQKHeads); + static constexpr int kLocalVDim = local_v_dim(kLocalVHeads); + static constexpr int kLocalMixedQKVDim = local_mixed_qkv_dim(kLocalQKHeads, kLocalVHeads); + + // Decode shape policy. Head dimension is fixed at 128 for Qwen3.5, but keep + // these knobs with the local-head traits so future TP-shape tuning has one + // place to specialize. + static constexpr int kLayoutVec = 4; + static constexpr int kLayoutThreads = kHeadDimQK / kLayoutVec; + static constexpr int kKdaThreads = 128; + static constexpr int kKdaTileV = 16; + static constexpr int kKdaTileK = 16; + + static_assert(kLocalVHeads % kLocalQKHeads == 0); + static_assert(kHeadDimQK == kHeadDimV); + static_assert(kHeadDimQK % kLayoutVec == 0); + static_assert(kHeadDimV % kKdaTileV == 0); + static_assert(kHeadDimQK % kKdaTileK == 0); +}; + +struct ConvDecodeParams { + at::Tensor mixed_qkv; // [B, 1, local_conv_dim] + at::Tensor conv_state; // [B, local_conv_dim, 4] + at::Tensor conv_weight; // [local_conv_dim, 4] + at::Tensor out; // [B, 1, local_conv_dim] +}; + +struct LayoutDecodeParams { + at::Tensor mixed_qkv_conv; // [N, local_conv_dim] + at::Tensor a; // [N, local_v_heads] + at::Tensor b; // [N, local_v_heads] + at::Tensor q_rep; // [N, local_v_heads, 128] + at::Tensor k_rep; // [N, local_v_heads, 128] + at::Tensor v; // [N, local_v_heads, 128] + at::Tensor a_kernel; // [N, local_v_heads] + at::Tensor b_kernel; // [N, local_v_heads] +}; + +struct ScalarKdaDecodeParams { + // Dtype contract for the first implementation: + // - activations / outputs: half or bf16 + // q_rep, k_rep, v, a_kernel, b_kernel, out + // - recurrent parameters / state: float32 + // A_log, dt_bias, recurrent_state + at::Tensor q_rep; // [N, local_v_heads, 128] + at::Tensor k_rep; // [N, local_v_heads, 128] + at::Tensor v; // [N, local_v_heads, 128] + at::Tensor a_kernel; // [N, local_v_heads] + at::Tensor b_kernel; // [N, local_v_heads] + at::Tensor A_log; // [local_v_heads], float32 + at::Tensor dt_bias; // [local_v_heads], float32 + at::Tensor recurrent_state; // [pool, local_v_heads, 128, 128], float32 + at::Tensor pool_idx; // [N], int32 + at::Tensor out; // [N, local_v_heads, 128] +}; + +struct LayoutScalarKdaDecodeParams { + at::Tensor mixed_qkv_conv; // [N, local_conv_dim] + at::Tensor a; // [N, local_v_heads] + at::Tensor b; // [N, local_v_heads] + at::Tensor A_log; // [local_v_heads], float32 + at::Tensor dt_bias; // [local_v_heads], float32 + at::Tensor recurrent_state; // [pool, local_v_heads, 128, 128], float32 + at::Tensor pool_idx; // [N], int32 + at::Tensor out; // [N, local_v_heads, 128] +}; + +void run_qwen35_conv1d_decode(ConvDecodeParams& params); +void run_qwen35_layout_decode(LayoutDecodeParams& params); +void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params); +void run_qwen35_layout_scalar_kda_decode(LayoutScalarKdaDecodeParams& params); + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_layout_decode.cu b/csrc/qwen35/decode/qwen35_layout_decode.cu new file mode 100644 index 00000000..811e98b9 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_layout_decode.cu @@ -0,0 +1,167 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#include "qwen35_decode_common.cuh" +#include "qwen35_layout_kernel.hpp" + +#include +#include +#include +#include +#include +#include + +namespace { + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +void check_tensor_shape_2d(const at::Tensor& tensor, const char* name) { + TORCH_CHECK( + tensor.dim() == 2, + name, + " must have rank 2, but got rank ", + tensor.dim(), + "."); +} + +template +void launch_layout_decode_for_heads( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t batch_size) { + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + dim3 grid(Shape::kLocalVHeads, static_cast(batch_size), 1); + cula::qwen35::decode::qwen35_layout_decode_kernel_cute + <<>>( + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + batch_size); +} + +} // namespace + +namespace cula::qwen35::decode { + +void run_qwen35_layout_decode(LayoutDecodeParams& params) { + const at::Tensor& mixed_qkv_conv = params.mixed_qkv_conv; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& q_rep = params.q_rep; + const at::Tensor& k_rep = params.k_rep; + const at::Tensor& v = params.v; + const at::Tensor& a_kernel = params.a_kernel; + const at::Tensor& b_kernel = params.b_kernel; + + TORCH_CHECK(mixed_qkv_conv.is_cuda(), "mixed_qkv_conv must be a CUDA tensor."); + TORCH_CHECK(mixed_qkv_conv.is_contiguous(), "mixed_qkv_conv must be contiguous."); + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == a.scalar_type() && + mixed_qkv_conv.scalar_type() == b.scalar_type() && + mixed_qkv_conv.scalar_type() == q_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == k_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == v.scalar_type() && + mixed_qkv_conv.scalar_type() == a_kernel.scalar_type() && + mixed_qkv_conv.scalar_type() == b_kernel.scalar_type(), + "All layout decode tensors must share the same dtype."); + + check_tensor_shape_2d(a, "a"); + check_tensor_shape_2d(b, "b"); + + const int64_t batch_size = mixed_qkv_conv.size(0); + const int64_t local_v_heads = a.size(1); + TORCH_CHECK(is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + const int local_qk_heads = local_qk_heads_from_v_heads(static_cast(local_v_heads)); + const int local_mixed_dim = local_mixed_qkv_dim(local_qk_heads, static_cast(local_v_heads)); + TORCH_CHECK( + mixed_qkv_conv.dim() == 2 && mixed_qkv_conv.size(1) == local_mixed_dim, + "mixed_qkv_conv must have shape [N, local_conv_dim=", local_mixed_dim, "], got ", + mixed_qkv_conv.sizes(), "."); + const at::Device device = mixed_qkv_conv.device(); + + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(q_rep, "q_rep", device); + check_tensor_device(k_rep, "k_rep", device); + check_tensor_device(v, "v", device); + check_tensor_device(a_kernel, "a_kernel", device); + check_tensor_device(b_kernel, "b_kernel", device); + + TORCH_CHECK(q_rep.is_contiguous(), "q_rep must be contiguous."); + TORCH_CHECK(k_rep.is_contiguous(), "k_rep must be contiguous."); + TORCH_CHECK(v.is_contiguous(), "v must be contiguous."); + TORCH_CHECK(a_kernel.is_contiguous(), "a_kernel must be contiguous."); + TORCH_CHECK(b_kernel.is_contiguous(), "b_kernel must be contiguous."); + + TORCH_CHECK( + q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({batch_size, local_v_heads, kHeadDimQK}), + "q_rep must have shape [N, local_v_heads, 128]."); + TORCH_CHECK( + k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({batch_size, local_v_heads, kHeadDimQK}), + "k_rep must have shape [N, local_v_heads, 128]."); + TORCH_CHECK( + v.dim() == 3 && v.sizes() == at::IntArrayRef({batch_size, local_v_heads, kHeadDimV}), + "v must have shape [N, local_v_heads, 128]."); + TORCH_CHECK( + a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({batch_size, local_v_heads}), + "a_kernel must have shape [N, local_v_heads]."); + TORCH_CHECK( + b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({batch_size, local_v_heads}), + "b_kernel must have shape [N, local_v_heads]."); + + TORCH_CHECK(a.sizes() == at::IntArrayRef({batch_size, local_v_heads}), "a must have shape [N, local_v_heads]."); + TORCH_CHECK(b.sizes() == at::IntArrayRef({batch_size, local_v_heads}), "b must have shape [N, local_v_heads]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + mixed_qkv_conv.scalar_type(), + "qwen35_layout_decode_kernel_cute", + [&] { + switch (local_v_heads) { + case 48: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + case 24: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + case 12: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + case 6: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_layout_kernel.hpp b/csrc/qwen35/decode/qwen35_layout_kernel.hpp new file mode 100644 index 00000000..e43ffc00 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_layout_kernel.hpp @@ -0,0 +1,131 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include "qwen35_decode_common.cuh" + +#include +#include +#include + +namespace cula::qwen35::decode { + +using namespace cute; + +template +CUTE_DEVICE void copy_vec_contiguous( + scalar_t* __restrict__ dst, + const scalar_t* __restrict__ src) { + constexpr int kBytes = sizeof(scalar_t) * kVec; + if constexpr (kBytes == 16 || kBytes == 8) { + using VecType = cutlass::AlignedArray; + auto dst_addr = reinterpret_cast(dst); + auto src_addr = reinterpret_cast(src); + if ((dst_addr % alignof(VecType) == 0) && (src_addr % alignof(VecType) == 0)) { + *reinterpret_cast(dst) = *reinterpret_cast(src); + return; + } + } + +#pragma unroll + for (int i = 0; i < kVec; ++i) { + dst[i] = src[i]; + } +} + +template +__global__ void qwen35_layout_decode_kernel_cute( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + scalar_t* __restrict__ q_rep, + scalar_t* __restrict__ k_rep, + scalar_t* __restrict__ v_out, + scalar_t* __restrict__ a_kernel, + scalar_t* __restrict__ b_kernel, + int64_t token_count) { + using Shape = Qwen35DecodeLocalShape; + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; + // TODO(qwen35-layout-opt): + // - Re-evaluate whether Vec=8 is profitable for bf16/fp16 on the target GPUs. + // - Push more of the q/k repeat mapping into compile-time CuTe layout transforms. + // - Revisit whether a shared-memory staging path is worthwhile after profiling. + // - Consider widening the a/b writeback path if it shows up in profiling. + constexpr int kVec = Shape::kLayoutVec; + static_assert(kHeadDimV % kVec == 0); + static_assert(kHeadDimQK == kHeadDimV); + static_assert(kHeadDimQK % kVec == 0); + + const int token_idx = static_cast(blockIdx.y); + const int hv = static_cast(blockIdx.x); + const int tid = static_cast(threadIdx.x); + + if (token_idx >= token_count || hv >= kLocalVHeads) { + return; + } + + const int mapped_h = hv / kRepeatFactor; + + auto qk_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto out_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + + const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kLocalMixedQKVDim; + const scalar_t* q_src_ptr = token_ptr; + const scalar_t* k_src_ptr = token_ptr + kLocalQDim; + const scalar_t* v_src_ptr = token_ptr + kLocalQDim + kLocalKDim; + + scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kLocalVHeads * kHeadDimV; + + // Current version uses a direct GMEM->GMEM vector copy path. This keeps the + // kernel simple while already removing the scalar-copy bottleneck from the + // first draft. More aggressive staging/copy strategies should be driven by + // profiling rather than added pre-emptively. + for (int vec_idx = tid; vec_idx < kHeadDimV / kVec; vec_idx += blockDim.x) { + const int d = vec_idx * kVec; + const int q_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int k_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int v_src_idx = crd2idx(make_coord(hv, d), v_src_layout); + const int dst_idx = crd2idx(make_coord(hv, d), out_layout); + + copy_vec_contiguous(q_dst_ptr + dst_idx, q_src_ptr + q_src_idx); + copy_vec_contiguous(k_dst_ptr + dst_idx, k_src_ptr + k_src_idx); + copy_vec_contiguous(v_dst_ptr + dst_idx, v_src_ptr + v_src_idx); + } + + if (tid == 0) { + // TODO(qwen35-layout-opt): If a/b copy becomes measurable, fuse a wider + // per-head copy path here instead of scalar head writes. + const int head_idx = crd2idx(make_coord(hv), head_layout); + const int64_t token_head_offset = static_cast(token_idx) * kLocalVHeads + head_idx; + a_kernel[token_head_offset] = a[token_head_offset]; + b_kernel[token_head_offset] = b[token_head_offset]; + } +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu new file mode 100644 index 00000000..ab822752 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu @@ -0,0 +1,293 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#include "qwen35_decode_common.cuh" +#include "qwen35_scalar_kda_kernel.hpp" + +#include +#include +#include +#include + +namespace cula::qwen35::decode { + +namespace { + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +template +void dispatch_scalar_decode_for_heads( + cudaStream_t stream, + const scalar_t* q_rep, + const scalar_t* k_rep, + const scalar_t* v, + const scalar_t* a_kernel, + const scalar_t* b_kernel, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + using Shape = Qwen35DecodeLocalShape; + kernel::launch_qwen35_scalar_kda_decode_kernel( + stream, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + +template +void dispatch_layout_scalar_decode_for_heads( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + using Shape = Qwen35DecodeLocalShape; + kernel::launch_qwen35_layout_scalar_kda_decode_kernel( + stream, + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + +} // namespace + +void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params) { + const at::Tensor& q_rep = params.q_rep; + const at::Tensor& k_rep = params.k_rep; + const at::Tensor& v = params.v; + const at::Tensor& a_kernel = params.a_kernel; + const at::Tensor& b_kernel = params.b_kernel; + const at::Tensor& A_log = params.A_log; + const at::Tensor& dt_bias = params.dt_bias; + const at::Tensor& recurrent_state = params.recurrent_state; + const at::Tensor& pool_idx = params.pool_idx; + const at::Tensor& out = params.out; + + TORCH_CHECK(q_rep.is_cuda(), "q_rep must be a CUDA tensor."); + const at::Device device = q_rep.device(); + + check_tensor_device(k_rep, "k_rep", device); + check_tensor_device(v, "v", device); + check_tensor_device(a_kernel, "a_kernel", device); + check_tensor_device(b_kernel, "b_kernel", device); + check_tensor_device(A_log, "A_log", device); + check_tensor_device(dt_bias, "dt_bias", device); + check_tensor_device(recurrent_state, "recurrent_state", device); + check_tensor_device(pool_idx, "pool_idx", device); + check_tensor_device(out, "out", device); + + TORCH_CHECK(q_rep.is_contiguous(), "q_rep must be contiguous."); + TORCH_CHECK(k_rep.is_contiguous(), "k_rep must be contiguous."); + TORCH_CHECK(v.is_contiguous(), "v must be contiguous."); + TORCH_CHECK(a_kernel.is_contiguous(), "a_kernel must be contiguous."); + TORCH_CHECK(b_kernel.is_contiguous(), "b_kernel must be contiguous."); + TORCH_CHECK(A_log.is_contiguous(), "A_log must be contiguous."); + TORCH_CHECK(dt_bias.is_contiguous(), "dt_bias must be contiguous."); + TORCH_CHECK(recurrent_state.is_contiguous(), "recurrent_state must be contiguous."); + TORCH_CHECK(pool_idx.is_contiguous(), "pool_idx must be contiguous."); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous."); + + TORCH_CHECK( + q_rep.scalar_type() == k_rep.scalar_type() && q_rep.scalar_type() == v.scalar_type() && + q_rep.scalar_type() == a_kernel.scalar_type() && q_rep.scalar_type() == b_kernel.scalar_type() && + q_rep.scalar_type() == out.scalar_type(), + "q_rep/k_rep/v/a_kernel/b_kernel/out must share the same dtype."); + TORCH_CHECK(A_log.scalar_type() == at::kFloat, "A_log must be float32."); + TORCH_CHECK(dt_bias.scalar_type() == at::kFloat, "dt_bias must be float32."); + TORCH_CHECK(recurrent_state.scalar_type() == at::kFloat, "recurrent_state must be float32."); + TORCH_CHECK(pool_idx.scalar_type() == at::kInt, "pool_idx must be int32."); + + TORCH_CHECK(q_rep.dim() == 3, "q_rep must have shape [N, local_v_heads, 128]."); + const int64_t token_count = q_rep.size(0); + const int64_t local_v_heads = q_rep.size(1); + TORCH_CHECK(is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + TORCH_CHECK( + q_rep.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimQK}), + "q_rep must have shape [N, local_v_heads, 128]."); + TORCH_CHECK( + k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimQK}), + "k_rep must have shape [N, local_v_heads, 128]."); + TORCH_CHECK( + v.dim() == 3 && v.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimV}), + "v must have shape [N, local_v_heads, 128]."); + TORCH_CHECK( + a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "a_kernel must have shape [N, local_v_heads]."); + TORCH_CHECK( + b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "b_kernel must have shape [N, local_v_heads]."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == local_v_heads, "A_log must have shape [local_v_heads]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == local_v_heads, "dt_bias must have shape [local_v_heads]."); + TORCH_CHECK( + recurrent_state.dim() == 4 && + recurrent_state.size(1) == local_v_heads && + recurrent_state.size(2) == kHeadDimQK && + recurrent_state.size(3) == kHeadDimV, + "recurrent_state must have shape [pool, local_v_heads, 128, 128]."); + TORCH_CHECK(pool_idx.dim() == 1 && pool_idx.size(0) == token_count, "pool_idx must have shape [N]."); + TORCH_CHECK( + out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimV}), + "out must have shape [N, local_v_heads, 128]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + q_rep.scalar_type(), + "launch_qwen35_scalar_kda_decode_kernel", + [&] { + switch (local_v_heads) { + case 48: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 24: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 12: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 6: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void run_qwen35_layout_scalar_kda_decode(LayoutScalarKdaDecodeParams& params) { + const at::Tensor& mixed_qkv_conv = params.mixed_qkv_conv; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& A_log = params.A_log; + const at::Tensor& dt_bias = params.dt_bias; + const at::Tensor& recurrent_state = params.recurrent_state; + const at::Tensor& pool_idx = params.pool_idx; + const at::Tensor& out = params.out; + + TORCH_CHECK(mixed_qkv_conv.is_cuda(), "mixed_qkv_conv must be a CUDA tensor."); + const at::Device device = mixed_qkv_conv.device(); + + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(A_log, "A_log", device); + check_tensor_device(dt_bias, "dt_bias", device); + check_tensor_device(recurrent_state, "recurrent_state", device); + check_tensor_device(pool_idx, "pool_idx", device); + check_tensor_device(out, "out", device); + + TORCH_CHECK(mixed_qkv_conv.is_contiguous(), "mixed_qkv_conv must be contiguous."); + TORCH_CHECK(a.is_contiguous(), "a must be contiguous."); + TORCH_CHECK(b.is_contiguous(), "b must be contiguous."); + TORCH_CHECK(A_log.is_contiguous(), "A_log must be contiguous."); + TORCH_CHECK(dt_bias.is_contiguous(), "dt_bias must be contiguous."); + TORCH_CHECK(recurrent_state.is_contiguous(), "recurrent_state must be contiguous."); + TORCH_CHECK(pool_idx.is_contiguous(), "pool_idx must be contiguous."); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous."); + + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == a.scalar_type() && + mixed_qkv_conv.scalar_type() == b.scalar_type() && + mixed_qkv_conv.scalar_type() == out.scalar_type(), + "mixed_qkv_conv/a/b/out must share the same dtype."); + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == at::kHalf || mixed_qkv_conv.scalar_type() == at::kBFloat16, + "mixed_qkv_conv must be float16 or bfloat16."); + TORCH_CHECK(A_log.scalar_type() == at::kFloat, "A_log must be float32."); + TORCH_CHECK(dt_bias.scalar_type() == at::kFloat, "dt_bias must be float32."); + TORCH_CHECK(recurrent_state.scalar_type() == at::kFloat, "recurrent_state must be float32."); + TORCH_CHECK(pool_idx.scalar_type() == at::kInt, "pool_idx must be int32."); + + TORCH_CHECK(mixed_qkv_conv.dim() == 2, "mixed_qkv_conv must have shape [N, local_conv_dim]."); + TORCH_CHECK(a.dim() == 2, "a must have shape [N, local_v_heads]."); + const int64_t token_count = mixed_qkv_conv.size(0); + const int64_t local_v_heads = a.size(1); + TORCH_CHECK(is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + const int local_qk_heads = local_qk_heads_from_v_heads(static_cast(local_v_heads)); + const int local_mixed_dim = local_mixed_qkv_dim(local_qk_heads, static_cast(local_v_heads)); + TORCH_CHECK( + mixed_qkv_conv.sizes() == at::IntArrayRef({token_count, local_mixed_dim}), + "mixed_qkv_conv must have shape [N, local_conv_dim]."); + TORCH_CHECK( + a.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "a must have shape [N, local_v_heads]."); + TORCH_CHECK( + b.dim() == 2 && b.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "b must have shape [N, local_v_heads]."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == local_v_heads, "A_log must have shape [local_v_heads]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == local_v_heads, "dt_bias must have shape [local_v_heads]."); + TORCH_CHECK( + recurrent_state.dim() == 4 && + recurrent_state.size(1) == local_v_heads && + recurrent_state.size(2) == kHeadDimQK && + recurrent_state.size(3) == kHeadDimV, + "recurrent_state must have shape [pool, local_v_heads, 128, 128]."); + TORCH_CHECK(pool_idx.dim() == 1 && pool_idx.size(0) == token_count, "pool_idx must have shape [N]."); + TORCH_CHECK( + out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimV}), + "out must have shape [N, local_v_heads, 128]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + mixed_qkv_conv.scalar_type(), + "launch_qwen35_layout_scalar_kda_decode_kernel", + [&] { + switch (local_v_heads) { + case 48: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 24: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 12: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 6: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp new file mode 100644 index 00000000..25e0aa7a --- /dev/null +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -0,0 +1,1079 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include "qwen35_decode_common.cuh" +#include "qwen35_scalar_kda_mainloop.hpp" + +#include +#include + +namespace cula::qwen35::decode::kernel { + +using namespace cute; + +template +CUTE_DEVICE void cp_async_ca_shared_global(void* smem_ptr, const void* gmem_ptr) { + static_assert(kBytes == 16, "Only 16-byte cp.async copies are supported here."); + const unsigned smem_addr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" ::"r"(smem_addr), "l"(gmem_ptr)); +} + +CUTE_DEVICE void cp_async_bulk_shared_global( + void* smem_ptr, + const void* gmem_ptr, + uint32_t bytes, + cutlass::arch::ClusterTransactionBarrier::ValueType* barrier) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + const uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + const uint32_t barrier_addr = cute::cast_smem_ptr_to_uint(barrier); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " + "[%0], [%1], %2, [%3];\n" + : + : "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(barrier_addr) + : "memory"); +#else + (void)smem_ptr; + (void)gmem_ptr; + (void)bytes; + (void)barrier; +#endif +} + +CUTE_DEVICE void cp_async_commit_group() { + asm volatile("cp.async.commit_group;\n" ::); +} + +CUTE_DEVICE void cp_async_wait_all() { + asm volatile("cp.async.wait_group 0;\n" ::); +} + +CUTE_DEVICE void cp_async_wait_group_1() { + asm volatile("cp.async.wait_group 1;\n" ::); +} + +template +struct Qwen35ScalarKdaDecodeKernel { + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); + // Decode-first design: + // - 1 CTA owns 1 (token_idx, hv) + // - 1 warpgroup (128 threads) per CTA + // - recurrent state stays fp32 and is traversed as 16x16 tiles over the + // internal [V, K] view + // - the intended optimized path is fp32 FFMA on CUDA cores, not a forced + // Tensor Core lowering + static constexpr int kThreads = Shape::kKdaThreads; + static constexpr int kWarpGroupThreads = Shape::kKdaThreads; + static constexpr int kTileV = Shape::kKdaTileV; + static constexpr int kTileK = Shape::kKdaTileK; + static constexpr int kTilesPerV = kHeadDimV / kTileV; + static constexpr int kTilesPerK = kHeadDimQK / kTileK; + + static_assert(kLocalQKHeads < kLocalVHeads); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kThreads == kWarpGroupThreads); + static_assert(kHeadDimV % kTileV == 0); + static_assert(kHeadDimQK % kTileK == 0); + + struct SharedStorage { + // Shared staging plan for the fp32 decode path: + // - q/k/v are staged once per CTA + // - proj/out intermediates remain in fp32 + // - recurrent state itself remains in fp32 global storage + alignas(16) float q_smem[kHeadDimQK]; + alignas(16) float k_smem[kHeadDimQK]; + alignas(16) scalar_t v_smem[kHeadDimV]; + alignas(16) float norm_smem[2]; + alignas(16) float proj_smem[kHeadDimV]; + alignas(16) float out_smem[kHeadDimV]; + }; + + static dim3 block_shape() { + return dim3(kThreads, 1, 1); + } + + static dim3 grid_shape(int token_count) { + // One block owns one (token_idx, hv) pair in the first implementation. + return dim3(static_cast(Shape::kLocalVHeads), static_cast(token_count), 1); + } + + template + CUTE_DEVICE static void run_device( + const scalar_t* __restrict__ q_rep, + const scalar_t* __restrict__ k_rep, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a_kernel, + const scalar_t* __restrict__ b_kernel, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count, + SharedStorage& storage) { + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + if (token_idx >= token_count || hv >= kLocalVHeads) { + return; + } + + // Internal tensor-view contract fixed for the first implementation pass: + // + // 1. q_rep / k_rep / v / out stay in their external contiguous layouts: + // - q_rep : [N, HV, K] with stride (HV*K, K, 1) + // - k_rep : [N, HV, K] with stride (HV*K, K, 1) + // - v : [N, HV, V] with stride (HV*V, V, 1) + // - out : [N, HV, V] with stride (HV*V, V, 1) + // + // 2. a_kernel / b_kernel are treated as: + // - [N, HV] with stride (HV, 1) + // + // 3. A_log / dt_bias are treated as: + // - [HV] with stride (1) + // + // 4. recurrent_state keeps the external physical storage contract: + // - [pool, HV, K, V] + // but the kernel's main computation will use an internal VK view: + // - [pool, HV, V, K] + // + // This lets the recurrent update consume one V-row of state against q/k + // more naturally in the first mainloop design, while preserving the + // existing external state ABI. + // + // The current block owns exactly one (token_idx, hv) pair. That means one + // warpgroup-sized CTA updates one 128x128 recurrent-state tile for one + // v-head. + // + // TODO(qwen35-scalar-kda-opt): + // - Likely next optimization path: keep one CTA per (token_idx, hv), but + // tile the 128x128 state more aggressively inside the block (for example + // along V tiles or KxV subtiles assigned per warp). + // - More complex alternative: split one (token_idx, hv) tile across + // multiple CTAs and coordinate updates. Not a first-pass target. + // - After the fp32 decode path is stable, evaluate warp specialization: + // dedicated producer/load warp(s) vs consumer/compute warp(s), instead + // of introducing that complexity before the math path itself is stable. + + auto q_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimQK, kHeadDimQK, Int<1>{})); + auto v_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_kv = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + auto gQ = make_tensor(make_gmem_ptr(q_rep), q_layout); + auto gK = make_tensor(make_gmem_ptr(k_rep), q_layout); + auto gV = make_tensor(make_gmem_ptr(v), v_layout); + auto gO = make_tensor(make_gmem_ptr(out), v_layout); + auto gA = make_tensor(make_gmem_ptr(a_kernel), head_layout); + auto gB = make_tensor(make_gmem_ptr(b_kernel), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_kv = make_tensor(make_gmem_ptr(recurrent_state), state_layout_kv); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + (void)gH_kv; // Keep the physical KV view documented and available. + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, hv, _); + auto k_vec = gK(token_idx, hv, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto a_scalar = gA(token_idx, hv); + auto b_scalar = gB(token_idx, hv); + auto A_log_scalar = gAlog(hv); + auto dt_bias_scalar = gDt(hv); + auto state_vk = gH_vk(state_row, hv, _, _); + + Mainloop::run( + q_vec, + k_vec, + v_vec, + a_scalar, + b_scalar, + A_log_scalar, + dt_bias_scalar, + state_vk, + out_vec, + storage, + tid, + kThreads); + } + + template + CUTE_DEVICE static void run_layout_device( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count, + SharedStorage& storage) { + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; + + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + if (token_idx >= token_count || hv >= kLocalVHeads) { + return; + } + + const int mapped_h = hv / kRepeatFactor; + + auto qk_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimQK, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimV, Int<1>{})); + auto out_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_kv = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + const scalar_t* q_src = mixed_qkv_conv; + const scalar_t* k_src = mixed_qkv_conv + kLocalQDim; + const scalar_t* v_src = mixed_qkv_conv + kLocalQDim + kLocalKDim; + + auto gQ = make_tensor(make_gmem_ptr(q_src), qk_src_layout); + auto gK = make_tensor(make_gmem_ptr(k_src), qk_src_layout); + auto gV = make_tensor(make_gmem_ptr(v_src), v_src_layout); + auto gO = make_tensor(make_gmem_ptr(out), out_layout); + auto gA = make_tensor(make_gmem_ptr(a), head_layout); + auto gB = make_tensor(make_gmem_ptr(b), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_kv = make_tensor(make_gmem_ptr(recurrent_state), state_layout_kv); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + (void)gH_kv; + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, mapped_h, _); + auto k_vec = gK(token_idx, mapped_h, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto a_scalar = gA(token_idx, hv); + auto b_scalar = gB(token_idx, hv); + auto A_log_scalar = gAlog(hv); + auto dt_bias_scalar = gDt(hv); + auto state_vk = gH_vk(state_row, hv, _, _); + + Mainloop::run( + q_vec, + k_vec, + v_vec, + a_scalar, + b_scalar, + A_log_scalar, + dt_bias_scalar, + state_vk, + out_vec, + storage, + tid, + kThreads); + } +}; + +template > +__global__ void qwen35_scalar_kda_decode_kernel( + const scalar_t* __restrict__ q_rep, + const scalar_t* __restrict__ k_rep, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a_kernel, + const scalar_t* __restrict__ b_kernel, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; + Qwen35ScalarKdaDecodeKernel::template run_device( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count, + storage); +} + +template > +void launch_qwen35_scalar_kda_decode_kernel( + cudaStream_t stream, + const scalar_t* q_rep, + const scalar_t* k_rep, + const scalar_t* v, + const scalar_t* a_kernel, + const scalar_t* b_kernel, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + if (token_count >= 64) { + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_scalar_kda_decode_kernel< + scalar_t, + kLocalQKHeads, + kLocalVHeads, + Qwen35ScalarKdaDecodeLongMainloop><<>>( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); + return; + } + + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_scalar_kda_decode_kernel<<>>( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + +template > +__global__ void qwen35_layout_scalar_kda_decode_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; + Qwen35ScalarKdaDecodeKernel::template run_layout_device( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count, + storage); +} + +template +__global__ void qwen35_layout_scalar_kda_decode_long_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; + constexpr int kWarpTileV = 32; + constexpr int kWarpSize = 32; + constexpr int kThreads = 128; + constexpr int kWarps = kThreads / kWarpSize; + constexpr int kPipeTileK = kPipeTileK_; + constexpr int kPipeStages = 2; + constexpr int kVecFloats = 4; + constexpr int kStatePipeStrideV = kHeadDimV + 4; + + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kHeadDimV % kWarpTileV == 0); + static_assert(kWarps * kWarpTileV == kHeadDimV); + static_assert(kHeadDimQK % kPipeTileK == 0); + static_assert(kPipeTileK == 16 || kPipeTileK == 32); + + __shared__ float q_smem[kHeadDimQK]; + __shared__ float k_smem[kHeadDimQK]; + __shared__ float norm_smem[2 * kWarps]; + __shared__ float state_pipe[kPipeStages][kPipeTileK][kStatePipeStrideV]; + + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + const int lane = tid & (kWarpSize - 1); + const int warp_id = tid / kWarpSize; + const int mapped_h = hv / kRepeatFactor; + const int v_row = warp_id * kWarpTileV + lane; + + auto qk_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimQK, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimV, Int<1>{})); + auto out_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + const scalar_t* q_src = mixed_qkv_conv; + const scalar_t* k_src = mixed_qkv_conv + kLocalQDim; + const scalar_t* v_src = mixed_qkv_conv + kLocalQDim + kLocalKDim; + + auto gQ = make_tensor(make_gmem_ptr(q_src), qk_src_layout); + auto gK = make_tensor(make_gmem_ptr(k_src), qk_src_layout); + auto gV = make_tensor(make_gmem_ptr(v_src), v_src_layout); + auto gO = make_tensor(make_gmem_ptr(out), out_layout); + auto gA = make_tensor(make_gmem_ptr(a), head_layout); + auto gB = make_tensor(make_gmem_ptr(b), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, mapped_h, _); + auto k_vec = gK(token_idx, mapped_h, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto state_vk = gH_vk(state_row, hv, _, _); + + const float a_val = static_cast(gA(token_idx, hv)); + const float b_val = static_cast(gB(token_idx, hv)); + const float g = -expf(static_cast(gAlog(hv))) * + Qwen35ScalarKdaDecodeMainloop::softplusf_approx(a_val + static_cast(gDt(hv))); + const float decay = expf(g); + const float beta = 1.f / (1.f + expf(-b_val)); + + const float q_raw = static_cast(q_vec(tid)); + const float k_raw = static_cast(k_vec(tid)); + q_smem[tid] = q_raw; + k_smem[tid] = k_raw; + + float q_norm_sq = q_raw * q_raw; + float k_norm_sq = k_raw * k_raw; + q_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_norm_sq); + k_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_norm_sq); + if (lane == 0) { + norm_smem[warp_id] = q_norm_sq; + norm_smem[kWarps + warp_id] = k_norm_sq; + } + __syncthreads(); + + float q_block_sum = lane < kWarps ? norm_smem[lane] : 0.f; + float k_block_sum = lane < kWarps ? norm_smem[kWarps + lane] : 0.f; + if (warp_id == 0) { + q_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_block_sum); + k_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_block_sum); + if (lane == 0) { + norm_smem[0] = rsqrtf(q_block_sum + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); + norm_smem[1] = rsqrtf(k_block_sum + 1e-6f); + } + } + __syncthreads(); + + const float q_normed = q_raw * norm_smem[0]; + const float k_normed = k_raw * norm_smem[1]; + q_smem[tid] = q_normed; + k_smem[tid] = k_normed; + + float qk_dot = q_normed * k_normed; + qk_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_dot); + if (lane == 0) { + norm_smem[warp_id] = qk_dot; + } + __syncthreads(); + float qk_block_sum = lane < kWarps ? norm_smem[lane] : 0.f; + if (warp_id == 0) { + qk_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_block_sum); + if (lane == 0) { + norm_smem[2] = qk_block_sum; + } + } + __syncthreads(); + + auto load_state_pipe_tile = [&](int stage, int k_base) { +#pragma unroll 1 + for (int elem = tid * kVecFloats; elem < kPipeTileK * kHeadDimV; elem += kThreads * kVecFloats) { + const int k_local = elem / kHeadDimV; + const int v_base = elem - k_local * kHeadDimV; + cp_async_ca_shared_global<16>( + &state_pipe[stage][k_local][v_base], + &state_vk(v_base, k_base + k_local)); + } + }; + + float proj_acc0 = 0.f; + float proj_acc1 = 0.f; + float proj_acc2 = 0.f; + float proj_acc3 = 0.f; + float out_acc0 = 0.f; + float out_acc1 = 0.f; + float out_acc2 = 0.f; + float out_acc3 = 0.f; + int pipe_stage = 0; + load_state_pipe_tile(0, 0); + cp_async_commit_group(); + if (kPipeTileK < kHeadDimQK) { + load_state_pipe_tile(1, kPipeTileK); + cp_async_commit_group(); + } + +#pragma unroll 1 + for (int k_base = 0; k_base < kHeadDimQK; k_base += kPipeTileK) { + const int next_k_base = k_base + kPipeTileK; + const int next_stage = pipe_stage ^ 1; + const int prefetch_k_base = k_base + 2 * kPipeTileK; + if (next_k_base < kHeadDimQK) { + cp_async_wait_group_1(); + } else { + cp_async_wait_all(); + } + __syncthreads(); + + float q_regs[kPipeTileK]; + float k_regs[kPipeTileK]; +#pragma unroll + for (int kk = 0; kk < kPipeTileK; ++kk) { + q_regs[kk] = q_smem[k_base + kk]; + k_regs[kk] = k_smem[k_base + kk]; + } + +#pragma unroll + for (int k_local = 0; k_local < kPipeTileK; k_local += 4) { + const float state0 = state_pipe[pipe_stage][k_local + 0][v_row]; + const float state1 = state_pipe[pipe_stage][k_local + 1][v_row]; + const float state2 = state_pipe[pipe_stage][k_local + 2][v_row]; + const float state3 = state_pipe[pipe_stage][k_local + 3][v_row]; + proj_acc0 += state0 * k_regs[k_local + 0]; + proj_acc1 += state1 * k_regs[k_local + 1]; + proj_acc2 += state2 * k_regs[k_local + 2]; + proj_acc3 += state3 * k_regs[k_local + 3]; + out_acc0 += state0 * q_regs[k_local + 0]; + out_acc1 += state1 * q_regs[k_local + 1]; + out_acc2 += state2 * q_regs[k_local + 2]; + out_acc3 += state3 * q_regs[k_local + 3]; + } + if (prefetch_k_base < kHeadDimQK) { + __syncthreads(); + load_state_pipe_tile(pipe_stage, prefetch_k_base); + cp_async_commit_group(); + } + pipe_stage = next_stage; + } + + const float proj_row = (proj_acc0 + proj_acc1) + (proj_acc2 + proj_acc3); + const float out_old_row = (out_acc0 + out_acc1) + (out_acc2 + out_acc3); + const float v_val = static_cast(v_vec(v_row)); + const float v_new_row = beta * (v_val - decay * proj_row); + out_vec(v_row) = static_cast(decay * out_old_row + v_new_row * norm_smem[2]); + + pipe_stage = 0; + load_state_pipe_tile(0, 0); + cp_async_commit_group(); + if (kPipeTileK < kHeadDimQK) { + load_state_pipe_tile(1, kPipeTileK); + cp_async_commit_group(); + } + +#pragma unroll 1 + for (int k_base = 0; k_base < kHeadDimQK; k_base += kPipeTileK) { + const int next_k_base = k_base + kPipeTileK; + const int next_stage = pipe_stage ^ 1; + const int prefetch_k_base = k_base + 2 * kPipeTileK; + if (next_k_base < kHeadDimQK) { + cp_async_wait_group_1(); + } else { + cp_async_wait_all(); + } + __syncthreads(); + + float k_regs[kPipeTileK]; +#pragma unroll + for (int kk = 0; kk < kPipeTileK; ++kk) { + k_regs[kk] = k_smem[k_base + kk]; + } + +#pragma unroll + for (int k_local = 0; k_local < kPipeTileK; k_local += 4) { + const float state_new0 = decay * state_pipe[pipe_stage][k_local + 0][v_row] + v_new_row * k_regs[k_local + 0]; + const float state_new1 = decay * state_pipe[pipe_stage][k_local + 1][v_row] + v_new_row * k_regs[k_local + 1]; + const float state_new2 = decay * state_pipe[pipe_stage][k_local + 2][v_row] + v_new_row * k_regs[k_local + 2]; + const float state_new3 = decay * state_pipe[pipe_stage][k_local + 3][v_row] + v_new_row * k_regs[k_local + 3]; + state_vk(v_row, k_base + k_local + 0) = state_new0; + state_vk(v_row, k_base + k_local + 1) = state_new1; + state_vk(v_row, k_base + k_local + 2) = state_new2; + state_vk(v_row, k_base + k_local + 3) = state_new3; + } + if (prefetch_k_base < kHeadDimQK) { + __syncthreads(); + load_state_pipe_tile(pipe_stage, prefetch_k_base); + cp_async_commit_group(); + } + pipe_stage = next_stage; + } + +} + + +template +__global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; + constexpr int kWarpSize = 32; + constexpr int kThreads = kTileV; + constexpr int kWarps = kThreads / kWarpSize; + constexpr int kVTiles = kHeadDimV / kTileV; + constexpr int kKPerThread = kHeadDimQK / kThreads; + + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kTileV == 32 || kTileV == 64 || kTileV == 128); + static_assert(kHeadDimV % kTileV == 0); + static_assert(kHeadDimQK % kThreads == 0); + + __shared__ float q_smem[kHeadDimQK]; + __shared__ float k_smem[kHeadDimQK]; + __shared__ float state_smem[kHeadDimQK][kTileV]; + __shared__ float norm_smem[3]; + __shared__ float warp_reduce_smem[3 * kWarps]; + __shared__ cutlass::arch::ClusterTransactionBarrier::ValueType state_barrier; + + const int hv_tile = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int hv = hv_tile / kVTiles; + const int v_tile = hv_tile - hv * kVTiles; + const int tid = static_cast(threadIdx.x); + const int lane = tid & (kWarpSize - 1); + const int warp_id = tid / kWarpSize; + const int mapped_h = hv / kRepeatFactor; + const int v_row = v_tile * kTileV + tid; + + auto qk_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimQK, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimV, Int<1>{})); + auto out_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + const scalar_t* q_src = mixed_qkv_conv; + const scalar_t* k_src = mixed_qkv_conv + kLocalQDim; + const scalar_t* v_src = mixed_qkv_conv + kLocalQDim + kLocalKDim; + + auto gQ = make_tensor(make_gmem_ptr(q_src), qk_src_layout); + auto gK = make_tensor(make_gmem_ptr(k_src), qk_src_layout); + auto gV = make_tensor(make_gmem_ptr(v_src), v_src_layout); + auto gO = make_tensor(make_gmem_ptr(out), out_layout); + auto gA = make_tensor(make_gmem_ptr(a), head_layout); + auto gB = make_tensor(make_gmem_ptr(b), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, mapped_h, _); + auto k_vec = gK(token_idx, mapped_h, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto state_vk = gH_vk(state_row, hv, _, _); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Start the full-state transfer before the Q/K normalization and gate + // arithmetic. Those independent instructions hide a portion of the HBM + // latency for the long-token path. + if (tid == 0) { + cutlass::arch::ClusterTransactionBarrier::init(&state_barrier, 1); + cutlass::arch::ClusterTransactionBarrier::arrive_and_expect_tx( + &state_barrier, kHeadDimQK * kTileV * sizeof(float)); + } + __syncthreads(); + if constexpr (kTileV == kHeadDimV) { + constexpr int kStateBytes = kHeadDimQK * kHeadDimV * sizeof(float); + constexpr int kBulkChunkBytes = 32 * 1024; + constexpr int kBulkChunkFloats = kBulkChunkBytes / sizeof(float); + constexpr int kBulkChunks = kStateBytes / kBulkChunkBytes; + if (tid == 0) { +#pragma unroll + for (int chunk = 0; chunk < kBulkChunks; ++chunk) { + cp_async_bulk_shared_global( + &state_smem[0][0] + chunk * kBulkChunkFloats, + &state_vk(0, 0) + chunk * kBulkChunkFloats, + kBulkChunkBytes, + &state_barrier); + } + } + } else { +#pragma unroll 1 + for (int k_idx = tid; k_idx < kHeadDimQK; k_idx += kThreads) { + cp_async_bulk_shared_global( + &state_smem[k_idx][0], + &state_vk(v_tile * kTileV, k_idx), + kTileV * sizeof(float), + &state_barrier); + } + } +#endif + + float q_norm_sq = 0.f; + float k_norm_sq = 0.f; + float qk_raw_dot = 0.f; +#pragma unroll + for (int i = 0; i < kKPerThread; ++i) { + const int k_idx = i * kThreads + tid; + const float q_raw = static_cast(q_vec(k_idx)); + const float k_raw = static_cast(k_vec(k_idx)); + q_smem[k_idx] = q_raw; + k_smem[k_idx] = k_raw; + q_norm_sq += q_raw * q_raw; + k_norm_sq += k_raw * k_raw; + qk_raw_dot += q_raw * k_raw; + } + q_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_norm_sq); + k_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_norm_sq); + qk_raw_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_raw_dot); + if (lane == 0) { + warp_reduce_smem[warp_id] = q_norm_sq; + warp_reduce_smem[kWarps + warp_id] = k_norm_sq; + warp_reduce_smem[2 * kWarps + warp_id] = qk_raw_dot; + } + __syncthreads(); + if (warp_id == 0) { + float q_block_sum = lane < kWarps ? warp_reduce_smem[lane] : 0.f; + float k_block_sum = lane < kWarps ? warp_reduce_smem[kWarps + lane] : 0.f; + float qk_block_sum = lane < kWarps ? warp_reduce_smem[2 * kWarps + lane] : 0.f; + q_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_block_sum); + k_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_block_sum); + qk_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_block_sum); + if (lane == 0) { + norm_smem[0] = rsqrtf(q_block_sum + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); + norm_smem[1] = rsqrtf(k_block_sum + 1e-6f); + norm_smem[2] = qk_block_sum * norm_smem[0] * norm_smem[1]; + } + } + __syncthreads(); + +#pragma unroll + for (int i = 0; i < kKPerThread; ++i) { + const int k_idx = i * kThreads + tid; + const float q_normed = q_smem[k_idx] * norm_smem[0]; + const float k_normed = k_smem[k_idx] * norm_smem[1]; + q_smem[k_idx] = q_normed; + k_smem[k_idx] = k_normed; + } + + const float a_val = static_cast(gA(token_idx, hv)); + const float b_val = static_cast(gB(token_idx, hv)); + const float g = -expf(static_cast(gAlog(hv))) * + Qwen35ScalarKdaDecodeMainloop::softplusf_approx(a_val + static_cast(gDt(hv))); + const float decay = expf(g); + const float beta = 1.f / (1.f + expf(-b_val)); + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cutlass::arch::ClusterTransactionBarrier::wait(&state_barrier, 0); + __syncthreads(); +#else +#pragma unroll 1 + for (int elem = tid * 4; elem < kHeadDimQK * kTileV; elem += kThreads * 4) { + const int k_idx = elem / kTileV; + const int v_base = elem - k_idx * kTileV; + cp_async_ca_shared_global<16>( + &state_smem[k_idx][v_base], + &state_vk(v_tile * kTileV + v_base, k_idx)); + } + cp_async_commit_group(); + cp_async_wait_all(); + __syncthreads(); +#endif + + float proj_acc0 = 0.f; + float proj_acc1 = 0.f; + float proj_acc2 = 0.f; + float proj_acc3 = 0.f; + float proj_acc4 = 0.f; + float proj_acc5 = 0.f; + float proj_acc6 = 0.f; + float proj_acc7 = 0.f; + float out_acc0 = 0.f; + float out_acc1 = 0.f; + float out_acc2 = 0.f; + float out_acc3 = 0.f; + float out_acc4 = 0.f; + float out_acc5 = 0.f; + float out_acc6 = 0.f; + float out_acc7 = 0.f; +#pragma unroll 1 + for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 8) { + const float state0 = state_smem[k_idx + 0][tid]; + const float state1 = state_smem[k_idx + 1][tid]; + const float state2 = state_smem[k_idx + 2][tid]; + const float state3 = state_smem[k_idx + 3][tid]; + const float state4 = state_smem[k_idx + 4][tid]; + const float state5 = state_smem[k_idx + 5][tid]; + const float state6 = state_smem[k_idx + 6][tid]; + const float state7 = state_smem[k_idx + 7][tid]; + const float k0 = k_smem[k_idx + 0]; + const float k1 = k_smem[k_idx + 1]; + const float k2 = k_smem[k_idx + 2]; + const float k3 = k_smem[k_idx + 3]; + const float k4 = k_smem[k_idx + 4]; + const float k5 = k_smem[k_idx + 5]; + const float k6 = k_smem[k_idx + 6]; + const float k7 = k_smem[k_idx + 7]; + const float q0 = q_smem[k_idx + 0]; + const float q1 = q_smem[k_idx + 1]; + const float q2 = q_smem[k_idx + 2]; + const float q3 = q_smem[k_idx + 3]; + const float q4 = q_smem[k_idx + 4]; + const float q5 = q_smem[k_idx + 5]; + const float q6 = q_smem[k_idx + 6]; + const float q7 = q_smem[k_idx + 7]; + proj_acc0 += state0 * k0; + proj_acc1 += state1 * k1; + proj_acc2 += state2 * k2; + proj_acc3 += state3 * k3; + proj_acc4 += state4 * k4; + proj_acc5 += state5 * k5; + proj_acc6 += state6 * k6; + proj_acc7 += state7 * k7; + out_acc0 += state0 * q0; + out_acc1 += state1 * q1; + out_acc2 += state2 * q2; + out_acc3 += state3 * q3; + out_acc4 += state4 * q4; + out_acc5 += state5 * q5; + out_acc6 += state6 * q6; + out_acc7 += state7 * q7; + } + + const float proj_row = + ((proj_acc0 + proj_acc1) + (proj_acc2 + proj_acc3)) + + ((proj_acc4 + proj_acc5) + (proj_acc6 + proj_acc7)); + const float out_old_row = + ((out_acc0 + out_acc1) + (out_acc2 + out_acc3)) + + ((out_acc4 + out_acc5) + (out_acc6 + out_acc7)); + const float v_val = static_cast(v_vec(v_row)); + const float v_new_row = beta * (v_val - decay * proj_row); + out_vec(v_row) = static_cast(decay * out_old_row + v_new_row * norm_smem[2]); + +#pragma unroll 1 + for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 8) { + const float state0 = state_smem[k_idx + 0][tid]; + const float state1 = state_smem[k_idx + 1][tid]; + const float state2 = state_smem[k_idx + 2][tid]; + const float state3 = state_smem[k_idx + 3][tid]; + const float state4 = state_smem[k_idx + 4][tid]; + const float state5 = state_smem[k_idx + 5][tid]; + const float state6 = state_smem[k_idx + 6][tid]; + const float state7 = state_smem[k_idx + 7][tid]; + const float state_new0 = decay * state0 + v_new_row * k_smem[k_idx + 0]; + const float state_new1 = decay * state1 + v_new_row * k_smem[k_idx + 1]; + const float state_new2 = decay * state2 + v_new_row * k_smem[k_idx + 2]; + const float state_new3 = decay * state3 + v_new_row * k_smem[k_idx + 3]; + const float state_new4 = decay * state4 + v_new_row * k_smem[k_idx + 4]; + const float state_new5 = decay * state5 + v_new_row * k_smem[k_idx + 5]; + const float state_new6 = decay * state6 + v_new_row * k_smem[k_idx + 6]; + const float state_new7 = decay * state7 + v_new_row * k_smem[k_idx + 7]; + state_vk(v_row, k_idx + 0) = state_new0; + state_vk(v_row, k_idx + 1) = state_new1; + state_vk(v_row, k_idx + 2) = state_new2; + state_vk(v_row, k_idx + 3) = state_new3; + state_vk(v_row, k_idx + 4) = state_new4; + state_vk(v_row, k_idx + 5) = state_new5; + state_vk(v_row, k_idx + 6) = state_new6; + state_vk(v_row, k_idx + 7) = state_new7; + } +} + +template +void launch_qwen35_layout_scalar_kda_decode_long_kernel( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + constexpr int kWarpTileV = 32; + (void)kWarpTileV; + if (token_count == 64 || token_count == 128) { + constexpr int kLongTileV = 128; + dim3 grid(kLocalVHeads * (kHeadDimV / kLongTileV), token_count, 1); + dim3 block(kLongTileV, 1, 1); + qwen35_layout_scalar_kda_decode_long_vtile_kernel + <<>>( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); + return; + } + + dim3 grid(kLocalVHeads, token_count, 1); + dim3 block(128, 1, 1); + qwen35_layout_scalar_kda_decode_long_kernel<<>>( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + +template > +void launch_qwen35_layout_scalar_kda_decode_kernel( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + if (token_count >= 64) { + launch_qwen35_layout_scalar_kda_decode_long_kernel( + stream, + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); + return; + } + + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_layout_scalar_kda_decode_kernel<<>>( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + +} // namespace cula::qwen35::decode::kernel diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp new file mode 100644 index 00000000..74460077 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp @@ -0,0 +1,492 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include "qwen35_decode_common.cuh" + +#include +#include +#include +#include + +namespace cula::qwen35::decode::kernel { + +using namespace cute; + +template +struct Qwen35ScalarKdaDecodeMainloop { + // Decode design decision: + // - recurrent_state remains fp32 both physically and mathematically + // - decode is treated as a register-level recurrent GEMV/rank-1-update + // problem, not as a Tensor Core GEMM problem + // - you can think of the target implementation style as a + // flash_linear_decode_kernel: pure CUDA Core math, warp-shuffle vector + // sharing, and fp32 state kept live as long as possible during one token + // update + // + // Reason: + // - state participates in a long recurrent chain; lowering master-state + // precision is risky and usually not worth it + // - decode operates on a single-token q/k vector, so the dominant kernels + // are GEMV-like: + // proj = state @ k + // out = state' @ q + // This is typically register / memory bound rather than Tensor-Core bound + // + // Practical consequence: + // - the first production-worthy decode path should be built around fp32 FFMA + // on CUDA cores + // - tile structure is still useful, but it should serve register ownership, + // reduction, and cache behavior instead of forcing a GMMA lowering + // + // The remaining work for this kernel is therefore: + // 1. tighten thread ownership of the fp32 state tile + // 2. optimize proj / update / out reductions + // 3. evaluate warp-specialized load/compute roles only after the fp32 path + // is stable and measured + static constexpr int kTileV = kTileV_; + static constexpr int kTileK = kTileK_; + static constexpr int kTilesPerV = kHeadDimV / kTileV; + static constexpr int kTilesPerK = kHeadDimQK / kTileK; + static constexpr int kWarpSize = 32; + static constexpr int kRowsPerTile = kTileV; + static constexpr int kWarpsPerCta = 4; + static constexpr int kRowsPerWarp = kWarpSize; + static constexpr int kRowsPerThread = 1; + + static_assert(kHeadDimV == 128); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV % kTileV == 0); + static_assert(kHeadDimQK % kTileK == 0); + static_assert(kWarpsPerCta * kRowsPerWarp == kHeadDimV); + + // First concrete decode threading plan: + // + // - 1 CTA = 1 (token, hv) + // - 128 threads = 4 warps + // - 1 thread owns exactly 1 V-row of the 128x128 recurrent state + // - Therefore one CTA covers all 128 V-rows exactly once + // + // For the owned row, the thread streams over K in 16-wide tiles: + // + // state_row[0:15] -> registers + // state_row[16:31] -> registers + // ... + // state_row[112:127]-> registers + // + // This means the first concrete fp32 path does NOT attempt to keep the + // whole 128-float row resident in registers at once. Instead it keeps the + // current K tile resident: + // + // - state_regs[16] : current fp32 state tile + // - k_regs[16] : current key tile + // - q_regs[16] : current query tile + // + // plus a handful of scalar accumulators: + // + // - proj_row + // - out_row + // - v_new_row + // - gate scalars + // + // This is a practical first step toward the user's desired "state stays in + // registers for the current token" behavior while keeping register pressure + // manageable. + // + // Reduction policy for this first concrete plan: + // + // - proj/out are row-local, so no warp reduction is required + // - each row is fully owned by one thread across all K tiles + // - warp shuffle is reserved for future vector-broadcast refinements if we + // decide to move q/k staging from shared memory into warp-register paths + struct ThreadRowPlan { + int warp_id; + int lane_id; + int v_row; + bool owns_row; + }; + + struct TileCoords { + int v_base; + int k_base; + }; + + CUTE_DEVICE static ThreadRowPlan make_thread_row_plan(int tid) { + const int warp_id = tid / kWarpSize; + const int lane_id = tid % kWarpSize; + const int v_row = warp_id * kRowsPerWarp + lane_id; + const bool owns_row = v_row < kHeadDimV; + return ThreadRowPlan{warp_id, lane_id, v_row, owns_row}; + } + + template + CUTE_DEVICE static void load_vec_tile_to_regs( + TensorVec const& vec, + TileCoords coords, + float (®s)[kTileK]) { +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + regs[kk] = static_cast(vec(coords.k_base + kk)); + } + } + + template + CUTE_DEVICE static void load_state_row_tile_to_regs( + TensorState const& state_vk, + int v_row, + TileCoords coords, + float (&state_regs)[kTileK]) { +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + state_regs[kk] = static_cast(state_vk(v_row, coords.k_base + kk)); + } + } + + template + CUTE_DEVICE static void store_state_row_tile_from_regs( + TensorState& state_vk, + int v_row, + TileCoords coords, + float const (&state_regs)[kTileK]) { +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + state_vk(v_row, coords.k_base + kk) = state_regs[kk]; + } + } + + struct RowTileProjPlan { + int v_base; + int k_base; + int warp_id; + int lane_id; + bool owns_row; + int row_in_tile; + int v_row; + }; + + CUTE_DEVICE static TileCoords make_tile_coords(int tile_v, int tile_k) { + return TileCoords{tile_v * kTileV, tile_k * kTileK}; + } + + CUTE_DEVICE static RowTileProjPlan make_row_tile_proj_plan( + TileCoords coords, + int warp_id, + int lane_id) { + const bool owns_row = lane_id < kTileV; + const int row_in_tile = lane_id; + const int v_row = coords.v_base + row_in_tile; + return RowTileProjPlan{ + coords.v_base, + coords.k_base, + warp_id, + lane_id, + owns_row, + row_in_tile, + v_row, + }; + } + + struct RowTileUpdatePlan { + TileCoords coords; + int warp_id; + int lane_id; + bool owns_row; + int row_in_tile; + int v_row; + }; + + CUTE_DEVICE static RowTileUpdatePlan make_row_tile_update_plan( + TileCoords coords, + int warp_id, + int lane_id) { + const bool owns_row = lane_id < kTileV; + const int row_in_tile = lane_id; + const int v_row = coords.v_base + row_in_tile; + return RowTileUpdatePlan{ + coords, + warp_id, + lane_id, + owns_row, + row_in_tile, + v_row, + }; + } + + template + CUTE_DEVICE static float accumulate_proj_row_tile( + TensorState const& state_vk, + TensorKTile const& k_smem, + int v_row, + TileCoords coords) { + float state_regs[kTileK]; + float k_regs[kTileK]; + load_state_row_tile_to_regs(state_vk, v_row, coords, state_regs); + load_vec_tile_to_regs(k_smem, coords, k_regs); + + float accum = 0.f; +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + accum += state_regs[kk] * k_regs[kk]; + } + return accum; + } + + template + CUTE_DEVICE static float update_state_row_tile_and_accumulate_out( + TensorState& state_vk, + TensorKTile const& k_smem, + TensorQTile const& q_smem, + int v_row, + TileCoords coords, + float decay, + float v_new) { + float state_regs[kTileK]; + float k_regs[kTileK]; + float q_regs[kTileK]; + load_state_row_tile_to_regs(state_vk, v_row, coords, state_regs); + load_vec_tile_to_regs(k_smem, coords, k_regs); + load_vec_tile_to_regs(q_smem, coords, q_regs); + + float out_acc = 0.f; +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + const float state_new = decay * state_regs[kk] + v_new * k_regs[kk]; + state_regs[kk] = state_new; + out_acc += state_new * q_regs[kk]; + } + store_state_row_tile_from_regs(state_vk, v_row, coords, state_regs); + return out_acc; + } + + template + CUTE_DEVICE static float project_row_tile( + TensorState const& state_vk, + TensorKTile const& k_smem, + int v_row, + TileCoords coords) { + // Current decode path: + // - one thread owns one full V-row + // - this helper computes the row-local proj contribution for one K tile + // - no cross-thread reduction is needed + return accumulate_proj_row_tile(state_vk, k_smem, v_row, coords); + } + + template + CUTE_DEVICE static float project_row_tile( + TensorState const& state_vk, + TensorKTile const& k_smem, + RowTileProjPlan const& plan) { + if (!plan.owns_row) { + return 0.f; + } + return project_row_tile( + state_vk, k_smem, plan.v_row, TileCoords{plan.v_base, plan.k_base}); + } + + template + CUTE_DEVICE static float update_and_output_row_tile( + TensorState& state_vk, + TensorKTile const& k_smem, + TensorQTile const& q_smem, + int v_row, + TileCoords coords, + float decay, + float v_new) { + // Current decode path: + // - read one 16-wide state tile for the owned row into registers + // - apply decay and rank-1 update in fp32 + // - accumulate the matching out contribution against q + // - write the updated state tile back + return update_state_row_tile_and_accumulate_out( + state_vk, k_smem, q_smem, v_row, coords, decay, v_new); + } + + template + CUTE_DEVICE static float update_and_output_row_tile( + TensorState& state_vk, + TensorKTile const& k_smem, + TensorQTile const& q_smem, + RowTileUpdatePlan const& plan, + float decay, + float v_new) { + if (!plan.owns_row) { + return 0.f; + } + return update_and_output_row_tile( + state_vk, + k_smem, + q_smem, + plan.v_row, + plan.coords, + decay, + v_new); + } + + CUTE_DEVICE static float softplusf_approx(float x) { + return x > 20.f ? x : log1pf(expf(x)); + } + + CUTE_DEVICE static float warp_sum(float value) { + constexpr unsigned int kFullMask = 0xffffffffu; +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(kFullMask, value, offset); + } + return value; + } + + template < + typename TensorQ, + typename TensorK, + typename TensorV, + typename TensorA, + typename TensorB, + typename TensorAlog, + typename TensorDt, + typename TensorHvk, + typename TensorOut, + typename SharedStorage> + CUTE_DEVICE static void run( + TensorQ const& q_vec, + TensorK const& k_vec, + TensorV const& v_vec, + TensorA const& a_scalar, + TensorB const& b_scalar, + TensorAlog const& A_log_scalar, + TensorDt const& dt_bias_scalar, + TensorHvk& state_vk, + TensorOut& out_vec, + SharedStorage& storage, + int tid, + int num_threads) { + // Decode organization: + // - 1 warpgroup owns the full [128, 128] state tile for one (token, hv) + // - state is traversed as 16x16 tiles over the internal VK view + // - q/k/v are staged once into shared memory + // - proj/out are accumulated over K tiles + // - rank-1 update is applied tile-by-tile in the same traversal order + // + // This pass establishes the tile-first organization for the final fp32 + // decode kernel. The next implementation step should optimize the scalar + // inner loops with better register ownership / reductions rather than + // forcing Tensor Core math. + // + // TODO(qwen35-decode-fp32): + // - evaluate whether q/k should move from shared-memory staging to + // warp-shuffle broadcast + // - evaluate whether one thread should own more than one V-row + // - evaluate whether some parts of the state row can remain resident in + // registers across both proj and update/out passes with acceptable + // register pressure + + const float a_val = static_cast(a_scalar); + const float b_val = static_cast(b_scalar); + const float A_log_val = static_cast(A_log_scalar); + const float dt_bias_val = static_cast(dt_bias_scalar); + + const float g = -expf(A_log_val) * softplusf_approx(a_val + dt_bias_val); + const float decay = expf(g); + const float beta = 1.f / (1.f + expf(-b_val)); + + auto q_smem = make_tensor(make_smem_ptr(storage.q_smem), make_layout(make_shape(Int{}))); + auto k_smem = make_tensor(make_smem_ptr(storage.k_smem), make_layout(make_shape(Int{}))); + auto v_smem = make_tensor(make_smem_ptr(storage.v_smem), make_layout(make_shape(Int{}))); + auto norm_smem = make_tensor(make_smem_ptr(storage.norm_smem), make_layout(make_shape(Int<2>{}))); + auto proj_smem = make_tensor(make_smem_ptr(storage.proj_smem), make_layout(make_shape(Int{}))); + auto out_smem = make_tensor(make_smem_ptr(storage.out_smem), make_layout(make_shape(Int{}))); + + // Stage q/k/v once per CTA for the current decode token. + for (int idx = tid; idx < kHeadDimQK; idx += num_threads) { + q_smem(idx) = static_cast(q_vec(idx)); + k_smem(idx) = static_cast(k_vec(idx)); + } + for (int idx = tid; idx < kHeadDimV; idx += num_threads) { + v_smem(idx) = v_vec(idx); + proj_smem(idx) = 0.f; + out_smem(idx) = 0.f; + } + __syncthreads(); + + if (tid == 0) { + float q_norm_sq = 0.f; + float k_norm_sq = 0.f; +#pragma unroll + for (int idx = 0; idx < kHeadDimQK; ++idx) { + const float q_val = q_smem(idx); + const float k_val = k_smem(idx); + q_norm_sq += q_val * q_val; + k_norm_sq += k_val * k_val; + } + norm_smem(0) = rsqrtf(q_norm_sq + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); + norm_smem(1) = rsqrtf(k_norm_sq + 1e-6f); + } + __syncthreads(); + + for (int idx = tid; idx < kHeadDimQK; idx += num_threads) { + q_smem(idx) = q_smem(idx) * norm_smem(0); + k_smem(idx) = k_smem(idx) * norm_smem(1); + } + __syncthreads(); + + ThreadRowPlan row_plan = make_thread_row_plan(tid); + + // First concrete ownership model: + // - each thread owns one full state row across all 128 K columns + // - the row is streamed tile-by-tile through registers + // - no cross-thread reduction is needed for proj/out because the full row + // stays with one thread for the duration of the token update + if (row_plan.owns_row) { + float proj_row = 0.f; + for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { + TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; + RowTileProjPlan proj_plan = make_row_tile_proj_plan(coords, row_plan.warp_id, row_plan.lane_id); + proj_plan.v_row = row_plan.v_row; + proj_plan.owns_row = true; + proj_row += project_row_tile(state_vk, k_smem, proj_plan); + } + + proj_smem(row_plan.v_row) = proj_row; + + const float v_val = static_cast(v_smem(row_plan.v_row)); + const float decayed_proj_row = decay * proj_row; + const float v_new_row = beta * (v_val - decayed_proj_row); + + float out_row = 0.f; + for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { + TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; + RowTileUpdatePlan update_plan = make_row_tile_update_plan(coords, row_plan.warp_id, row_plan.lane_id); + update_plan.v_row = row_plan.v_row; + update_plan.owns_row = true; + out_row += update_and_output_row_tile( + state_vk, k_smem, q_smem, update_plan, decay, v_new_row); + } + + out_smem(row_plan.v_row) = out_row; + } + __syncthreads(); + + for (int idx = tid; idx < kHeadDimV; idx += num_threads) { + out_vec(idx) = static_cast(out_smem(idx)); + } + } +}; + +template +struct Qwen35ScalarKdaDecodeLongMainloop : public Qwen35ScalarKdaDecodeMainloop { + using Base = Qwen35ScalarKdaDecodeMainloop; + using Base::run; +}; + +} // namespace cula::qwen35::decode::kernel diff --git a/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100.hpp b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100.hpp new file mode 100644 index 00000000..6c2c8d5c --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100.hpp @@ -0,0 +1,639 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +// This header is intentionally independent from +// qwen35_scalar_kda_prefill_kernel.hpp. It is a Blackwell-only prototype for +// replacing that file's WMMA chunk state/output stage without perturbing the +// recurrent fallback or the in-flight scalar-kernel work. + +#if defined(CULA_SM100_ENABLED) + +#include +#include +#include +#include +#include +#include + +#include "kerutils/kerutils.cuh" + +namespace cula::qwen35::prefill::kernel::sm100_ts { + +using namespace cute; + +using bf16 = kerutils::bf16; + +struct alignas(16) Bf16x8 { + bf16 values[8]; +}; + +static constexpr int kHeadDim = 128; +static constexpr int kChunk = 64; +#ifndef CULA_QWEN35_TS_VALUE_TILE +#define CULA_QWEN35_TS_VALUE_TILE 128 +#endif +static constexpr int kValueTile = CULA_QWEN35_TS_VALUE_TILE; +static_assert(kValueTile == 64 || kValueTile == 128); +static constexpr int kTmemThreads = 128; +#ifndef CULA_QWEN35_TS_THREADS +#define CULA_QWEN35_TS_THREADS 352 +#endif +static constexpr int kThreads = CULA_QWEN35_TS_THREADS; + +// TMEM is addressed in 32-bit columns. TS-UMMA requires its M=64 accumulator +// to start at datapath zero, so projection and output use separate 64-column +// regions. (The DP16 packing accepted by the SS recompute mainloop is not a +// legal destination for this TS instruction.) +struct TmemAllocation { + static constexpr uint32_t kStateF32 = 0; // DP 0..31, 128 columns + static constexpr uint32_t kStateBf16 = 128; // DP 0..31, 64 columns + static constexpr uint32_t kResult = 192; // DP 0..31, 64 columns + static constexpr uint32_t kOutput = 256; // DP 0..31, 64 columns + static constexpr int kColumns = 512; +}; + +using SmemLayout64x128K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayout64x64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +// Shape-only layouts used to construct the TMEM A fragments. They are kept +// separate from the 64-row B layouts above because the full-value path uses +// an M=128 TS-UMMA while W/Qg/Aqk still have 64 token rows. +using SmemLayoutValuex128K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutValuex64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +// Logical shape [N=128, K=64]. N/MN-major storage makes both the source KG +// loads (KG is physically [token, K]) and the SMEM stores coalesced while UMMA +// performs the required transpose internally. +using SmemLayout128x64MN = decltype(coalesce( + tile_to_shape( + UMMA::Layout_MN_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using TiledMma64x64K = decltype(make_tiled_mma( + SM100_MMA_F16BF16_TS< + bf16, + bf16, + float, + kValueTile, + kChunk, + UMMA::Major::K, + UMMA::Major::K>{})); + +using TiledMma64x128MN = decltype(make_tiled_mma( + SM100_MMA_F16BF16_TS< + bf16, + bf16, + float, + kValueTile, + kHeadDim, + UMMA::Major::K, + UMMA::Major::MN>{})); + +struct alignas(128) Qwen35ChunkStateOutputSm100Shared { + // Two buffers are required because the two independent TS contractions are + // issued before a single UMMA completion wait. + alignas(128) bf16 operand_b0[kChunk * kHeadDim]; + alignas(128) bf16 operand_b1[kChunk * kHeadDim]; + float gate_exp[kChunk]; + alignas(16) cute::uint64_t mma_barrier; + alignas(16) cute::uint32_t tmem_base_ptr; +}; + +static_assert( + sizeof(Qwen35ChunkStateOutputSm100Shared) <= 48 * 1024, + "SM100 state/output prototype should remain below 48 KiB shared memory"); + +CUTE_DEVICE uint32_t pack_bf16_pair(float x0, float x1) { + union Bf16Bits { + __nv_bfloat16 value; + uint16_t bits; + } lo{}, hi{}; + lo.value = __float2bfloat16_rn(x0); + hi.value = __float2bfloat16_rn(x1); + return static_cast(lo.bits) | (static_cast(hi.bits) << 16); +} + +// Store a [K=128, V=64] FP32 state tile in its transposed TMEM representation +// [V=64, K=128], and create the packed BF16 operand-A shadow at the same time. +CUTE_DEVICE void initialize_state_tmem( + uint32_t tmem_base, + const float* __restrict__ initial_state, + int state_global_base, + int v_base, + bool has_initial_state) { + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + constexpr int kValuesPerWarp = kValueTile / 4; + const bool active = lane < kValuesPerWarp; + const int vv = warp * kValuesPerWarp + (lane & (kValuesPerWarp - 1)); + +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + uint32_t state_bf16[8]; +#pragma unroll + for (int item = 0; item < 16; ++item) { + state_values[item] = active && has_initial_state + ? initial_state[state_global_base + (kk0 + item) * kHeadDim + v_base + vv] + : 0.0f; + } +#pragma unroll + for (int item = 0; item < 8; ++item) { + state_bf16[item] = pack_bf16_pair(state_values[2 * item], state_values[2 * item + 1]); + } + kerutils::tmem_st_32dp32bNx<16>(tmem_base + TmemAllocation::kStateF32 + kk0, state_values); + kerutils::tmem_st_32dp32bNx<8>(tmem_base + TmemAllocation::kStateBf16 + kk0 / 2, state_bf16); + } + cutlass::arch::fence_view_async_tmem_store(); + kerutils::tcgen05_before_thread_sync(); +} + +template +__global__ __launch_bounds__(kThreads, 1) void qwen35_chunk_state_output_sm100_ts_kernel( + const __nv_bfloat16* __restrict__ q_norm, + const float* __restrict__ g, + const __nv_bfloat16* __restrict__ Aqk, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ u, + const __nv_bfloat16* __restrict__ kg, + const float* __restrict__ initial_state, + __nv_bfloat16* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + extern __shared__ char shared_bytes[]; + // CUDA only guarantees the base alignment of an untyped dynamic-shared + // declaration. UMMA SW128 descriptors require 128-byte alignment, so align + // the struct explicitly instead of relying on alignas to move the runtime + // base address. + const uintptr_t shared_addr = reinterpret_cast(shared_bytes); + const uintptr_t shared_aligned_addr = (shared_addr + 127u) & ~uintptr_t(127u); + auto& shared = *reinterpret_cast(shared_aligned_addr); + + const int tid = static_cast(threadIdx.x); + const int lane = tid & 31; + const int warp = tid >> 5; + + int work = static_cast(blockIdx.x); + const int value_tile = work % (kHeadDim / kValueTile); + work /= (kHeadDim / kValueTile); + const int hv = work % kLocalVHeads; + const int seq = work / kLocalVHeads; + if (seq >= batch_size) { + return; + } + + const int heads_per_group = kLocalVHeads / qk_heads; + const int qk_h = hv / heads_per_group; + const int v_base = value_tile * kValueTile; + const int state_global_base = (seq * kLocalVHeads + hv) * kHeadDim * kHeadDim; + + cute::TMEM::Allocator1Sm tmem_allocator{}; + if (warp == 0) { + tmem_allocator.allocate(TmemAllocation::kColumns, &shared.tmem_base_ptr); + // Do not hold the per-SM allocation permit for the whole kernel. + tmem_allocator.release_allocation_lock(); + } + if (tid == 0) { + cute::initialize_barrier(shared.mma_barrier, 1); + } + __syncthreads(); + + const uint32_t tmem_base = shared.tmem_base_ptr; + if (tid < kTmemThreads) { + initialize_state_tmem( + tmem_base, + initial_state, + state_global_base, + v_base, + has_initial_state); + } + __syncthreads(); + kerutils::tcgen05_after_thread_sync(); + + TiledMma64x64K mma_64x64; + TiledMma64x128MN mma_64x128; + + // Construct correctly-shaped TMEM A fragments from fake SMEM tensors. The + // fragment data pointers are then redirected to the explicit TMEM plan. + // A TS operand is a TMEM fragment; the SMEM tensor is shape-only. It must + // use a null SMEM pointer so no real shared-memory base/swizzle offset leaks + // into the generated TMEM fragment layout. + auto fake_state = make_tensor( + make_smem_ptr(static_cast(nullptr)), SmemLayoutValuex128K{}); + auto fake_vnew = make_tensor( + make_smem_ptr(static_cast(nullptr)), SmemLayoutValuex64K{}); + auto t_state_a = mma_64x64.get_slice(_0{}).partition_fragment_A(fake_state); + t_state_a.data() = tmem_base + TmemAllocation::kStateBf16; + auto t_vnew_a_64 = mma_64x64.get_slice(_0{}).partition_fragment_A(fake_vnew); + t_vnew_a_64.data() = tmem_base + TmemAllocation::kStateBf16; + auto t_vnew_a_128 = mma_64x128.get_slice(_0{}).partition_fragment_A(fake_vnew); + t_vnew_a_128.data() = tmem_base + TmemAllocation::kStateBf16; + + auto t_projection = partition_fragment_C( + mma_64x64, Shape, Int>{}); + t_projection.data() = tmem_base + TmemAllocation::kResult; + auto t_output = partition_fragment_C( + mma_64x64, Shape, Int>{}); + t_output.data() = tmem_base + TmemAllocation::kOutput; + auto t_state_acc = partition_fragment_C( + mma_64x128, Shape, Int>{}); + t_state_acc.data() = tmem_base + TmemAllocation::kStateF32; + + int barrier_phase = 0; + const int chunk_count = (seq_len + kChunk - 1) / kChunk; + const float q_scale = rsqrtf(static_cast(kHeadDim)); + const auto* q_norm_bf16 = reinterpret_cast(q_norm); + const auto* Aqk_bf16 = reinterpret_cast(Aqk); + const auto* w_bf16 = reinterpret_cast(w); + const auto* kg_bf16 = reinterpret_cast(kg); + + if constexpr (kPrefetchGate) { + // Seed the first chunk's scalar gate. Later chunks are prefetched while + // the first UMMA pair is in flight, which removes one block barrier from + // every recurrent chunk transition. + if (tid < kChunk) { + if (tid < min(kChunk, seq_len)) { + const int token = seq * seq_len + tid; + shared.gate_exp[tid] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[tid] = 0.0f; + } + } + __syncthreads(); + } + + for (int chunk = 0; chunk < chunk_count; ++chunk) { + const int chunk_start = chunk * kChunk; + const int valid_rows = min(kChunk, seq_len - chunk_start); + + if constexpr (!kPrefetchGate) { + // Keep the short-sequence specialization identical to the lower-latency + // original path; prefetching only pays back after several chunks. + if (tid < kChunk) { + if (tid < valid_rows) { + const int token = seq * seq_len + chunk_start + tid; + shared.gate_exp[tid] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[tid] = 0.0f; + } + } + __syncthreads(); + } + + // First pair: transpose(W @ state) and transpose(Qg @ state). + auto s_w = make_tensor(make_smem_ptr(shared.operand_b0), SmemLayout64x128K{}); + auto s_qg = make_tensor(make_smem_ptr(shared.operand_b1), SmemLayout64x128K{}); + constexpr int kBf16PerVector = 8; + constexpr int kHeadVectors = kHeadDim / kBf16PerVector; + for (int vector_idx = tid; + vector_idx < kChunk * kHeadVectors; + vector_idx += kThreads) { + const int row = vector_idx / kHeadVectors; + const int kk = (vector_idx % kHeadVectors) * kBf16PerVector; + Bf16x8 w_values{}; + Bf16x8 qg_values{}; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + w_values = *reinterpret_cast( + w_bf16 + (token * kLocalVHeads + hv) * kHeadDim + kk); + const Bf16x8 q_values = *reinterpret_cast( + q_norm_bf16 + (token * qk_heads + qk_h) * kHeadDim + kk); +#pragma unroll + for (int item = 0; item < kBf16PerVector; ++item) { + qg_values.values[item] = bf16( + static_cast(q_values.values[item]) * + shared.gate_exp[row] * q_scale); + } + } + *reinterpret_cast(&s_w(row, kk)) = w_values; + *reinterpret_cast(&s_qg(row, kk)) = qg_values; + } + __syncthreads(); + + if (warp == 0) { + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ts(mma_64x64, t_state_a, s_w, t_projection, true); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ts(mma_64x64, t_state_a, s_qg, t_output, true); + cutlass::arch::umma_arrive(&shared.mma_barrier); + } + // Once W/Qg have been staged, gate_exp is dead for this chunk. Use two + // non-issuer warps to prepare the next chunk while UMMA consumes SMEM. + // The epilogue's block barrier below makes these writes visible before + // the next iteration starts loading Qg. + if constexpr (kPrefetchGate) { + if (chunk + 1 < chunk_count && tid >= 32 && tid < 32 + kChunk) { + const int next_row = tid - 32; + const int next_start = chunk_start + kChunk; + const int next_valid_rows = min(kChunk, seq_len - next_start); + if (next_row < next_valid_rows) { + const int token = seq * seq_len + next_start + next_row; + shared.gate_exp[next_row] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[next_row] = 0.0f; + } + } + } + cute::wait_barrier(shared.mma_barrier, barrier_phase); + barrier_phase ^= 1; + + // Load the transposed projection in 16-column slices and create Vnew. + constexpr int kValuesPerWarp = kValueTile / 4; + const bool active_value = lane < kValuesPerWarp; + const int vv = warp * kValuesPerWarp + (lane & (kValuesPerWarp - 1)); + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); +#pragma unroll + for (int row0 = 0; row0 < kChunk; row0 += 16) { + float pair_values[16]; + uint32_t vnew_bf16[8]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kResult + row0, + pair_values); + cutlass::arch::fence_view_async_tmem_load(); +#pragma unroll + for (int item = 0; item < 8; ++item) { + float v0 = 0.0f; + float v1 = 0.0f; + if (active_value) { + const int row_a = row0 + 2 * item; + const int row_b = row_a + 1; + if (row_a < valid_rows) { + const int token = seq * seq_len + chunk_start + row_a; + v0 = __bfloat162float( + u[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv]) - + pair_values[2 * item]; + } + if (row_b < valid_rows) { + const int token = seq * seq_len + chunk_start + row_b; + v1 = __bfloat162float( + u[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv]) - + pair_values[2 * item + 1]; + } + } + vnew_bf16[item] = pack_bf16_pair(v0, v1); + } + kerutils::tmem_st_32dp32bNx<8>( + tmem_base + TmemAllocation::kStateBf16 + row0 / 2, + vnew_bf16); + } + + // Apply the chunk decay to the persistent FP32 state before accumulating + // the KG^T @ Vnew update. DP 16..31 are unused for this tile. + const int last_token = seq * seq_len + chunk_start + valid_rows - 1; + const float chunk_decay = exp2f(g[last_token * kLocalVHeads + hv]); +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + cutlass::arch::fence_view_async_tmem_load(); +#pragma unroll + for (int item = 0; item < 16; ++item) { + state_values[item] *= chunk_decay; + } + kerutils::tmem_st_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + } + cutlass::arch::fence_view_async_tmem_store(); + kerutils::tcgen05_before_thread_sync(); + } + __syncthreads(); + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); + } + + // Second pair: Aqk @ Vnew accumulates into output, while KG^T @ Vnew + // accumulates directly into the decayed FP32 state tile. + auto s_aqk = make_tensor(make_smem_ptr(shared.operand_b0), SmemLayout64x64K{}); + auto s_kg = make_tensor(make_smem_ptr(shared.operand_b1), SmemLayout128x64MN{}); + constexpr int kChunkVectors = kChunk / kBf16PerVector; + for (int vector_idx = tid; + vector_idx < kChunk * kChunkVectors; + vector_idx += kThreads) { + const int row = vector_idx / kChunkVectors; + const int col = (vector_idx % kChunkVectors) * kBf16PerVector; + Bf16x8 values{}; + if (row < valid_rows && col < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + values = *reinterpret_cast( + Aqk_bf16 + (token * kLocalVHeads + hv) * kChunk + col); + } + *reinterpret_cast(&s_aqk(row, col)) = values; + } + // Iterate in KG's physical [token, K] order; MN-major B storage keeps the + // destination N coordinate contiguous too. + for (int vector_idx = tid; + vector_idx < kChunk * kHeadVectors; + vector_idx += kThreads) { + const int row = vector_idx / kHeadVectors; + const int kk = (vector_idx % kHeadVectors) * kBf16PerVector; + Bf16x8 values{}; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + values = *reinterpret_cast( + kg_bf16 + (token * kLocalVHeads + hv) * kHeadDim + kk); + } +#pragma unroll + for (int item = 0; item < kBf16PerVector; ++item) { + s_kg(kk + item, row) = values.values[item]; + } + } + __syncthreads(); + + if (warp == 0) { + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ts(mma_64x64, t_vnew_a_64, s_aqk, t_output, false); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ts(mma_64x128, t_vnew_a_128, s_kg, t_state_acc, false); + cutlass::arch::umma_arrive(&shared.mma_barrier); + } + cute::wait_barrier(shared.mma_barrier, barrier_phase); + barrier_phase ^= 1; + + // Store the completed transposed output. Within each warp, lower-half + // lanes write consecutive V columns for a fixed token. + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); +#pragma unroll + for (int row0 = 0; row0 < kChunk; row0 += 16) { + float pair_values[16]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kOutput + row0, + pair_values); + cutlass::arch::fence_view_async_tmem_load(); + if (active_value) { +#pragma unroll + for (int item = 0; item < 16; ++item) { + const int row = row0 + item; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + out[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv] = + __float2bfloat16_rn(pair_values[item]); + } + } + } + } + + // Refresh the BF16 state shadow for the next chunk. The final chunk also + // writes the exact FP32 persistent state to the public output tensor, but + // does not need a shadow refresh because no later chunk consumes it. + const bool is_last_chunk = chunk + 1 == chunk_count; + if (!is_last_chunk) { +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + uint32_t state_bf16[8]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + cutlass::arch::fence_view_async_tmem_load(); +#pragma unroll + for (int item = 0; item < 8; ++item) { + state_bf16[item] = pack_bf16_pair(state_values[2 * item], state_values[2 * item + 1]); + } + kerutils::tmem_st_32dp32bNx<8>( + tmem_base + TmemAllocation::kStateBf16 + kk0 / 2, + state_bf16); + } + cutlass::arch::fence_view_async_tmem_store(); + kerutils::tcgen05_before_thread_sync(); + } else { +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + cutlass::arch::fence_view_async_tmem_load(); + if (active_value) { +#pragma unroll + for (int item = 0; item < 16; ++item) { + final_state[state_global_base + (kk0 + item) * kHeadDim + v_base + vv] = + state_values[item]; + } + } + } + } + } + __syncthreads(); + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); + } + } + + __syncthreads(); + if (warp == 0) { + tmem_allocator.free(tmem_base, TmemAllocation::kColumns); + } +} + +template +inline void launch_qwen35_chunk_state_output_sm100_ts_variant( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + auto kernel_fn = + &qwen35_chunk_state_output_sm100_ts_kernel; + constexpr size_t shared_bytes = sizeof(Qwen35ChunkStateOutputSm100Shared) + 127; + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes); + const int grid = batch_size * kLocalVHeads * (kHeadDim / kValueTile); + kernel_fn<<>>( + q_norm, + g, + Aqk, + w, + u, + kg, + initial_state, + out, + final_state, + batch_size, + seq_len, + qk_heads, + has_initial_state); +} + +template +inline void launch_qwen35_chunk_state_output_sm100_ts( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + if (seq_len >= 256) { + launch_qwen35_chunk_state_output_sm100_ts_variant( + stream, q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, + batch_size, seq_len, qk_heads, has_initial_state); + } else { + launch_qwen35_chunk_state_output_sm100_ts_variant( + stream, q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, + batch_size, seq_len, qk_heads, has_initial_state); + } +} + +} // namespace cula::qwen35::prefill::kernel::sm100_ts + +#endif // CULA_SM100_ENABLED diff --git a/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100_ss.hpp b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100_ss.hpp new file mode 100644 index 00000000..a520ad0a --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100_ss.hpp @@ -0,0 +1,499 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +// Standalone Blackwell SS-UMMA prototype for the state/output portion of the +// Qwen3.5 scalar prefill path. Keeping this header independent makes it +// possible to compile and inspect the replacement without changing the +// production WMMA kernel or its launcher. + +#if defined(CULA_SM100_ENABLED) + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kerutils/kerutils.cuh" + +namespace cula::qwen35::prefill::kernel::sm100_ss { + +using namespace cute; + +using bf16 = kerutils::bf16; + +static constexpr int kHeadDim = 128; +static constexpr int kChunk = 64; +static constexpr int kValueTile = 64; +static constexpr int kThreads = 128; +// Columns 0..63 hold the two 64-row output accumulators (lower/upper DP +// halves); columns 64..127 hold the independent M128 state update. +static constexpr int kTmemColumns = 128; +static constexpr uint32_t kOutputUpperDp = 16u * 65536u; +static constexpr uint32_t kStateUpdateColumn = 64u; + +// UMMA sees both operands as logical matrices. W/Qg/Aqk/KgT are A +// operands, so K-major is the natural row-major representation. State and +// Vnew are B operands with logical shape [N, K]; MN-major is required here, +// rather than treating their physical [K, V] input representation as an +// ordinary K-major matrix. +using SmemLayoutA64x128K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutStateB64x128MN = decltype(coalesce( + tile_to_shape( + UMMA::Layout_MN_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutA64x64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutVnewB64x64MN = decltype(coalesce( + tile_to_shape( + UMMA::Layout_MN_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutKgT128x64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +static_assert(cosize_v == kChunk * kHeadDim); +static_assert(cosize_v == kValueTile * kHeadDim); +static_assert(cosize_v == kChunk * kChunk); +static_assert(cosize_v == kValueTile * kChunk); +static_assert(cosize_v == kHeadDim * kChunk); + +using TiledMma64x64 = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS< + bf16, + bf16, + float, + kChunk, + kValueTile, + UMMA::Major::K, + UMMA::Major::MN>{})); + +// State update is Kg^T[M=128,K=64] @ Vnew^T[N=64,K=64]. Keeping M=128 in +// one instruction is important: splitting it into two M64 updates repeats +// the Vnew descriptor traffic and completion synchronization. +using TiledMma128x64 = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS< + bf16, + bf16, + float, + kHeadDim, + kValueTile, + UMMA::Major::K, + UMMA::Major::MN>{})); + +using CompletionPipeline = cutlass::PipelineUmmaAsync<1>; +using CompletionPipelineState = cutlass::PipelineState; +using ClusterShape = Shape<_1, _1, _1>; + +struct alignas(128) Qwen35ChunkStateOutputSm100SsShared { + // Persistent state is laid out as B[N=V,K=head_dim]. + alignas(128) bf16 state[kValueTile * kHeadDim]; + // W/Aqk use this K-major A buffer. It is reused as Aqk after the first + // dual UMMA pair has completed. + alignas(128) bf16 operand_a[kHeadDim * kChunk]; + // Qg/Kg^T use a second K-major A buffer. Keeping the two A operands + // separate allows each pair of contractions to share one completion wait. + alignas(128) bf16 operand_a_aux[kHeadDim * kChunk]; + // Vnew is the MN-major B operand for both Aqk and the M128 state update. + alignas(128) bf16 vnew[kValueTile * kChunk]; + float gate_exp[kChunk]; + alignas(16) uint32_t tmem_base_ptr; + alignas(16) typename CompletionPipeline::SharedStorage completion; +}; + +static_assert( + sizeof(Qwen35ChunkStateOutputSm100SsShared) <= 64 * 1024, + "SS-UMMA state/output prototype must remain below 64 KiB shared memory"); + +CUTE_DEVICE void release_ss_mma_result( + CompletionPipeline& completion, + CompletionPipelineState& consumer_state) { + // TMEM loads performed by the epilogue must become visible before the + // consumer marks the single pipeline stage reusable. + kerutils::tcgen05_before_thread_sync(); + completion.consumer_release(consumer_state); + ++consumer_state; + __syncthreads(); +} + +template +__global__ __launch_bounds__(kThreads, 1) void qwen35_chunk_state_output_sm100_ss_kernel( + const __nv_bfloat16* __restrict__ q_norm, + const float* __restrict__ g, + const __nv_bfloat16* __restrict__ Aqk, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ u, + const __nv_bfloat16* __restrict__ kg, + const float* __restrict__ initial_state, + __nv_bfloat16* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + extern __shared__ char shared_bytes[]; + auto& shared = *reinterpret_cast(shared_bytes); + + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + + int work = static_cast(blockIdx.x); + const int value_tile = work % (kHeadDim / kValueTile); + work /= (kHeadDim / kValueTile); + const int hv = work % kLocalVHeads; + const int seq = work / kLocalVHeads; + if (seq >= batch_size) { + return; + } + + const int qk_h = hv / (kLocalVHeads / qk_heads); + const int v_base = value_tile * kValueTile; + const int state_global_base = (seq * kLocalVHeads + hv) * kHeadDim * kHeadDim; + + auto s_state = make_tensor( + make_smem_ptr(shared.state), SmemLayoutStateB64x128MN{}); + auto s_a_64x128 = make_tensor( + make_smem_ptr(shared.operand_a), SmemLayoutA64x128K{}); + auto s_a_qg = make_tensor( + make_smem_ptr(shared.operand_a_aux), SmemLayoutA64x128K{}); + auto s_a_aqk = make_tensor( + make_smem_ptr(shared.operand_a), SmemLayoutA64x64K{}); + auto s_a_kg = make_tensor( + make_smem_ptr(shared.operand_a_aux), SmemLayoutKgT128x64K{}); + auto s_a_128x64 = make_tensor( + make_smem_ptr(shared.operand_a), SmemLayoutKgT128x64K{}); + auto s_vnew = make_tensor( + make_smem_ptr(shared.vnew), SmemLayoutVnewB64x64MN{}); + + for (int index = tid; index < kValueTile * kHeadDim; index += kThreads) { + const int vv = index / kHeadDim; + const int kk = index % kHeadDim; + const float value = has_initial_state + ? initial_state[state_global_base + kk * kHeadDim + v_base + vv] + : 0.0f; + s_state(vv, kk) = bf16(value); + } + + // One completion pipeline is sufficient because the four contractions are + // issued as two dual-UMMA pairs. Warp 0 both issues UMMA and participates + // in the 128-thread epilogue. + typename CompletionPipeline::Params completion_params; + completion_params.producer_arv_count = 1; + completion_params.consumer_arv_count = kThreads; + completion_params.initializing_warp = 0; + completion_params.role = warp == 0 + ? CompletionPipeline::ThreadCategory::ProducerConsumer + : CompletionPipeline::ThreadCategory::Consumer; + CompletionPipeline completion( + shared.completion, completion_params, ClusterShape{}); + + cute::TMEM::Allocator1Sm tmem_allocator{}; + if (warp == 0) { + tmem_allocator.allocate(kTmemColumns, &shared.tmem_base_ptr); + tmem_allocator.release_allocation_lock(); + } + __syncthreads(); + + TiledMma64x64 mma_64x64; + TiledMma128x64 mma_128x64; + auto t_acc_64x64_lower = partition_fragment_C( + mma_64x64, Shape, Int>{}); + auto t_acc_64x64_upper = partition_fragment_C( + mma_64x64, Shape, Int>{}); + auto t_acc_128x64_state = partition_fragment_C( + mma_128x64, Shape, Int>{}); + t_acc_64x64_lower.data() = shared.tmem_base_ptr; + t_acc_64x64_upper.data() = shared.tmem_base_ptr + kOutputUpperDp; + t_acc_128x64_state.data() = + shared.tmem_base_ptr + kStateUpdateColumn; + + auto c64 = make_identity_tensor( + Shape, Int>{}); + auto c128 = make_identity_tensor( + Shape, Int>{}); + auto t_c64 = mma_64x64.get_slice(_0{}).partition_C(c64); + auto t_c128 = mma_128x64.get_slice(_0{}).partition_C(c128); + + CompletionPipelineState producer_state = + cutlass::make_producer_start_state(); + CompletionPipelineState consumer_state; + + const auto* q_bf16 = reinterpret_cast(q_norm); + const auto* aqk_bf16 = reinterpret_cast(Aqk); + const auto* w_bf16 = reinterpret_cast(w); + const auto* u_bf16 = reinterpret_cast(u); + const auto* kg_bf16 = reinterpret_cast(kg); + const int chunk_count = (seq_len + kChunk - 1) / kChunk; + const float q_scale = rsqrtf(static_cast(kHeadDim)); + + for (int chunk = 0; chunk < chunk_count; ++chunk) { + const int chunk_start = chunk * kChunk; + const int valid_rows = min(kChunk, seq_len - chunk_start); + + if (tid < kChunk) { + if (tid < valid_rows) { + const int token = seq * seq_len + chunk_start + tid; + shared.gate_exp[tid] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[tid] = 0.0f; + } + } + + // Pair 1: projection = W[64,128] @ state^T[128,64] in the lower DP + // half, and output_base = Qg[64,128] @ state^T in the upper DP half. + // Both A operands are staged before the single UMMA completion signal. + for (int index = tid; index < kChunk * kHeadDim; index += kThreads) { + const int row = index / kHeadDim; + const int kk = index % kHeadDim; + bf16 value = bf16(0.0f); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = w_bf16[(token * kLocalVHeads + hv) * kHeadDim + kk]; + } + s_a_64x128(row, kk) = value; + float q_value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + q_value = static_cast( + q_bf16[(token * qk_heads + qk_h) * kHeadDim + kk]) * + shared.gate_exp[row] * q_scale; + } + s_a_qg(row, kk) = bf16(q_value); + } + __syncthreads(); + + if (warp == 0) { + completion.producer_acquire(producer_state); + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ss( + mma_64x64, + s_a_64x128, + s_state, + t_acc_64x64_lower, + true); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ss( + mma_64x64, + s_a_qg, + s_state, + t_acc_64x64_upper, + true); + completion.producer_commit(producer_state); + ++producer_state; + } + completion.consumer_wait(consumer_state); + kerutils::tcgen05_after_thread_sync(); + + { + auto tiled_t2r = make_tmem_copy( + SM100_TMEM_LOAD_16dp256b8x{}, t_acc_64x64_lower); + auto thr_t2r = tiled_t2r.get_slice(tid); + auto t_src = thr_t2r.partition_S(t_acc_64x64_lower); + auto t_coord = thr_t2r.partition_D(t_c64); + auto r_acc = make_tensor(shape(t_coord)); + copy(tiled_t2r, t_src, r_acc); + cutlass::arch::fence_view_async_tmem_load(); + CUTE_UNROLL + for (int item = 0; item < size(r_acc); ++item) { + const auto coord = t_coord(item); + const int row = static_cast(get<0>(coord)); + const int vv = static_cast(get<1>(coord)); + float value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = static_cast( + u_bf16[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv]) - + r_acc(item); + } + s_vnew(vv, row) = bf16(value); + } + } + cutlass::arch::fence_view_async_shared(); + release_ss_mma_result(completion, consumer_state); + + // Pair 2: output += Aqk[64,64] @ Vnew^T in the upper DP half, while the + // state update Kg^T[128,64] @ Vnew^T is accumulated in a separate M128 + // TMEM fragment. The two independent UMMAs share one completion wait. + for (int index = tid; index < kChunk * kChunk; index += kThreads) { + const int row = index / kChunk; + const int col = index % kChunk; + bf16 value = bf16(0.0f); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = aqk_bf16[(token * kLocalVHeads + hv) * kChunk + col]; + } + s_a_aqk(row, col) = value; + } + for (int index = tid; index < kHeadDim * kChunk; index += kThreads) { + const int kk = index / kChunk; + const int row = index % kChunk; + bf16 value = bf16(0.0f); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = kg_bf16[(token * kLocalVHeads + hv) * kHeadDim + kk]; + } + s_a_kg(kk, row) = value; + } + const int last_token = seq * seq_len + chunk_start + valid_rows - 1; + const float chunk_decay = exp2f(g[last_token * kLocalVHeads + hv]); + __syncthreads(); + + if (warp == 0) { + completion.producer_acquire(producer_state); + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ss( + mma_64x64, + s_a_aqk, + s_vnew, + t_acc_64x64_upper, + false); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ss( + mma_128x64, + s_a_kg, + s_vnew, + t_acc_128x64_state, + true); + completion.producer_commit(producer_state); + ++producer_state; + } + completion.consumer_wait(consumer_state); + kerutils::tcgen05_after_thread_sync(); + + { + auto tiled_t2r = make_tmem_copy( + SM100_TMEM_LOAD_16dp256b8x{}, t_acc_64x64_upper); + auto thr_t2r = tiled_t2r.get_slice(tid); + auto t_src = thr_t2r.partition_S(t_acc_64x64_upper); + auto t_coord = thr_t2r.partition_D(t_c64); + auto r_acc = make_tensor(shape(t_coord)); + copy(tiled_t2r, t_src, r_acc); + cutlass::arch::fence_view_async_tmem_load(); + CUTE_UNROLL + for (int item = 0; item < size(r_acc); ++item) { + const auto coord = t_coord(item); + const int row = static_cast(get<0>(coord)); + const int vv = static_cast(get<1>(coord)); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + out[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv] = + __float2bfloat16_rn(r_acc(item)); + } + } + } + + { + auto tiled_t2r = make_tmem_copy( + SM100_TMEM_LOAD_32dp32b16x{}, t_acc_128x64_state); + auto thr_t2r = tiled_t2r.get_slice(tid); + auto t_src = thr_t2r.partition_S(t_acc_128x64_state); + auto t_coord = thr_t2r.partition_D(t_c128); + auto r_acc = make_tensor(shape(t_coord)); + copy(tiled_t2r, t_src, r_acc); + cutlass::arch::fence_view_async_tmem_load(); + const bool is_last_chunk = chunk + 1 == chunk_count; + CUTE_UNROLL + for (int item = 0; item < size(r_acc); ++item) { + const auto coord = t_coord(item); + const int kk = static_cast(get<0>(coord)); + const int vv = static_cast(get<1>(coord)); + const float updated = + r_acc(item) + chunk_decay * static_cast(s_state(vv, kk)); + const bf16 quantized = bf16(updated); + s_state(vv, kk) = quantized; + if (is_last_chunk) { + final_state[state_global_base + kk * kHeadDim + v_base + vv] = + static_cast(quantized); + } + } + } + cutlass::arch::fence_view_async_shared(); + release_ss_mma_result(completion, consumer_state); + } + + __syncthreads(); + if (warp == 0) { + tmem_allocator.free(shared.tmem_base_ptr, kTmemColumns); + } +} + +template +inline void launch_qwen35_chunk_state_output_sm100_ss( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + auto kernel_fn = &qwen35_chunk_state_output_sm100_ss_kernel; + constexpr size_t shared_bytes = + sizeof(Qwen35ChunkStateOutputSm100SsShared); + cudaFuncSetAttribute( + kernel_fn, + cudaFuncAttributeMaxDynamicSharedMemorySize, + shared_bytes); + const int grid = + batch_size * kLocalVHeads * (kHeadDim / kValueTile); + kernel_fn<<>>( + q_norm, + g, + Aqk, + w, + u, + kg, + initial_state, + out, + final_state, + batch_size, + seq_len, + qk_heads, + has_initial_state); +} + +} // namespace cula::qwen35::prefill::kernel::sm100_ss + +#endif // CULA_SM100_ENABLED diff --git a/csrc/qwen35/prefill/qwen35_layout_prefill.cu b/csrc/qwen35/prefill/qwen35_layout_prefill.cu new file mode 100644 index 00000000..5cdf1218 --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_layout_prefill.cu @@ -0,0 +1,188 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#include "qwen35_layout_prefill_kernel.hpp" +#include "qwen35_prefill_common.cuh" + +#include +#include +#include +#include + +namespace cula::qwen35::prefill { + +namespace { + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +void check_rank_2(const at::Tensor& tensor, const char* name) { + TORCH_CHECK(tensor.dim() == 2, name, " must be rank 2, got rank ", tensor.dim(), "."); +} + +template +void dispatch_layout_prefill_for_heads( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t token_count) { + constexpr int kLocalQKHeads = decode::local_qk_heads_from_v_heads(kLocalVHeads); + kernel::launch_qwen35_layout_prefill_kernel( + stream, + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + token_count); +} + +template +void dispatch_layout_prefill( + int64_t local_v_heads, + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t token_count) { + switch (local_v_heads) { + case 48: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + case 24: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + case 12: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + case 6: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + } +} + +} // namespace + +void run_qwen35_layout_prefill(LayoutPrefillParams& params) { + const at::Tensor& mixed_qkv_conv = params.mixed_qkv_conv; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& q_rep = params.q_rep; + const at::Tensor& k_rep = params.k_rep; + const at::Tensor& v = params.v; + const at::Tensor& a_kernel = params.a_kernel; + const at::Tensor& b_kernel = params.b_kernel; + + TORCH_CHECK(mixed_qkv_conv.is_cuda(), "mixed_qkv_conv must be a CUDA tensor."); + const at::Device device = mixed_qkv_conv.device(); + + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(q_rep, "q_rep", device); + check_tensor_device(k_rep, "k_rep", device); + check_tensor_device(v, "v", device); + check_tensor_device(a_kernel, "a_kernel", device); + check_tensor_device(b_kernel, "b_kernel", device); + + TORCH_CHECK(mixed_qkv_conv.is_contiguous(), "mixed_qkv_conv must be contiguous."); + TORCH_CHECK(a.is_contiguous(), "a must be contiguous."); + TORCH_CHECK(b.is_contiguous(), "b must be contiguous."); + TORCH_CHECK(q_rep.is_contiguous(), "q_rep must be contiguous."); + TORCH_CHECK(k_rep.is_contiguous(), "k_rep must be contiguous."); + TORCH_CHECK(v.is_contiguous(), "v must be contiguous."); + TORCH_CHECK(a_kernel.is_contiguous(), "a_kernel must be contiguous."); + TORCH_CHECK(b_kernel.is_contiguous(), "b_kernel must be contiguous."); + + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == a.scalar_type() && + mixed_qkv_conv.scalar_type() == b.scalar_type() && + mixed_qkv_conv.scalar_type() == q_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == k_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == v.scalar_type() && + mixed_qkv_conv.scalar_type() == a_kernel.scalar_type() && + mixed_qkv_conv.scalar_type() == b_kernel.scalar_type(), + "All layout prefill tensors must share the same dtype."); + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == at::kHalf || mixed_qkv_conv.scalar_type() == at::kBFloat16, + "mixed_qkv_conv must be float16 or bfloat16."); + + check_rank_2(mixed_qkv_conv, "mixed_qkv_conv"); + check_rank_2(a, "a"); + check_rank_2(b, "b"); + + const int64_t token_count = mixed_qkv_conv.size(0); + const int64_t local_v_heads = a.size(1); + TORCH_CHECK(decode::is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + const int local_qk_heads = decode::local_qk_heads_from_v_heads(static_cast(local_v_heads)); + const int local_mixed_dim = decode::local_mixed_qkv_dim(local_qk_heads, static_cast(local_v_heads)); + TORCH_CHECK(mixed_qkv_conv.size(1) == local_mixed_dim, "mixed_qkv_conv must be [N, local_conv_dim=", local_mixed_dim, "]."); + TORCH_CHECK(a.sizes() == at::IntArrayRef({token_count, local_v_heads}), "a must be [N, local_v_heads]."); + TORCH_CHECK(b.sizes() == at::IntArrayRef({token_count, local_v_heads}), "b must be [N, local_v_heads]."); + TORCH_CHECK( + q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimQK}), + "q_rep must be [N, local_v_heads, 128]."); + TORCH_CHECK(k_rep.sizes() == q_rep.sizes(), "k_rep must match q_rep shape."); + TORCH_CHECK(v.sizes() == q_rep.sizes(), "v must match q_rep shape."); + TORCH_CHECK(a_kernel.sizes() == a.sizes(), "a_kernel must match a shape."); + TORCH_CHECK(b_kernel.sizes() == b.sizes(), "b_kernel must match b shape."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + if (mixed_qkv_conv.scalar_type() == at::kHalf) { + dispatch_layout_prefill( + local_v_heads, + stream, + mixed_qkv_conv.data_ptr(), + a.data_ptr(), + b.data_ptr(), + q_rep.data_ptr(), + k_rep.data_ptr(), + v.data_ptr(), + a_kernel.data_ptr(), + b_kernel.data_ptr(), + token_count); + } else { + dispatch_layout_prefill( + local_v_heads, + stream, + mixed_qkv_conv.data_ptr(), + a.data_ptr(), + b.data_ptr(), + q_rep.data_ptr(), + k_rep.data_ptr(), + v.data_ptr(), + a_kernel.data_ptr(), + b_kernel.data_ptr(), + token_count); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::prefill diff --git a/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp new file mode 100644 index 00000000..7dbcae3d --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp @@ -0,0 +1,144 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include "qwen35_prefill_common.cuh" + +#include +#include +#include +#include + +namespace cula::qwen35::prefill::kernel { + +using namespace cute; + +template +CUTE_DEVICE void copy_prefill_vec_contiguous( + scalar_t* __restrict__ dst, + const scalar_t* __restrict__ src) { + constexpr int kBytes = sizeof(scalar_t) * kVec; + if constexpr (kBytes == 16 || kBytes == 8) { + using VecType = cutlass::AlignedArray; + const auto dst_addr = reinterpret_cast(dst); + const auto src_addr = reinterpret_cast(src); + if ((dst_addr % alignof(VecType) == 0) && (src_addr % alignof(VecType) == 0)) { + *reinterpret_cast(dst) = *reinterpret_cast(src); + return; + } + } + +#pragma unroll + for (int i = 0; i < kVec; ++i) { + dst[i] = src[i]; + } +} + +template +__global__ void qwen35_layout_prefill_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + scalar_t* __restrict__ q_rep, + scalar_t* __restrict__ k_rep, + scalar_t* __restrict__ v_out, + scalar_t* __restrict__ a_kernel, + scalar_t* __restrict__ b_kernel, + int64_t token_count) { + static_assert(kLocalVHeads % kLocalQKHeads == 0); + static_assert(kHeadDimQK == kHeadDimV); + constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; + constexpr int kLocalQDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalKDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalMixedQKVDim = 2 * kLocalQDim + kLocalVHeads * kHeadDimV; + constexpr int kVec = 4; + static_assert(kHeadDimQK % kVec == 0); + + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + if (token_idx >= token_count || hv >= kLocalVHeads) { + return; + } + + const int mapped_h = hv / kRepeatFactor; + + auto qk_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto hv_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + + const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kLocalMixedQKVDim; + const scalar_t* q_src_ptr = token_ptr; + const scalar_t* k_src_ptr = token_ptr + kLocalQDim; + const scalar_t* v_src_ptr = token_ptr + kLocalQDim + kLocalKDim; + + scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kLocalVHeads * kHeadDimV; + + for (int vec_idx = tid; vec_idx < kHeadDimQK / kVec; vec_idx += blockDim.x) { + const int d = vec_idx * kVec; + const int q_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int k_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int v_src_idx = crd2idx(make_coord(hv, d), v_src_layout); + const int dst_idx = crd2idx(make_coord(hv, d), hv_layout); + + copy_prefill_vec_contiguous(q_dst_ptr + dst_idx, q_src_ptr + q_src_idx); + copy_prefill_vec_contiguous(k_dst_ptr + dst_idx, k_src_ptr + k_src_idx); + copy_prefill_vec_contiguous(v_dst_ptr + dst_idx, v_src_ptr + v_src_idx); + } + + if (tid == 0) { + const int head_idx = crd2idx(make_coord(hv), head_layout); + const int64_t token_head_offset = static_cast(token_idx) * kLocalVHeads + head_idx; + a_kernel[token_head_offset] = a[token_head_offset]; + b_kernel[token_head_offset] = b[token_head_offset]; + } +} + +template +void launch_qwen35_layout_prefill_kernel( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t token_count) { + constexpr int kThreads = 32; + dim3 grid(kLocalVHeads, static_cast(token_count), 1); + qwen35_layout_prefill_kernel<<>>( + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + token_count); +} + +} // namespace cula::qwen35::prefill::kernel diff --git a/csrc/qwen35/prefill/qwen35_prefill_common.cuh b/csrc/qwen35/prefill/qwen35_prefill_common.cuh new file mode 100644 index 00000000..ca5c2e16 --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_prefill_common.cuh @@ -0,0 +1,93 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include "qwen35/decode/qwen35_decode_common.cuh" + +#include + +namespace cula::qwen35::prefill { + +using decode::kHeadDimQK; +using decode::kHeadDimV; +using decode::kKDim; +using decode::kMixedQKVDim; +using decode::kNumQKHeads; +using decode::kNumVHeads; +using decode::kQDim; +using decode::kVDim; + +struct LayoutPrefillParams { + at::Tensor mixed_qkv_conv; // [N, local_conv_dim] + at::Tensor a; // [N, local_v_heads] + at::Tensor b; // [N, local_v_heads] + at::Tensor q_rep; // [N, local_v_heads, 128] + at::Tensor k_rep; // [N, local_v_heads, 128] + at::Tensor v; // [N, local_v_heads, 128] + at::Tensor a_kernel; // [N, local_v_heads] + at::Tensor b_kernel; // [N, local_v_heads] +}; + +struct ScalarKdaPrefillParams { + at::Tensor q; // [B, T, local_qk_heads, 128] + at::Tensor k; // [B, T, local_qk_heads, 128] + at::Tensor v; // [B, T, local_v_heads, 128] + at::Tensor a; // [B, T, local_v_heads] + at::Tensor b; // [B, T, local_v_heads] + at::Tensor A_log; // [local_v_heads], float32 + at::Tensor dt_bias; // [local_v_heads], float32 + at::Tensor initial_state; // [N, local_v_heads, 128, 128], float32, may be empty + at::Tensor cu_seqlens; // [N + 1], int32, may be empty + at::Tensor out; // [B, T, local_v_heads, 128] + at::Tensor final_state; // [N, local_v_heads, 128, 128], float32 +}; + +// Core-only ABI used for apples-to-apples comparison with SGLang's +// TritonGDNKernel.extend. g/beta are the already materialized per-token +// scalar gate and beta tensors; q/k normalization and chunk-local gate scan +// remain inside the CUDA prefill calculation. +struct ScalarKdaPrefillCoreParams { + at::Tensor q; // [B, T, local_qk_heads, 128], bf16 + at::Tensor k; // [B, T, local_qk_heads, 128], bf16 + at::Tensor v; // [B, T, local_v_heads, 128], bf16 + at::Tensor g; // [B, T, local_v_heads], float32, natural-log gate + at::Tensor beta; // [B, T, local_v_heads], float32 + at::Tensor initial_state; // [N, local_v_heads, 128, 128], float32, may be empty + at::Tensor cu_seqlens; // [N + 1], int32, may be empty + at::Tensor out; // [B, T, local_v_heads, 128], bf16 + at::Tensor final_state; // [N, local_v_heads, 128, 128], float32 +}; + +// All local V-head counts produced by the downloaded Qwen3.5/Qwen3.6 +// configurations at TP={1,2,4,8}. The scalar path accepts compact native-GVA +// Q/K heads and maps each V head to its Q/K group inside the kernel. +inline constexpr bool is_supported_scalar_prefill_v_heads(int local_v_heads) { + return local_v_heads == 64 || local_v_heads == 48 || local_v_heads == 32 || + local_v_heads == 24 || local_v_heads == 16 || local_v_heads == 12 || + local_v_heads == 8 || local_v_heads == 6 || local_v_heads == 4 || + local_v_heads == 2; +} + +void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params); +void run_qwen35_scalar_kda_prefill_core(ScalarKdaPrefillCoreParams& params); +void run_qwen35_layout_prefill(LayoutPrefillParams& params); + +} // namespace cula::qwen35::prefill + +namespace cula::qwen35::prefill::sm90 { + +void qwen35_chunk_qk_prefill_sm90(at::Tensor q, at::Tensor k, at::Tensor out); + +} // namespace cula::qwen35::prefill::sm90 diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu new file mode 100644 index 00000000..b18b4e80 --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu @@ -0,0 +1,889 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#include "qwen35_prefill_common.cuh" +#include "qwen35_scalar_kda_prefill_kernel.hpp" +#ifdef CULA_SM90A_ENABLED +#include "kda/sm90/prefill_kernel.hpp" +#endif +#ifdef CULA_SM100_ENABLED +#include "kda/sm100/kda_fwd_common.cuh" +#include "qwen35_chunk_state_output_sm100.hpp" +#include "qwen35_chunk_state_output_sm100_ss.hpp" +#endif + +#include +#include +#include +#include + +#include + +namespace cula::qwen35::prefill { + +namespace { + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + if (tensor.defined() && tensor.numel() > 0) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); + } +} + +void check_contiguous(const at::Tensor& tensor, const char* name) { + if (tensor.defined() && tensor.numel() > 0) { + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous."); + } +} + +template +void dispatch_scalar_prefill_for_heads( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + kernel::launch_qwen35_scalar_kda_prefill_kernel( + stream, + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + batch_size, + seq_len, + qk_heads, + sequence_count, + is_varlen, + has_initial_state); +} + +template +void dispatch_scalar_prefill( + int64_t local_v_heads, + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + switch (local_v_heads) { + case 64: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 48: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 32: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 24: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 16: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 12: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 8: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 6: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 4: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 2: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + default: + TORCH_CHECK(false, "unsupported scalar prefill local V-head count: ", local_v_heads); + } +} + +template +void dispatch_scalar_prefill_precomputed_fallback( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const float* g, + const float* beta, + const float* initial_state, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + int local_v_heads, + bool has_initial_state, + const int32_t* unsafe_gate_flags) { +#define CULA_QWEN35_FALLBACK_CASE(HV) \ + case HV: \ + kernel::launch_qwen35_scalar_kda_prefill_precomputed_fallback( \ + stream, q, k, v, g, beta, initial_state, out, final_state, \ + batch_size, seq_len, qk_heads, has_initial_state, unsafe_gate_flags); \ + break + switch (local_v_heads) { + CULA_QWEN35_FALLBACK_CASE(64); + CULA_QWEN35_FALLBACK_CASE(48); + CULA_QWEN35_FALLBACK_CASE(32); + CULA_QWEN35_FALLBACK_CASE(24); + CULA_QWEN35_FALLBACK_CASE(16); + CULA_QWEN35_FALLBACK_CASE(12); + CULA_QWEN35_FALLBACK_CASE(8); + CULA_QWEN35_FALLBACK_CASE(6); + CULA_QWEN35_FALLBACK_CASE(4); + CULA_QWEN35_FALLBACK_CASE(2); + default: + TORCH_CHECK(false, "Unsupported local V head count: ", local_v_heads, "."); + } +#undef CULA_QWEN35_FALLBACK_CASE +} + +#ifdef CULA_SM90A_ENABLED +void run_qwen35_chunk_prefill_sm90_bf16( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& a, + const at::Tensor& b, + const at::Tensor& A_log, + const at::Tensor& dt_bias, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + int local_v_heads, + bool has_initial_state, + cudaStream_t stream, + const at::Tensor* precomputed_g = nullptr, + const at::Tensor* precomputed_beta = nullptr) { + const bool use_sm90_chunk = local_v_heads % 4 == 0; + TORCH_CHECK( + (precomputed_g == nullptr) == (precomputed_beta == nullptr), + "precomputed gate and beta must be supplied together."); + constexpr int chunk_size = kernel::kChunkSize; + const int chunks_per_sequence = (seq_len + chunk_size - 1) / chunk_size; + const int total_chunks = batch_size * chunks_per_sequence; + const int64_t total_tokens = static_cast(batch_size) * seq_len; + const auto bf16_options = q.options().dtype(at::kBFloat16); + const auto fp32_options = q.options().dtype(at::kFloat); + const auto int_options = q.options().dtype(at::kInt); + + at::Tensor q_norm = at::empty_like(q, bf16_options); + at::Tensor k_norm = at::empty_like(k, bf16_options); + at::Tensor g = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + at::Tensor g_raw = precomputed_g == nullptr + ? at::empty({batch_size, seq_len, local_v_heads}, fp32_options) + : *precomputed_g; + // Always materialize beta into private workspace. The core ABI input may be + // aliased or reused concurrently and must remain read-only. + at::Tensor beta = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + // HV=6/2 TP shards cannot satisfy Hopper TMA's four-adjacent-head scalar + // gate transaction. Keep the core ABI correct by marking every head for the + // exact recurrent path; divisible-by-four shapes use speculative SM90 KDA. + at::Tensor unsafe_gate_flags = use_sm90_chunk + ? at::zeros({batch_size, local_v_heads}, int_options) + : at::ones({batch_size, local_v_heads}, int_options); + at::Tensor cu_work = at::empty({batch_size + 1}, int_options); + at::Tensor chunk_indices = at::empty({total_chunks, 2}, int_options); + + kernel::launch_qwen35_chunk_preprocess( + stream, + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + precomputed_g == nullptr + ? reinterpret_cast(a.data_ptr()) + : nullptr, + precomputed_g == nullptr + ? reinterpret_cast(b.data_ptr()) + : nullptr, + precomputed_g == nullptr ? A_log.data_ptr() : nullptr, + precomputed_g == nullptr ? dt_bias.data_ptr() : nullptr, + precomputed_g == nullptr ? nullptr : precomputed_g->data_ptr(), + precomputed_beta == nullptr ? nullptr : precomputed_beta->data_ptr(), + precomputed_g != nullptr, + reinterpret_cast<__nv_bfloat16*>(q_norm.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(k_norm.data_ptr()), + g.data_ptr(), + precomputed_g == nullptr ? g_raw.data_ptr() : nullptr, + beta.data_ptr(), + unsafe_gate_flags.data_ptr(), + cu_work.data_ptr(), + chunk_indices.data_ptr(), + batch_size, + seq_len, + qk_heads, + local_v_heads); + + if (use_sm90_chunk) { + const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + at::Tensor workspace = + at::empty({static_cast(sm_count) * 128}, q.options().dtype(at::kByte)); + kda::sm90::launch_qwen35_scalar_kda_fwd_prefill_kernel( + stream, + out.data_ptr(), + final_state.data_ptr(), + q_norm.data_ptr(), + k_norm.data_ptr(), + v.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + g.data_ptr(), + beta.data_ptr(), + cu_work.data_ptr(), + workspace.data_ptr(), + batch_size, + qk_heads, + local_v_heads, + kHeadDimQK, + total_tokens, + rsqrtf(static_cast(kHeadDimQK)), + has_initial_state, + sm_count); + } + + // The SM90 fully-fused safe-gate algebra assumes raw log gates in [-5, 0]. + // Preprocess marks unsafe (sequence, V-head) pairs while it already reads + // the raw gates. A lightweight recurrent launch returns immediately for + // safe heads; unsafe heads overwrite the speculative fast result exactly, + // without a host sync or Python-side branch. + dispatch_scalar_prefill_precomputed_fallback( + stream, + q_norm.data_ptr(), + k_norm.data_ptr(), + v.data_ptr(), + g_raw.data_ptr(), + beta.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + out.data_ptr(), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + local_v_heads, + has_initial_state, + unsafe_gate_flags.data_ptr()); +} +#endif + +#ifdef CULA_SM100_ENABLED +// The TS-UMMA implementation is the default SM100 chunk state/output path. +// Keep compile-time escape hatches for A/B comparisons with the WMMA and +// standalone SS prototypes without changing the Python ABI or launch args. +#ifndef CULA_QWEN35_USE_WMMA_CHUNK +#define CULA_QWEN35_USE_WMMA_CHUNK 0 +#endif +#ifndef CULA_QWEN35_USE_TS_CHUNK +#define CULA_QWEN35_USE_TS_CHUNK 1 +#endif + +template +void launch_chunk_state_output_for_heads( + cudaStream_t stream, + const at::Tensor& q_norm, + const at::Tensor& g, + const at::Tensor& Aqk, + const at::Tensor& w, + const at::Tensor& u, + const at::Tensor& kg, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { +#if CULA_QWEN35_USE_WMMA_CHUNK + kernel::launch_qwen35_chunk_state_output( + stream, + reinterpret_cast(q_norm.data_ptr()), + g.data_ptr(), + reinterpret_cast(Aqk.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(u.data_ptr()), + reinterpret_cast(kg.data_ptr()), + has_initial_state ? initial_state.data_ptr() : nullptr, + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + has_initial_state); +#elif CULA_QWEN35_USE_TS_CHUNK + kernel::sm100_ts::launch_qwen35_chunk_state_output_sm100_ts( + stream, + reinterpret_cast(q_norm.data_ptr()), + g.data_ptr(), + reinterpret_cast(Aqk.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(u.data_ptr()), + reinterpret_cast(kg.data_ptr()), + has_initial_state ? initial_state.data_ptr() : nullptr, + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + has_initial_state); +#else + kernel::sm100_ss::launch_qwen35_chunk_state_output_sm100_ss( + stream, + reinterpret_cast(q_norm.data_ptr()), + g.data_ptr(), + reinterpret_cast(Aqk.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(u.data_ptr()), + reinterpret_cast(kg.data_ptr()), + has_initial_state ? initial_state.data_ptr() : nullptr, + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + has_initial_state); +#endif +} + +void launch_chunk_state_output( + int64_t local_v_heads, + cudaStream_t stream, + const at::Tensor& q_norm, + const at::Tensor& g, + const at::Tensor& Aqk, + const at::Tensor& w, + const at::Tensor& u, + const at::Tensor& kg, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { +#define CULA_QWEN35_CHUNK_HEAD_CASE(HV) \ + case HV: \ + launch_chunk_state_output_for_heads( \ + stream, q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, batch_size, seq_len, \ + qk_heads, has_initial_state); \ + break + switch (local_v_heads) { + CULA_QWEN35_CHUNK_HEAD_CASE(64); + CULA_QWEN35_CHUNK_HEAD_CASE(48); + CULA_QWEN35_CHUNK_HEAD_CASE(32); + CULA_QWEN35_CHUNK_HEAD_CASE(24); + CULA_QWEN35_CHUNK_HEAD_CASE(16); + CULA_QWEN35_CHUNK_HEAD_CASE(12); + CULA_QWEN35_CHUNK_HEAD_CASE(8); + CULA_QWEN35_CHUNK_HEAD_CASE(6); + CULA_QWEN35_CHUNK_HEAD_CASE(4); + CULA_QWEN35_CHUNK_HEAD_CASE(2); + default: + TORCH_CHECK(false, "unsupported chunk prefill local V-head count: ", local_v_heads); + } +#undef CULA_QWEN35_CHUNK_HEAD_CASE +} + +void run_qwen35_chunk_prefill_bf16( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& a, + const at::Tensor& b, + const at::Tensor& A_log, + const at::Tensor& dt_bias, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + int local_v_heads, + bool has_initial_state, + cudaStream_t stream, + const at::Tensor* precomputed_g = nullptr, + const at::Tensor* precomputed_beta = nullptr) { + TORCH_CHECK( + (precomputed_g == nullptr) == (precomputed_beta == nullptr), + "precomputed gate and beta must be supplied together."); + constexpr int chunk_size = kernel::kChunkSize; + const int chunks_per_sequence = (seq_len + chunk_size - 1) / chunk_size; + const int total_chunks = batch_size * chunks_per_sequence; + const int64_t total_tokens = static_cast(batch_size) * seq_len; + const auto bf16_options = q.options().dtype(at::kBFloat16); + const auto fp32_options = q.options().dtype(at::kFloat); + const auto int_options = q.options().dtype(at::kInt); + + // These tensors are genuine CUDA workspaces; no Python/reference operation + // participates in the numerical result. They are deliberately explicit + // while the chunk path is stabilized, and can later be supplied by an + // inference workspace pool without changing the kernels. + at::Tensor q_norm = at::empty({batch_size, seq_len, qk_heads, kHeadDimQK}, bf16_options); + at::Tensor k_norm = at::empty_like(q_norm); + // Qwen GDN has one scalar gate per token/value-head. Keep it compact and + // route only this adapter through the scalar-G KDA specializations. + at::Tensor g = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + at::Tensor beta = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + at::Tensor cu_work = at::empty({batch_size + 1}, int_options); + at::Tensor chunk_indices = at::empty({total_chunks, 2}, int_options); + at::Tensor Aqk = at::empty({batch_size, seq_len, local_v_heads, chunk_size}, bf16_options); + at::Tensor Akk = at::empty_like(Aqk); + at::Tensor w = at::empty({batch_size, seq_len, local_v_heads, kHeadDimQK}, bf16_options); + at::Tensor u = at::empty_like(w); + at::Tensor kg = at::empty_like(w); + + kernel::launch_qwen35_chunk_preprocess( + stream, + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + precomputed_g == nullptr + ? reinterpret_cast(a.data_ptr()) + : nullptr, + precomputed_g == nullptr + ? reinterpret_cast(b.data_ptr()) + : nullptr, + precomputed_g == nullptr ? A_log.data_ptr() : nullptr, + precomputed_g == nullptr ? dt_bias.data_ptr() : nullptr, + precomputed_g == nullptr ? nullptr : precomputed_g->data_ptr(), + precomputed_beta == nullptr ? nullptr : precomputed_beta->data_ptr(), + precomputed_g != nullptr && precomputed_beta != nullptr, + reinterpret_cast<__nv_bfloat16*>(q_norm.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(k_norm.data_ptr()), + g.data_ptr(), + nullptr, + beta.data_ptr(), + nullptr, + cu_work.data_ptr(), + chunk_indices.data_ptr(), + batch_size, + seq_len, + qk_heads, + local_v_heads); + + auto* device_prop = at::cuda::getCurrentDeviceProperties(); + const int scheduler_tiles = total_chunks * local_v_heads; + const int scheduler_sms = std::min(device_prop->multiProcessorCount, scheduler_tiles); + KDA_fwd_intra_params intra{}; + intra.total_q_len = static_cast(total_tokens); + intra.b = batch_size; + intra.h_qk = qk_heads; + intra.h_v = local_v_heads; + intra.heads_per_group = local_v_heads / qk_heads; + intra.d = kHeadDimQK; + intra.chunk_size = chunk_size; + intra.scale = rsqrtf(static_cast(kHeadDimQK)); + intra.use_tf32_inverse = false; + intra.unified_gref = true; + intra.is_beta_bf16 = false; + intra.q_ptr = q_norm.data_ptr(); + intra.k_ptr = k_norm.data_ptr(); + intra.g_ptr = g.data_ptr(); + intra.beta_ptr = beta.data_ptr(); + intra.Aqk_out_ptr = Aqk.data_ptr(); + intra.Akk_out_ptr = Akk.data_ptr(); + intra.cu_seqlens_ptr = cu_work.data_ptr(); + intra.chunk_indices_ptr = chunk_indices.data_ptr(); + intra.shape_Akk = cute::make_shape(intra.total_q_len, chunk_size, local_v_heads); + intra.stride_Akk = cute::make_stride(chunk_size * local_v_heads, cute::_1{}, chunk_size); + intra.num_sm = scheduler_sms; + intra.tile_scheduler_params = StaticPersistentTileScheduler::Params{ + total_chunks, + local_v_heads, + intra.heads_per_group, + intra.num_sm, + nullptr}; + kda::sm100::run_kda_fwd_intra_sm100_qwen_scalar_g(intra, stream); + + KDA_fwd_recomp_w_u_params recomp{}; + recomp.total_len = static_cast(total_tokens); + recomp.b = batch_size; + recomp.h_qk = qk_heads; + recomp.h_v = local_v_heads; + recomp.heads_per_group = local_v_heads / qk_heads; + recomp.d = kHeadDimQK; + recomp.chunk_size = chunk_size; + recomp.is_beta_bf16 = false; + recomp.k_ptr = k_norm.data_ptr(); + recomp.v_ptr = v.data_ptr(); + recomp.q_ptr = q_norm.data_ptr(); + recomp.beta_ptr = beta.data_ptr(); + recomp.A_ptr = Akk.data_ptr(); + recomp.g_ptr = g.data_ptr(); + recomp.cu_seqlens_ptr = cu_work.data_ptr(); + recomp.chunk_indices_ptr = chunk_indices.data_ptr(); + recomp.w_out_ptr = w.data_ptr(); + recomp.u_out_ptr = u.data_ptr(); + recomp.kg_out_ptr = kg.data_ptr(); + recomp.qg_out_ptr = nullptr; + recomp.store_qg = false; + recomp.shape_wukg = cute::make_shape(recomp.total_len, kHeadDimQK, local_v_heads); + recomp.stride_wukg = cute::make_stride(kHeadDimQK * local_v_heads, cute::_1{}, kHeadDimQK); + recomp.num_sm = scheduler_sms; + recomp.tile_scheduler_params = StaticPersistentTileScheduler::Params{ + total_chunks, local_v_heads, recomp.heads_per_group, recomp.num_sm, nullptr}; + kda::sm100::run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g(recomp, stream); + + launch_chunk_state_output( + local_v_heads, + stream, + q_norm, + g, + Aqk, + w, + u, + kg, + initial_state, + out, + final_state, + batch_size, + seq_len, + qk_heads, + has_initial_state); +} +#endif + +} // namespace + +void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { + const at::Tensor& q = params.q; + const at::Tensor& k = params.k; + const at::Tensor& v = params.v; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& A_log = params.A_log; + const at::Tensor& dt_bias = params.dt_bias; + const at::Tensor& initial_state = params.initial_state; + const at::Tensor& cu_seqlens = params.cu_seqlens; + const at::Tensor& out = params.out; + const at::Tensor& final_state = params.final_state; + + TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor."); + const at::Device device = q.device(); + + check_tensor_device(k, "k", device); + check_tensor_device(v, "v", device); + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(A_log, "A_log", device); + check_tensor_device(dt_bias, "dt_bias", device); + check_tensor_device(initial_state, "initial_state", device); + check_tensor_device(cu_seqlens, "cu_seqlens", device); + check_tensor_device(out, "out", device); + check_tensor_device(final_state, "final_state", device); + + check_contiguous(q, "q"); + check_contiguous(k, "k"); + check_contiguous(v, "v"); + check_contiguous(a, "a"); + check_contiguous(b, "b"); + check_contiguous(A_log, "A_log"); + check_contiguous(dt_bias, "dt_bias"); + check_contiguous(initial_state, "initial_state"); + check_contiguous(cu_seqlens, "cu_seqlens"); + check_contiguous(out, "out"); + check_contiguous(final_state, "final_state"); + + TORCH_CHECK( + q.scalar_type() == k.scalar_type() && q.scalar_type() == v.scalar_type() && + q.scalar_type() == a.scalar_type() && q.scalar_type() == b.scalar_type() && + q.scalar_type() == out.scalar_type(), + "q/k/v/a/b/out must share the same dtype."); + TORCH_CHECK(q.scalar_type() == at::kHalf || q.scalar_type() == at::kBFloat16, "q must be float16 or bfloat16."); + TORCH_CHECK(A_log.scalar_type() == at::kFloat, "A_log must be float32."); + TORCH_CHECK(dt_bias.scalar_type() == at::kFloat, "dt_bias must be float32."); + TORCH_CHECK(final_state.scalar_type() == at::kFloat, "final_state must be float32."); + TORCH_CHECK( + !initial_state.defined() || initial_state.numel() == 0 || initial_state.scalar_type() == at::kFloat, + "initial_state must be float32 when provided."); + TORCH_CHECK( + !cu_seqlens.defined() || cu_seqlens.numel() == 0 || cu_seqlens.scalar_type() == at::kInt, + "cu_seqlens must be int32 when provided."); + + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, "q/k/v must be 4D."); + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t qk_heads = q.size(2); + const int64_t local_v_heads = v.size(2); + TORCH_CHECK(qk_heads > 0 && local_v_heads > 0, "q/k/v head counts must be positive."); + TORCH_CHECK(local_v_heads % qk_heads == 0, "local V heads must be divisible by local Q/K heads."); + TORCH_CHECK( + is_supported_scalar_prefill_v_heads(static_cast(local_v_heads)), + "unsupported Qwen scalar prefill local V-head count: ", local_v_heads, "."); + TORCH_CHECK( + q.sizes() == at::IntArrayRef({B, T, qk_heads, kHeadDimQK}), + "q must have shape [B, T, local_qk_heads, 128]."); + TORCH_CHECK(k.sizes() == q.sizes(), "k must match q shape."); + TORCH_CHECK( + v.sizes() == at::IntArrayRef({B, T, local_v_heads, kHeadDimV}), + "v must have shape [B, T, local_v_heads, 128]."); + TORCH_CHECK(a.dim() == 3 && a.sizes() == at::IntArrayRef({B, T, local_v_heads}), "a must be [B, T, local_v_heads]."); + TORCH_CHECK(b.sizes() == a.sizes(), "b must match a shape."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == local_v_heads, "A_log must be [local_v_heads]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == local_v_heads, "dt_bias must be [local_v_heads]."); + TORCH_CHECK(out.sizes() == v.sizes(), "out must match v shape."); + + const bool is_varlen = cu_seqlens.defined() && cu_seqlens.numel() > 0; + const int64_t sequence_count = is_varlen ? cu_seqlens.numel() - 1 : B; + TORCH_CHECK(sequence_count > 0, "sequence_count must be positive."); + if (is_varlen) { + TORCH_CHECK(B == 1, "cu_seqlens mode expects flattened q/k/v with batch size 1."); + } + + TORCH_CHECK( + final_state.dim() == 4 && + final_state.sizes() == at::IntArrayRef({sequence_count, local_v_heads, kHeadDimQK, kHeadDimV}), + "final_state must be [N, local_v_heads, 128, 128]."); + const bool has_initial_state = initial_state.defined() && initial_state.numel() > 0; + if (has_initial_state) { + TORCH_CHECK(initial_state.sizes() == final_state.sizes(), "initial_state must match final_state shape."); + } + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM90A_ENABLED) + // A single packed sequence is equivalent to the fixed-length B=1 layout, + // so it can use the same chunk scheduler without reading cu_seqlens back to + // the host. True multi-sequence varlen remains on the recurrent fallback. + const bool fixed_like_layout = !is_varlen || sequence_count == 1; +#ifdef CULA_SM90A_ENABLED + const bool chunk_head_supported = local_v_heads % 4 == 0; +#else + constexpr bool chunk_head_supported = true; +#endif + if (q.scalar_type() == at::kBFloat16 && T >= 32 && fixed_like_layout && chunk_head_supported) { +#ifdef CULA_SM100_ENABLED + run_qwen35_chunk_prefill_bf16( +#else + run_qwen35_chunk_prefill_sm90_bf16( +#endif + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + out, + final_state, + static_cast(B), + static_cast(T), + static_cast(qk_heads), + static_cast(local_v_heads), + has_initial_state, + stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } +#endif + + if (q.scalar_type() == at::kHalf) { + dispatch_scalar_prefill( + local_v_heads, + stream, + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + a.data_ptr(), + b.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + is_varlen ? cu_seqlens.data_ptr() : nullptr, + out.data_ptr(), + final_state.data_ptr(), + static_cast(B), + static_cast(T), + static_cast(qk_heads), + static_cast(sequence_count), + is_varlen, + has_initial_state); + } else { + dispatch_scalar_prefill( + local_v_heads, + stream, + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + a.data_ptr(), + b.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + is_varlen ? cu_seqlens.data_ptr() : nullptr, + out.data_ptr(), + final_state.data_ptr(), + static_cast(B), + static_cast(T), + static_cast(qk_heads), + static_cast(sequence_count), + is_varlen, + has_initial_state); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void run_qwen35_scalar_kda_prefill_core(ScalarKdaPrefillCoreParams& params) { +#if !defined(CULA_SM100_ENABLED) && !defined(CULA_SM90A_ENABLED) + TORCH_CHECK(false, "Qwen scalar GDN prefill core requires an SM90 or SM100 build."); +#else + const at::Tensor& q = params.q; + const at::Tensor& k = params.k; + const at::Tensor& v = params.v; + const at::Tensor& gate_raw = params.g; + const at::Tensor& beta_raw = params.beta; + const at::Tensor& initial_state = params.initial_state; + const at::Tensor& cu_seqlens = params.cu_seqlens; + const at::Tensor& out = params.out; + const at::Tensor& final_state = params.final_state; + + TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor."); + const at::Device device = q.device(); + check_tensor_device(k, "k", device); + check_tensor_device(v, "v", device); + check_tensor_device(gate_raw, "g", device); + check_tensor_device(beta_raw, "beta", device); + check_tensor_device(initial_state, "initial_state", device); + check_tensor_device(cu_seqlens, "cu_seqlens", device); + check_tensor_device(out, "out", device); + check_tensor_device(final_state, "final_state", device); + + check_contiguous(q, "q"); + check_contiguous(k, "k"); + check_contiguous(v, "v"); + check_contiguous(gate_raw, "g"); + check_contiguous(beta_raw, "beta"); + check_contiguous(initial_state, "initial_state"); + check_contiguous(cu_seqlens, "cu_seqlens"); + check_contiguous(out, "out"); + check_contiguous(final_state, "final_state"); + + TORCH_CHECK( + q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && + v.scalar_type() == at::kBFloat16 && out.scalar_type() == at::kBFloat16, + "q/k/v/out must be bfloat16 for the chunk core path."); + TORCH_CHECK(gate_raw.scalar_type() == at::kFloat, "g must be float32."); + TORCH_CHECK(beta_raw.scalar_type() == at::kFloat, "beta must be float32."); + TORCH_CHECK(final_state.scalar_type() == at::kFloat, "final_state must be float32."); + TORCH_CHECK( + !initial_state.defined() || initial_state.numel() == 0 || initial_state.scalar_type() == at::kFloat, + "initial_state must be float32 when provided."); + TORCH_CHECK( + !cu_seqlens.defined() || cu_seqlens.numel() == 0 || cu_seqlens.scalar_type() == at::kInt, + "cu_seqlens must be int32 when provided."); + + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, "q/k/v must be 4D."); + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t qk_heads = q.size(2); + const int64_t local_v_heads = v.size(2); + TORCH_CHECK(T >= 32, "the chunk core path requires sequence length >= 32."); + TORCH_CHECK(qk_heads > 0 && local_v_heads > 0, "q/k/v head counts must be positive."); + TORCH_CHECK(local_v_heads % qk_heads == 0, "local V heads must be divisible by local Q/K heads."); + TORCH_CHECK( + is_supported_scalar_prefill_v_heads(static_cast(local_v_heads)), + "unsupported Qwen scalar prefill local V-head count: ", local_v_heads, "."); + TORCH_CHECK( + q.sizes() == at::IntArrayRef({B, T, qk_heads, kHeadDimQK}), + "q must have shape [B, T, local_qk_heads, 128]."); + TORCH_CHECK(k.sizes() == q.sizes(), "k must match q shape."); + TORCH_CHECK( + v.sizes() == at::IntArrayRef({B, T, local_v_heads, kHeadDimV}), + "v must have shape [B, T, local_v_heads, 128]."); + TORCH_CHECK( + gate_raw.sizes() == at::IntArrayRef({B, T, local_v_heads}), + "g must be [B, T, local_v_heads]."); + TORCH_CHECK(beta_raw.sizes() == gate_raw.sizes(), "beta must match g shape."); + TORCH_CHECK(out.sizes() == v.sizes(), "out must match v shape."); + + const bool is_varlen = cu_seqlens.defined() && cu_seqlens.numel() > 0; + const int64_t sequence_count = is_varlen ? cu_seqlens.numel() - 1 : B; + TORCH_CHECK(sequence_count > 0, "sequence_count must be positive."); + TORCH_CHECK(!is_varlen || (B == 1 && sequence_count == 1), + "the chunk core path supports fixed batches or one packed sequence."); + TORCH_CHECK( + final_state.dim() == 4 && + final_state.sizes() == at::IntArrayRef({sequence_count, local_v_heads, kHeadDimQK, kHeadDimV}), + "final_state must be [N, local_v_heads, 128, 128]."); + const bool has_initial_state = initial_state.defined() && initial_state.numel() > 0; + if (has_initial_state) { + TORCH_CHECK(initial_state.sizes() == final_state.sizes(), "initial_state must match final_state shape."); + } + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const at::Tensor empty; +#ifdef CULA_SM100_ENABLED + run_qwen35_chunk_prefill_bf16( +#else + run_qwen35_chunk_prefill_sm90_bf16( +#endif + q, + k, + v, + empty, + empty, + empty, + empty, + initial_state, + out, + final_state, + static_cast(B), + static_cast(T), + static_cast(qk_heads), + static_cast(local_v_heads), + has_initial_state, + stream, + &gate_raw, + &beta_raw); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +#endif +} + +} // namespace cula::qwen35::prefill diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp new file mode 100644 index 00000000..3c60b10d --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp @@ -0,0 +1,1208 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include "qwen35_prefill_common.cuh" + +#include +#include +#include +#include + +namespace cula::qwen35::prefill::kernel { + +using namespace cute; + +template +struct Qwen35ScalarKdaPrefillKernel { + static constexpr int kWarpSize = 32; + static constexpr int kWarps = kWarpsPerBlock; + static constexpr int kThreads = kWarps * kWarpSize; + static constexpr int kKValuesPerLane = kHeadDimQK / kWarpSize; + static constexpr int kHeadDim = kHeadDimQK; + static constexpr int kColumnsPerWarp = 1; + // Each warp owns one independent V column and keeps its complete recurrent + // state in registers. All warps in the CTA reuse one normalized Q/K vector + // and one scalar gate through shared memory. + static constexpr int kVTile = kWarps; + static constexpr int kNumVTiles = kHeadDimV / kVTile; + + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kThreads % kWarpSize == 0); + static_assert(kHeadDimQK % kWarpSize == 0); + static_assert(kHeadDimV % kVTile == 0); + + struct SharedStorage { + float q_norm[kHeadDimQK]; + float k_norm[kHeadDimQK]; + float decay; + float beta; + int unsafe_gate; + }; + + static dim3 block_shape() { + return dim3(kThreads, 1, 1); + } + + CUTE_HOST_DEVICE static auto make_v_work_tiles(int sequence_count) { + auto problem_layout = make_layout( + make_shape(Int{}, Int{}, sequence_count), + make_stride(Int<1>{}, Int{}, Int{})); + return zipped_divide(problem_layout, make_shape(Int{}, Int<1>{}, Int<1>{})); + } + + static dim3 grid_shape(int sequence_count) { + auto v_work_tiles = make_v_work_tiles(sequence_count); + return dim3(static_cast(size<1>(v_work_tiles)), 1, 1); + } + + CUTE_DEVICE static float load_as_float(scalar_t value) { + return static_cast(value); + } + + CUTE_DEVICE static scalar_t cast_output(float value) { + return static_cast(value); + } + + // A one-warp specialization avoids CTA barriers altogether. It is useful + // for the high-HV/small-T regime where the extra Q/K/gate work is cheaper + // than synchronizing a multi-warp V tile. + CUTE_DEVICE static void run_warp_only( + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ initial_state, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state, + const float* __restrict__ precomputed_g = nullptr, + const float* __restrict__ precomputed_beta = nullptr, + const int32_t* __restrict__ unsafe_gate_flags = nullptr) { + const int lane = static_cast(threadIdx.x) & 31; + int work = static_cast(blockIdx.x); + const int v_row = work % kHeadDimV; + work /= kHeadDimV; + const int hv = work % kLocalVHeads; + const int seq_idx = work / kLocalVHeads; + if (seq_idx >= sequence_count) { + return; + } + const int repeat = kLocalVHeads / qk_heads; + const int qk_h = hv / repeat; + const int token_begin = is_varlen ? static_cast(cu_seqlens[seq_idx]) : seq_idx * seq_len; + const int token_end = is_varlen ? static_cast(cu_seqlens[seq_idx + 1]) : token_begin + seq_len; + const int state_base = ((seq_idx * kLocalVHeads + hv) * kHeadDimQK) * kHeadDimV; + if (precomputed_g != nullptr) { + if (unsafe_gate_flags != nullptr) { + if (unsafe_gate_flags[seq_idx * kLocalVHeads + hv] == 0) { + return; + } + } else { + bool unsafe_gate = false; + for (int token = token_begin + lane; token < token_end; token += kWarpSize) { + const float gate = precomputed_g[token * kLocalVHeads + hv]; + unsafe_gate = unsafe_gate || !isfinite(gate) || gate < -5.0f || gate > 0.0f; + } + if (!__any_sync(0xffffffffu, unsafe_gate)) { + return; + } + } + } + float state_vals[kKValuesPerLane]; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + state_vals[item] = has_initial_state ? initial_state[state_base + kk * kHeadDimV + v_row] : 0.0f; + } + const float scale = rsqrtf(static_cast(kHeadDimQK)); + const float exp_A = precomputed_g == nullptr ? expf(A_log[hv]) : 0.0f; + const float dt = precomputed_g == nullptr ? dt_bias[hv] : 0.0f; + for (int token = token_begin; token < token_end; ++token) { + const int local_t = token - token_begin; + const int qk_base = ((token * qk_heads + qk_h) * kHeadDimQK); + const int v_input_base = ((token * kLocalVHeads + hv) * kHeadDimV); + float q_vals[kKValuesPerLane]; + float k_vals[kKValuesPerLane]; + float q_norm_sq = 0.0f; + float k_norm_sq = 0.0f; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + q_vals[item] = load_as_float(q[qk_base + kk]); + k_vals[item] = load_as_float(k[qk_base + kk]); + q_norm_sq += q_vals[item] * q_vals[item]; + k_norm_sq += k_vals[item] * k_vals[item]; + } + // The SM90 speculative path passes its BF16-normalized Q/K workspace to + // the exact overwrite. Reuse those values verbatim so both paths have + // identical normalization and rounding semantics. + const float q_rnorm = precomputed_g != nullptr + ? scale + : rsqrtf(fmaxf(warp_sum(q_norm_sq), 1.0e-20f)) * scale; + const float k_rnorm = precomputed_g != nullptr + ? 1.0f + : rsqrtf(fmaxf(warp_sum(k_norm_sq), 1.0e-20f)); + float decay = 0.0f; + float beta = 0.0f; + if (lane == 0) { + const int gate_base = token * kLocalVHeads + hv; + if (precomputed_g != nullptr) { + decay = expf(precomputed_g[gate_base]); + beta = precomputed_beta[gate_base]; + } else { + decay = expf(-exp_A * softplus(load_as_float(a[gate_base]) + dt)); + beta = 1.0f / (1.0f + expf(-load_as_float(b[gate_base]))); + } + } + decay = __shfl_sync(0xffffffffu, decay, 0); + beta = __shfl_sync(0xffffffffu, beta, 0); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + k_vals[item] *= k_rnorm; + q_vals[item] *= q_rnorm; + } + float proj_partial = 0.0f; + float out_partial = 0.0f; + const float v_val = __shfl_sync(0xffffffffu, lane == 0 ? load_as_float(v[v_input_base + v_row]) : 0.0f, 0); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + proj_partial += state_vals[item] * k_vals[item]; + } + const float proj = warp_sum(proj_partial); + const float v_new = beta * (v_val - decay * proj); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + state_vals[item] = decay * state_vals[item] + k_vals[item] * v_new; + out_partial += state_vals[item] * q_vals[item]; + } + const float out_acc = warp_sum(out_partial); + if (lane == 0) { + const int out_off = (token * kLocalVHeads + hv) * kHeadDimV + v_row; + out[out_off] = cast_output(out_acc); + } + } +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + final_state[state_base + kk * kHeadDimV + v_row] = state_vals[item]; + } + (void)batch_size; + } + + CUTE_DEVICE static float softplus(float x) { + return x > 20.0f ? x : log1pf(expf(x)); + } + + CUTE_DEVICE static float warp_sum(float value) { +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return __shfl_sync(0xffffffffu, value, 0); + } + + CUTE_DEVICE static void run_device( + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ initial_state, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state, + const float* __restrict__ precomputed_g, + const float* __restrict__ precomputed_beta, + const int32_t* __restrict__ unsafe_gate_flags, + SharedStorage& storage) { + if constexpr (kWarpsPerBlock == 1) { + run_warp_only( + q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, + batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state, + precomputed_g, precomputed_beta, unsafe_gate_flags); + return; + } + auto v_work_tiles = make_v_work_tiles(sequence_count); + auto work_layout = make_layout(get<1>(v_work_tiles.shape()), LayoutLeft{}); + auto work_coord = work_layout.get_hier_coord(static_cast(blockIdx.x)); + const int v_tile_idx = static_cast(get<0>(work_coord)); + const int hv = static_cast(get<1>(work_coord)); + const int seq_idx = static_cast(get<2>(work_coord)); + const int v_base = v_tile_idx * kVTile; + const int tid = static_cast(threadIdx.x); + const int warp = tid / kWarpSize; + const int lane = tid % kWarpSize; + + if (hv >= kLocalVHeads || seq_idx >= sequence_count) { + return; + } + + const int token_begin = is_varlen ? static_cast(cu_seqlens[seq_idx]) : seq_idx * seq_len; + const int token_end = is_varlen ? static_cast(cu_seqlens[seq_idx + 1]) : token_begin + seq_len; + const int state_base = ((seq_idx * kLocalVHeads + hv) * kHeadDimQK) * kHeadDimV; + const int qk_h = hv / (kLocalVHeads / qk_heads); + + // A compact SM90 safe-gate chunk is exact for per-token log gates in + // [-5, 0]. This recurrent fallback is launched after the fast kernel and + // overwrites only heads whose raw gate falls outside that domain. + if (precomputed_g != nullptr) { + if (unsafe_gate_flags != nullptr) { + if (unsafe_gate_flags[seq_idx * kLocalVHeads + hv] == 0) { + return; + } + } else { + if (tid == 0) { + storage.unsafe_gate = 0; + } + __syncthreads(); + for (int token = token_begin + tid; token < token_end; token += kThreads) { + const float gate = precomputed_g[token * kLocalVHeads + hv]; + if (!isfinite(gate) || gate < -5.0f || gate > 0.0f) { + atomicExch(&storage.unsafe_gate, 1); + } + } + __syncthreads(); + if (storage.unsafe_gate == 0) { + return; + } + } + } + + float state_vals[kColumnsPerWarp][kKValuesPerLane]; + +#pragma unroll + for (int column = 0; column < kColumnsPerWarp; ++column) { + const int v_row = v_base + warp * kColumnsPerWarp + column; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + const int state_off = state_base + kk * kHeadDimV + v_row; + state_vals[column][item] = has_initial_state ? initial_state[state_off] : 0.0f; + } + } + + const float scale = rsqrtf(static_cast(kHeadDimQK)); + const float exp_A = precomputed_g == nullptr ? expf(A_log[hv]) : 0.0f; + const float dt = precomputed_g == nullptr ? dt_bias[hv] : 0.0f; + + for (int token = token_begin; token < token_end; ++token) { + const int qk_base = ((token * qk_heads + qk_h) * kHeadDimQK); + const int v_base_input = ((token * kLocalVHeads + hv) * kHeadDimV); + const int gate_base = token * kLocalVHeads + hv; + + if (warp == 0) { + float q_vals_raw[kKValuesPerLane]; + float k_vals_raw[kKValuesPerLane]; + float q_norm_sq = 0.0f; + float k_norm_sq = 0.0f; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + q_vals_raw[item] = load_as_float(q[qk_base + kk]); + k_vals_raw[item] = load_as_float(k[qk_base + kk]); + q_norm_sq += q_vals_raw[item] * q_vals_raw[item]; + k_norm_sq += k_vals_raw[item] * k_vals_raw[item]; + } + q_norm_sq = warp_sum(q_norm_sq); + k_norm_sq = warp_sum(k_norm_sq); + const float q_rnorm = precomputed_g != nullptr + ? scale + : rsqrtf(fmaxf(q_norm_sq, 1.0e-20f)) * scale; + const float k_rnorm = precomputed_g != nullptr + ? 1.0f + : rsqrtf(fmaxf(k_norm_sq, 1.0e-20f)); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + storage.q_norm[kk] = q_vals_raw[item] * q_rnorm; + storage.k_norm[kk] = k_vals_raw[item] * k_rnorm; + } + if (lane == 0) { + if (precomputed_g != nullptr) { + storage.decay = expf(precomputed_g[gate_base]); + storage.beta = precomputed_beta[gate_base]; + } else { + storage.decay = expf(-exp_A * softplus(load_as_float(a[gate_base]) + dt)); + storage.beta = 1.0f / (1.0f + expf(-load_as_float(b[gate_base]))); + } + } + } + __syncthreads(); + + float q_vals[kKValuesPerLane]; + float k_vals[kKValuesPerLane]; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + q_vals[item] = storage.q_norm[kk]; + k_vals[item] = storage.k_norm[kk]; + } + const float decay = storage.decay; + const float beta = storage.beta; + +#pragma unroll + for (int column = 0; column < kColumnsPerWarp; ++column) { + const int v_row = v_base + warp * kColumnsPerWarp + column; + float proj_partial = 0.0f; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + proj_partial += state_vals[column][item] * k_vals[item]; + } + const float proj = warp_sum(proj_partial); + + float v_val = lane == 0 ? load_as_float(v[v_base_input + v_row]) : 0.0f; + v_val = __shfl_sync(0xffffffffu, v_val, 0); + const float v_new = beta * (v_val - decay * proj); + + float out_partial = 0.0f; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const float state_new = decay * state_vals[column][item] + k_vals[item] * v_new; + state_vals[column][item] = state_new; + out_partial += state_new * q_vals[item]; + } + const float out_acc = warp_sum(out_partial); + + if (lane == 0) { + const int out_off = (token * kLocalVHeads + hv) * kHeadDimV + v_row; + out[out_off] = cast_output(out_acc); + } + } + __syncthreads(); + } + +#pragma unroll + for (int column = 0; column < kColumnsPerWarp; ++column) { + const int v_row = v_base + warp * kColumnsPerWarp + column; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + const int state_off = state_base + kk * kHeadDimV + v_row; + final_state[state_off] = state_vals[column][item]; + } + } + + (void)batch_size; + } +}; + +// The recurrent scalar kernel above is still the lowest-latency path for very +// short prompts. For real prefill lengths, process time in 64-token chunks +// and use tensor cores for the state/output contractions. The preceding +// native-GVA intra/WY stages provide Aqk, W, U, and Kg for this kernel. +static constexpr int kChunkSize = 64; +static constexpr int kValueTile = 64; +static constexpr int kValueTilesPerBlock = kValueTile / 16; +static constexpr int kChunkOutputWarps = 16; +static constexpr int kChunkStateWarps = kChunkOutputWarps; + +struct alignas(128) Qwen35ChunkStateOutputShared { + __nv_bfloat16 state[kHeadDimQK * kValueTile]; + __nv_bfloat16 matrix[kChunkSize * kHeadDimQK]; + __nv_bfloat16 v_new[kChunkSize * kValueTile]; + float gate_exp[kChunkSize]; + float accum[kHeadDimQK * kValueTile]; + __nv_bfloat16 aqk[kChunkSize * kChunkSize]; +}; + +__device__ __forceinline__ float qwen35_bf16_to_float(__nv_bfloat16 value) { + return __bfloat162float(value); +} + +__device__ __forceinline__ __nv_bfloat16 qwen35_float_to_bf16(float value) { + return __float2bfloat16_rn(value); +} + +__device__ __forceinline__ float qwen35_warp_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__global__ void qwen35_chunk_qk_norm_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + __nv_bfloat16* __restrict__ q_norm, + __nv_bfloat16* __restrict__ k_norm, + int vector_count) { + const int vector_idx = static_cast(blockIdx.x); + const int lane = static_cast(threadIdx.x); + if (vector_idx >= vector_count) { + return; + } + const int base = vector_idx * kHeadDimQK; + float q_values[4]; + float k_values[4]; + float q_sq = 0.0f; + float k_sq = 0.0f; +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_values[item] = qwen35_bf16_to_float(q[base + kk]); + k_values[item] = qwen35_bf16_to_float(k[base + kk]); + q_sq += q_values[item] * q_values[item]; + k_sq += k_values[item] * k_values[item]; + } + const float q_rnorm = rsqrtf(qwen35_warp_sum(q_sq) + 1.0e-6f); + const float k_rnorm = rsqrtf(qwen35_warp_sum(k_sq) + 1.0e-6f); +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_norm[base + kk] = qwen35_float_to_bf16(q_values[item] * q_rnorm); + k_norm[base + kk] = qwen35_float_to_bf16(k_values[item] * k_rnorm); + } +} + +__device__ __forceinline__ float qwen35_softplus(float x) { + return x > 20.0f ? x : log1pf(expf(x)); +} + +__global__ void qwen35_chunk_gate_kernel( + const __nv_bfloat16* __restrict__ a, + const __nv_bfloat16* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ g, + float* __restrict__ beta, + int32_t* __restrict__ cu_seqlens, + int32_t* __restrict__ chunk_indices, + int batch_size, + int seq_len, + int v_heads, + int chunks_per_sequence) { + __shared__ float scan[kChunkSize]; + const int tid = static_cast(threadIdx.x); + int work = static_cast(blockIdx.x); + const int hv = work % v_heads; + work /= v_heads; + const int chunk = work % chunks_per_sequence; + const int seq = work / chunks_per_sequence; + const int local_t = chunk * kChunkSize + tid; + const bool valid = tid < kChunkSize && local_t < seq_len; + const int token = seq * seq_len + local_t; + + if (tid < kChunkSize) { + float log2_decay = 0.0f; + if (valid) { + const int gate_offset = token * v_heads + hv; + const float raw_a = qwen35_bf16_to_float(a[gate_offset]); + const float raw_b = qwen35_bf16_to_float(b[gate_offset]); + const float log_decay = -expf(A_log[hv]) * qwen35_softplus(raw_a + dt_bias[hv]); + log2_decay = log_decay * 1.4426950408889634f; + beta[gate_offset] = 1.0f / (1.0f + expf(-raw_b)); + } + scan[tid] = log2_decay; + } + __syncthreads(); + +#pragma unroll + for (int offset = 1; offset < kChunkSize; offset <<= 1) { + float addend = 0.0f; + if (tid < kChunkSize && tid >= offset) { + addend = scan[tid - offset]; + } + __syncthreads(); + if (tid < kChunkSize) { + scan[tid] += addend; + } + __syncthreads(); + } + + if (tid < kChunkSize) { + const int row_t = chunk * kChunkSize + tid; + if (row_t < seq_len) { + const int row_token = seq * seq_len + row_t; + g[row_token * v_heads + hv] = scan[tid]; + } + } + + if (hv == 0 && tid == 0) { + const int chunk_idx = seq * chunks_per_sequence + chunk; + chunk_indices[chunk_idx * 2] = seq; + chunk_indices[chunk_idx * 2 + 1] = chunk; + if (chunk == 0) { + cu_seqlens[seq] = seq * seq_len; + } + if (chunk == chunks_per_sequence - 1) { + cu_seqlens[seq + 1] = (seq + 1) * seq_len; + } + } +} + +// The Qwen prefill path always needs both normalization and scalar-gate +// preprocessing. Running the two small kernels back-to-back leaves roughly +// 2--3 us of avoidable serialization at short prefill lengths. This fused +// launcher assigns four warps to four Q/K vectors for the first block range, +// then reuses the same 128-thread block shape for the gate scan blocks. The +// two branches are block-uniform, so the gate barriers never involve norm +// threads from another branch and the prefix-sum semantics are unchanged. +#ifndef CULA_QWEN35_FAST_GATE_SCAN +#define CULA_QWEN35_FAST_GATE_SCAN 1 +#endif +#ifndef CULA_QWEN35_PREPROCESS_THREADS +#define CULA_QWEN35_PREPROCESS_THREADS 256 +#endif +#ifndef CULA_QWEN35_PREPROCESS_GATE_FIRST +#define CULA_QWEN35_PREPROCESS_GATE_FIRST 1 +#endif +static_assert(CULA_QWEN35_PREPROCESS_THREADS >= 64, + "Qwen35 preprocess requires at least two warps"); +static_assert(CULA_QWEN35_PREPROCESS_THREADS % 32 == 0, + "Qwen35 preprocess threads must be a multiple of 32"); +template +__global__ void qwen35_chunk_preprocess_fused_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ a, + const __nv_bfloat16* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ gate_raw, + const float* __restrict__ beta_in, + __nv_bfloat16* __restrict__ q_norm, + __nv_bfloat16* __restrict__ k_norm, + float* __restrict__ g, + float* __restrict__ g_raw_output, + float* __restrict__ beta, + int32_t* __restrict__ unsafe_gate_flags, + int32_t* __restrict__ cu_seqlens, + int32_t* __restrict__ chunk_indices, + int batch_size, + int seq_len, + int qk_heads, + int v_heads, + int vector_count, + int gate_blocks, + int norm_blocks, + int chunks_per_sequence) { + const int block = static_cast(blockIdx.x); +#if CULA_QWEN35_PREPROCESS_GATE_FIRST + if (block >= gate_blocks) { + const int norm_block = block - gate_blocks; +#else + if (block < norm_blocks) { + const int norm_block = block; +#endif + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int vector_idx = norm_block * (static_cast(blockDim.x) / 32) + warp; + if (vector_idx >= vector_count) { + return; + } + const int base = vector_idx * kHeadDimQK; + float q_values[4]; + float k_values[4]; + float q_sq = 0.0f; + float k_sq = 0.0f; +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_values[item] = qwen35_bf16_to_float(q[base + kk]); + k_values[item] = qwen35_bf16_to_float(k[base + kk]); + q_sq += q_values[item] * q_values[item]; + k_sq += k_values[item] * k_values[item]; + } + const float q_rnorm = rsqrtf(qwen35_warp_sum(q_sq) + 1.0e-6f); + const float k_rnorm = rsqrtf(qwen35_warp_sum(k_sq) + 1.0e-6f); +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_norm[base + kk] = qwen35_float_to_bf16(q_values[item] * q_rnorm); + k_norm[base + kk] = qwen35_float_to_bf16(k_values[item] * k_rnorm); + } + return; + } + + __shared__ float scan[kChunkSize + 2]; + const int tid = static_cast(threadIdx.x); +#if CULA_QWEN35_PREPROCESS_GATE_FIRST + int work = block; +#else + int work = block - norm_blocks; +#endif + const int hv = work % v_heads; + work /= v_heads; + const int chunk = work % chunks_per_sequence; + const int seq = work / chunks_per_sequence; + const int local_t = chunk * kChunkSize + tid; + const bool valid = tid < kChunkSize && local_t < seq_len; + const int token = seq * seq_len + local_t; + + if (tid < kChunkSize) { + float log2_decay = 0.0f; + if (valid) { + const int gate_offset = token * v_heads + hv; + float raw_log_decay; + if constexpr (UsePrecomputedGate) { + raw_log_decay = gate_raw[gate_offset]; + beta[gate_offset] = beta_in[gate_offset]; + } else { + const float raw_a = qwen35_bf16_to_float(a[gate_offset]); + const float raw_b = qwen35_bf16_to_float(b[gate_offset]); + raw_log_decay = -expf(A_log[hv]) * qwen35_softplus(raw_a + dt_bias[hv]); + beta[gate_offset] = 1.0f / (1.0f + expf(-raw_b)); + } + log2_decay = raw_log_decay * 1.4426950408889634f; + if (g_raw_output != nullptr) { + g_raw_output[gate_offset] = raw_log_decay; + } + if (unsafe_gate_flags != nullptr && + (!isfinite(raw_log_decay) || raw_log_decay < -5.0f || raw_log_decay > 0.0f)) { + atomicExch(&unsafe_gate_flags[seq * v_heads + hv], 1); + } + } + scan[tid] = log2_decay; + } + __syncthreads(); + +#if CULA_QWEN35_FAST_GATE_SCAN + if (tid < kChunkSize) { + const int lane = tid & 31; + const int warp = tid >> 5; + float prefix = scan[tid]; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const float addend = __shfl_up_sync(0xffffffffu, prefix, offset); + if (lane >= offset) { + prefix += addend; + } + } + if (lane == 31) { + scan[kChunkSize + warp] = prefix; + } + scan[tid] = prefix; + } + __syncthreads(); + if (tid >= 32 && tid < kChunkSize) { + scan[tid] += scan[kChunkSize]; + } + __syncthreads(); +#else +#pragma unroll + for (int offset = 1; offset < kChunkSize; offset <<= 1) { + float addend = 0.0f; + if (tid < kChunkSize && tid >= offset) { + addend = scan[tid - offset]; + } + __syncthreads(); + if (tid < kChunkSize) { + scan[tid] += addend; + } + __syncthreads(); + } +#endif + + if (tid < kChunkSize) { + const int row_t = chunk * kChunkSize + tid; + if (row_t < seq_len) { + const int row_token = seq * seq_len + row_t; + g[row_token * v_heads + hv] = scan[tid]; + } + } + + if (hv == 0 && tid == 0) { + const int chunk_idx = seq * chunks_per_sequence + chunk; + chunk_indices[chunk_idx * 2] = seq; + chunk_indices[chunk_idx * 2 + 1] = chunk; + if (chunk == 0) { + cu_seqlens[seq] = seq * seq_len; + } + if (chunk == chunks_per_sequence - 1) { + cu_seqlens[seq + 1] = (seq + 1) * seq_len; + } + } +} + +template +__global__ __launch_bounds__(kChunkStateWarps * 32, 1) void qwen35_chunk_state_output_kernel( + const __nv_bfloat16* __restrict__ q_norm, + const float* __restrict__ g, + const __nv_bfloat16* __restrict__ Aqk, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ u, + const __nv_bfloat16* __restrict__ kg, + const float* __restrict__ initial_state, + __nv_bfloat16* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + using namespace nvcuda; + extern __shared__ char shared_bytes[]; + auto& shared = *reinterpret_cast(shared_bytes); + const int tid = static_cast(threadIdx.x); + const int warp = tid / 32; + int work = static_cast(blockIdx.x); + const int value_tile = work % (kHeadDimV / kValueTile); + work /= (kHeadDimV / kValueTile); + const int hv = work % kLocalVHeads; + const int seq = work / kLocalVHeads; + if (seq >= batch_size) { + return; + } + const int qk_h = hv / (kLocalVHeads / qk_heads); + const int v_base = value_tile * kValueTile; + const int state_global_base = (seq * kLocalVHeads + hv) * kHeadDimQK * kHeadDimV; + + for (int index = tid; index < kHeadDimQK * kValueTile; index += static_cast(blockDim.x)) { + const int kk = index / kValueTile; + const int vv = index % kValueTile; + shared.state[index] = qwen35_float_to_bf16( + has_initial_state ? initial_state[state_global_base + kk * kHeadDimV + v_base + vv] : 0.0f); + } + __syncthreads(); + + const int chunk_count = (seq_len + kChunkSize - 1) / kChunkSize; + const float q_scale = rsqrtf(static_cast(kHeadDimQK)); + for (int chunk = 0; chunk < chunk_count; ++chunk) { + const int chunk_start = chunk * kChunkSize; + const int valid_rows = min(kChunkSize, seq_len - chunk_start); + + for (int row = tid; row < kChunkSize; row += static_cast(blockDim.x)) { + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.gate_exp[row] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[row] = 0.0f; + } + } + __syncthreads(); + for (int index = tid; index < kChunkSize * kHeadDimQK; index += static_cast(blockDim.x)) { + const int row = index / kHeadDimQK; + const int kk = index % kHeadDimQK; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.matrix[index] = w[(token * kLocalVHeads + hv) * kHeadDimQK + kk]; + } else { + shared.matrix[index] = qwen35_float_to_bf16(0.0f); + } + } + __syncthreads(); + + for (int tile = warp; tile < (kChunkSize / 16) * kValueTilesPerBlock; tile += kChunkOutputWarps) { + const int tile_m = tile / kValueTilesPerBlock; + const int tile_n = tile % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment frag_c; + wmma::fill_fragment(frag_c, 0.0f); +#pragma unroll + for (int kk = 0; kk < kHeadDimQK; kk += 16) { + wmma::load_matrix_sync(frag_a, shared.matrix + tile_m * 16 * kHeadDimQK + kk, kHeadDimQK); + wmma::load_matrix_sync(frag_b, shared.state + kk * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(frag_c, frag_a, frag_b, frag_c); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + frag_c, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + + for (int index = tid; index < kChunkSize * kValueTile; index += static_cast(blockDim.x)) { + const int row = index / kValueTile; + const int vv = index % kValueTile; + float value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = qwen35_bf16_to_float(u[(token * kLocalVHeads + hv) * kHeadDimV + v_base + vv]) - + shared.accum[index]; + } + shared.v_new[index] = qwen35_float_to_bf16(value); + } + for (int index = tid; index < kChunkSize * kHeadDimQK; index += static_cast(blockDim.x)) { + const int row = index / kHeadDimQK; + const int kk = index % kHeadDimQK; + float value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = qwen35_bf16_to_float(q_norm[(token * qk_heads + qk_h) * kHeadDimQK + kk]) * + shared.gate_exp[row] * q_scale; + } + shared.matrix[index] = qwen35_float_to_bf16(value); + } + __syncthreads(); + + for (int tile = warp; tile < (kChunkSize / 16) * kValueTilesPerBlock; tile += kChunkOutputWarps) { + const int tile_m = tile / kValueTilesPerBlock; + const int tile_n = tile % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment frag_c; + wmma::fill_fragment(frag_c, 0.0f); +#pragma unroll + for (int kk = 0; kk < kHeadDimQK; kk += 16) { + wmma::load_matrix_sync(frag_a, shared.matrix + tile_m * 16 * kHeadDimQK + kk, kHeadDimQK); + wmma::load_matrix_sync(frag_b, shared.state + kk * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(frag_c, frag_a, frag_b, frag_c); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + frag_c, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + + for (int index = tid; index < kChunkSize * kChunkSize; index += static_cast(blockDim.x)) { + const int row = index / kChunkSize; + const int col = index % kChunkSize; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.aqk[index] = Aqk[(token * kLocalVHeads + hv) * kChunkSize + col]; + } else { + shared.aqk[index] = qwen35_float_to_bf16(0.0f); + } + } + __syncthreads(); + + if (warp < kChunkOutputWarps) { + const int tile_m = warp / kValueTilesPerBlock; + const int tile_n = warp % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment output_acc; + wmma::load_matrix_sync( + output_acc, + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + kValueTile, + wmma::mem_row_major); +#pragma unroll + for (int kk = 0; kk < kChunkSize; kk += 16) { + wmma::load_matrix_sync(frag_a, shared.aqk + tile_m * 16 * kChunkSize + kk, kChunkSize); + wmma::load_matrix_sync(frag_b, shared.v_new + kk * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(output_acc, frag_a, frag_b, output_acc); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + output_acc, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + + for (int index = tid; index < valid_rows * kValueTile; index += static_cast(blockDim.x)) { + const int row = index / kValueTile; + const int vv = index % kValueTile; + const int token = seq * seq_len + chunk_start + row; + out[(token * kLocalVHeads + hv) * kHeadDimV + v_base + vv] = + qwen35_float_to_bf16(shared.accum[index]); + } + for (int index = tid; index < kChunkSize * kHeadDimQK; index += static_cast(blockDim.x)) { + const int row = index / kHeadDimQK; + const int kk = index % kHeadDimQK; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.matrix[index] = kg[(token * kLocalVHeads + hv) * kHeadDimQK + kk]; + } else { + shared.matrix[index] = qwen35_float_to_bf16(0.0f); + } + } + const int last_token = seq * seq_len + chunk_start + valid_rows - 1; + const float chunk_decay = exp2f(g[last_token * kLocalVHeads + hv]); + __syncthreads(); + + for (int tile = warp; tile < (kHeadDimQK / 16) * kValueTilesPerBlock; tile += kChunkStateWarps) { + const int tile_m = tile / kValueTilesPerBlock; + const int tile_n = tile % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment frag_c; + wmma::fill_fragment(frag_c, 0.0f); +#pragma unroll + for (int tt = 0; tt < kChunkSize; tt += 16) { + wmma::load_matrix_sync(frag_a, shared.matrix + tt * kHeadDimQK + tile_m * 16, kHeadDimQK); + wmma::load_matrix_sync(frag_b, shared.v_new + tt * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(frag_c, frag_a, frag_b, frag_c); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + frag_c, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + for (int index = tid; index < kHeadDimQK * kValueTile; index += static_cast(blockDim.x)) { + const float updated = shared.accum[index] + chunk_decay * qwen35_bf16_to_float(shared.state[index]); + shared.state[index] = qwen35_float_to_bf16(updated); + } + __syncthreads(); + } + + for (int index = tid; index < kHeadDimQK * kValueTile; index += static_cast(blockDim.x)) { + const int kk = index / kValueTile; + const int vv = index % kValueTile; + final_state[state_global_base + kk * kHeadDimV + v_base + vv] = qwen35_bf16_to_float(shared.state[index]); + } +} + +inline void launch_qwen35_chunk_preprocess( + cudaStream_t stream, + const __nv_bfloat16* q, + const __nv_bfloat16* k, + const __nv_bfloat16* a, + const __nv_bfloat16* b, + const float* A_log, + const float* dt_bias, + const float* gate_raw, + const float* beta_in, + bool use_precomputed_gate, + __nv_bfloat16* q_norm, + __nv_bfloat16* k_norm, + float* g, + float* g_raw_output, + float* beta, + int32_t* unsafe_gate_flags, + int32_t* cu_seqlens, + int32_t* chunk_indices, + int batch_size, + int seq_len, + int qk_heads, + int v_heads) { + const int vector_count = batch_size * seq_len * qk_heads; + const int chunks = (seq_len + kChunkSize - 1) / kChunkSize; + const int gate_blocks = batch_size * chunks * v_heads; + constexpr int kNormVectorsPerBlock = CULA_QWEN35_PREPROCESS_THREADS / 32; + const int norm_blocks = (vector_count + kNormVectorsPerBlock - 1) / kNormVectorsPerBlock; + if (use_precomputed_gate) { + qwen35_chunk_preprocess_fused_kernel<<< + norm_blocks + gate_blocks, CULA_QWEN35_PREPROCESS_THREADS, 0, stream>>>( + q, k, a, b, A_log, dt_bias, gate_raw, beta_in, + q_norm, k_norm, g, g_raw_output, beta, unsafe_gate_flags, + cu_seqlens, chunk_indices, batch_size, seq_len, qk_heads, v_heads, + vector_count, gate_blocks, norm_blocks, chunks); + } else { + qwen35_chunk_preprocess_fused_kernel<<< + norm_blocks + gate_blocks, CULA_QWEN35_PREPROCESS_THREADS, 0, stream>>>( + q, k, a, b, A_log, dt_bias, gate_raw, beta_in, + q_norm, k_norm, g, g_raw_output, beta, unsafe_gate_flags, + cu_seqlens, chunk_indices, batch_size, seq_len, qk_heads, v_heads, + vector_count, gate_blocks, norm_blocks, chunks); + } +} + +template +inline void launch_qwen35_chunk_state_output( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + auto kernel_fn = &qwen35_chunk_state_output_kernel; + constexpr size_t shared_bytes = sizeof(Qwen35ChunkStateOutputShared); + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes); + const int grid = batch_size * kLocalVHeads * (kHeadDimV / kValueTile); + kernel_fn<<>>( + q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, + batch_size, seq_len, qk_heads, has_initial_state); +} + + +template +__global__ void qwen35_scalar_kda_prefill_kernel( + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ initial_state, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state, + const float* __restrict__ precomputed_g, + const float* __restrict__ precomputed_beta, + const int32_t* __restrict__ unsafe_gate_flags) { + __shared__ typename Qwen35ScalarKdaPrefillKernel::SharedStorage storage; + Qwen35ScalarKdaPrefillKernel::run_device( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + batch_size, + seq_len, + qk_heads, + sequence_count, + is_varlen, + has_initial_state, + precomputed_g, + precomputed_beta, + unsafe_gate_flags, + storage); +} + +template +void launch_qwen35_scalar_kda_prefill_kernel_variant( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + using Kernel = Qwen35ScalarKdaPrefillKernel; + const auto grid = Kernel::grid_shape(sequence_count); + const auto block = Kernel::block_shape(); + qwen35_scalar_kda_prefill_kernel<<>>( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + batch_size, + seq_len, + qk_heads, + sequence_count, + is_varlen, + has_initial_state, + nullptr, + nullptr, + nullptr); +} + +template +void launch_qwen35_scalar_kda_prefill_precomputed_fallback( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const float* g, + const float* beta, + const float* initial_state, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state, + const int32_t* unsafe_gate_flags) { + // Safe inputs make this launch an early-return guard. A 16-warp CTA keeps + // that fixed cost to only 8 CTAs per V head, while still providing an exact + // recurrent overwrite for the rare unsafe head. + constexpr int kFallbackWarps = 16; + using Kernel = Qwen35ScalarKdaPrefillKernel; + qwen35_scalar_kda_prefill_kernel + <<>>( + q, + k, + v, + nullptr, + nullptr, + nullptr, + nullptr, + initial_state, + nullptr, + out, + final_state, + batch_size, + seq_len, + qk_heads, + batch_size, + false, + has_initial_state, + g, + beta, + unsafe_gate_flags); +} + +template +void launch_qwen35_scalar_kda_prefill_kernel( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + launch_qwen35_scalar_kda_prefill_kernel_variant( + stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, + batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); +} + +} // namespace cula::qwen35::prefill::kernel diff --git a/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu new file mode 100644 index 00000000..88b157b4 --- /dev/null +++ b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu @@ -0,0 +1,180 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#include "qwen35_chunk_prefill_traits_sm90.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cula::qwen35::prefill::sm90 { + +namespace { + +using DefaultTraits = Qwen35ChunkPrefillSm90DefaultTraits; + +static_assert(DefaultTraits::kBlockT == 64); +static_assert(DefaultTraits::kBlockV == 64); +static_assert(DefaultTraits::kStages == 2); +static_assert(size(typename DefaultTraits::TiledMmaQK{}) == 128); +static_assert(size(typename DefaultTraits::TiledMmaOV{}) == 128); +static_assert(cosize(typename DefaultTraits::SmemLayoutQ{}) > 0); +static_assert(cosize(typename DefaultTraits::SmemLayoutK{}) > 0); + +void check_cutlass_status(cutlass::Status status, const char* what) { + TORCH_CHECK(status == cutlass::Status::kSuccess, what, " failed with CUTLASS status ", static_cast(status)); +} + +template +void run_qwen35_chunk_qk_prefill_sm90_impl(const at::Tensor& q, const at::Tensor& k, const at::Tensor& out) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = float; + using ElementD = float; + using ElementAccumulator = float; + using ElementCompute = float; + + using LayoutA = cute::tuple; + using LayoutB = cute::tuple; + using LayoutC = cute::tuple; + using LayoutD = LayoutC; + + constexpr int kAlignmentA = 16 / sizeof(ElementA); + constexpr int kAlignmentB = 16 / sizeof(ElementB); + constexpr int kAlignmentC = 16 / sizeof(ElementC); + constexpr int kAlignmentD = 16 / sizeof(ElementD); + + using OperatorClass = cutlass::arch::OpClassTensorOp; + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) + using ArchTag = cutlass::arch::Sm100; + using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; + using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; +#else + using ArchTag = cutlass::arch::Sm90; + using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized; + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized; +#endif +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) + using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; +#else + using EpilogueTileType = decltype(cute::take<0, 2>(TileShape{})); +#endif + using FusionOperation = + typename cutlass::epilogue::fusion::LinearCombination; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + TileShape, + ClusterShape, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC, + LayoutC, + kAlignmentC, + ElementD, + LayoutD, + kAlignmentD, + EpilogueSchedule, + FusionOperation>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + TileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t HV = q.size(2); + constexpr int K = kHeadDimQK; + const int64_t L = B * HV; + + LayoutA stride_A{HV * K, cute::_1{}, K}; + LayoutB stride_B{HV * K, cute::_1{}, K}; + LayoutC stride_C{T, cute::_1{}, T * T}; + + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {static_cast(T), static_cast(T), K, static_cast(L)}, + { + reinterpret_cast(q.data_ptr()), + stride_A, + reinterpret_cast(k.data_ptr()), + stride_B, + }, + { + {1.0f, 0.0f}, + out.data_ptr(), + stride_C, + out.data_ptr(), + stride_C, + }, + }; + + Gemm gemm; + const size_t workspace_size = Gemm::get_workspace_size(arguments); + at::Tensor workspace = at::empty({static_cast(workspace_size)}, q.options().dtype(at::kByte)); + check_cutlass_status(gemm.can_implement(arguments), "qwen35_chunk_qk_prefill_sm90 can_implement"); + check_cutlass_status(gemm.initialize(arguments, workspace.data_ptr(), at::cuda::getCurrentCUDAStream(q.device().index())), "qwen35_chunk_qk_prefill_sm90 initialize"); + check_cutlass_status(gemm.run(at::cuda::getCurrentCUDAStream(q.device().index())), "qwen35_chunk_qk_prefill_sm90 run"); +} + +} // namespace + +void qwen35_chunk_qk_prefill_sm90(at::Tensor q, at::Tensor k, at::Tensor out) { + TORCH_CHECK(q.is_cuda(), "q must be CUDA"); + TORCH_CHECK(k.is_cuda(), "k must be CUDA"); + TORCH_CHECK(out.is_cuda(), "out must be CUDA"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16, "q must be bfloat16"); + TORCH_CHECK(k.scalar_type() == at::kBFloat16, "k must be bfloat16"); + TORCH_CHECK(out.scalar_type() == at::kFloat, "out must be float32"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B,T,HV,128]"); + TORCH_CHECK(k.is_contiguous(), "k must be contiguous [B,T,HV,128]"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B,HV,T,T]"); + TORCH_CHECK(q.dim() == 4, "q must be [B,T,HV,128]"); + TORCH_CHECK(k.sizes() == q.sizes(), "k must match q"); + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t HV = q.size(2); + TORCH_CHECK(decode::is_supported_local_v_heads(static_cast(HV)), "expected local HV in {48, 24, 12, 6}, got ", HV); + TORCH_CHECK(q.size(3) == kHeadDimQK, "expected D=128"); + TORCH_CHECK(out.sizes() == at::IntArrayRef({B, HV, T, T}), "out must be [B,HV,T,T]"); + + const at::cuda::OptionalCUDAGuard device_guard(q.device()); + run_qwen35_chunk_qk_prefill_sm90_impl(q, k, out); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::prefill::sm90 diff --git a/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_traits_sm90.hpp b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_traits_sm90.hpp new file mode 100644 index 00000000..c39e926e --- /dev/null +++ b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_traits_sm90.hpp @@ -0,0 +1,112 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// 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. + +#pragma once + +#include "qwen35/prefill/qwen35_prefill_common.cuh" + +#include +#include +#include +#include +#include +#include + +namespace cula::qwen35::prefill::sm90 { + +using namespace cute; + +// First SM90 chunk shape for Qwen3.5 prefill. +// +// This intentionally only describes the TMA/WGMMA tiles. The full chunk +// algorithm still needs a local chunk recurrence and inter-chunk state scan; +// those should be built on top of these traits instead of extending the scalar +// fallback kernel. +template +struct Qwen35ChunkPrefillSm90Traits { + static constexpr int kBlockT = kBlockT_; + static constexpr int kBlockV = kBlockV_; + static constexpr int kStages = kStages_; + + static_assert(kBlockT == 64 || kBlockT == 128, "GMMA chunk tiles expect BT=64 or BT=128."); + static_assert(kBlockV == 64 || kBlockV == 128, "V chunk tiles expect BV=64 or BV=128."); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + + using Element = cutlass::bfloat16_t; + using Accumulator = float; + static constexpr int kAlignment = 16 / sizeof(Element); + + using ClusterShape = Shape<_1, _1, _1>; + using StageCount = cutlass::gemm::collective::StageCount; + + // q/k/v are materialized by qwen35_layout_prefill as contiguous + // [total_tokens, 48, 128]. The TMA tensor view below exposes them as + // (token, dim, head), with dynamic strides: + // token stride = 48 * 128 + // dim stride = 1 + // head stride = 128 + using GmemStrideTDH = cute::tuple; + + using TileShapeQK = decltype(make_shape(Int{}, Int{}, Int{})); + using TileShapeOV = decltype(make_shape(Int{}, Int{}, Int{})); + + // Q @ K^T => [BT, BT]. CollectiveBuilder selects GMMA and TMA-compatible + // shared-memory layouts for SM90. + using CollectiveQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + GmemStrideTDH, + kAlignment, + Element, + GmemStrideTDH, + kAlignment, + Accumulator, + TileShapeQK, + ClusterShape, + StageCount, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaQK = typename CollectiveQK::TiledMma; + using SmemLayoutQ = typename CollectiveQK::SmemLayoutA; + using SmemLayoutK = typename CollectiveQK::SmemLayoutB; + using TmaQ = typename CollectiveQK::Params::TMA_A; + using TmaK = typename CollectiveQK::Params::TMA_B; + + // Q @ state / local_value => [BT, BV]. This is the second core WGMMA shape + // needed once chunk-local state summaries are available. + using CollectiveOV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + GmemStrideTDH, + kAlignment, + Element, + GmemStrideTDH, + kAlignment, + Accumulator, + TileShapeOV, + ClusterShape, + StageCount, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaOV = typename CollectiveOV::TiledMma; + using SmemLayoutOV_A = typename CollectiveOV::SmemLayoutA; + using SmemLayoutOV_B = typename CollectiveOV::SmemLayoutB; +}; + +using Qwen35ChunkPrefillSm90DefaultTraits = Qwen35ChunkPrefillSm90Traits<64, 64, 2>; + +} // namespace cula::qwen35::prefill::sm90 diff --git a/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py b/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py index 8e5ad3d4..5912b28d 100644 --- a/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py +++ b/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py @@ -56,6 +56,7 @@ import argparse import time +from types import SimpleNamespace import cuda.bindings.driver as cuda import cutlass @@ -68,7 +69,20 @@ from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Int32, Int64 -from fla.modules.l2norm import l2norm_fwd + +# CUTLASS DSL 4.3+ changed fence_proxy() from exported enum arguments to +# string literals. Keep the existing call sites compatible with both APIs. +if not hasattr(cute.arch, "ProxyKind"): + cute.arch.ProxyKind = SimpleNamespace(async_shared="async.shared") +if not hasattr(cute.arch, "SharedSpace"): + cute.arch.SharedSpace = SimpleNamespace(shared_cta="cta") + +try: + from fla.modules.l2norm import l2norm_fwd +except ImportError: + def l2norm_fwd(x: torch.Tensor): + rstd = torch.rsqrt(x.float().square().sum(dim=-1, keepdim=True).clamp_min(1.0e-12)) + return (x.float() * rstd).to(x.dtype), rstd from cula.utils import assert_blackwell @@ -90,6 +104,7 @@ class Constant: D = 128 # head dim HALF_D = 64 # half head dim for partitioned S2R SCALE = float(D) ** -0.5 + RCP_LN2 = 1.4426950408889634 BK_SC = 64 # tile size in subchunk MMA @@ -126,19 +141,32 @@ def __init__( io_dtype: type[cutlass.Numeric] = cutlass.BFloat16, scale: cutlass.Float32 = 1.0, safe_gate: bool = False, + scalar_gate: bool = False, + fuse_scalar_cumsum: bool = False, + split_value_tiles: bool = False, has_initial_state: bool = False, output_final_state: bool = False, is_varlen: bool = False, use_fast_math: bool = True, - # num_regs_cuda: int = 248, num_regs_cuda: int = 224, num_regs_subchunk: int = 192, - num_regs_others: int = 64, # Optimized: best config from comprehensive sweep + num_regs_others: int = 80, ): assert_blackwell() # make scale a constant self.scale = scale self.safe_gate = safe_gate + self.scalar_gate = scalar_gate + self.fuse_scalar_cumsum = fuse_scalar_cumsum + self.split_value_tiles = split_value_tiles + if scalar_gate and not safe_gate: + raise ValueError("native scalar gate currently requires safe_gate=True") + if fuse_scalar_cumsum and not scalar_gate: + raise ValueError("fuse_scalar_cumsum requires scalar_gate=True") + if split_value_tiles and (not scalar_gate or not safe_gate or is_varlen): + raise ValueError( + "split_value_tiles currently requires non-varlen safe scalar gate" + ) self.has_initial_state = has_initial_state self.output_final_state = output_final_state self.is_varlen = is_varlen @@ -169,6 +197,9 @@ def __init__( # K: (64, 128) # V: (64, 128) C, D = (Constant.C, Constant.D) + self.value_tile_size = D // 2 if split_value_tiles else D + self.num_value_tiles = D // self.value_tile_size + DV_TILE = self.value_tile_size HALF_D = Constant.HALF_D # (C, C, D) self.qk_mma_tiler = (C, C, D) # (M, N, K) @@ -176,15 +207,15 @@ def __init__( self.qk_mma_tiler_half = (C, C, HALF_D) # (M, N, K/2) self.kk_mma_tiler = (C, C, D) # (M, N, K) # (D, C, C) - self.vp_mma_tiler = (D, C, C) # (M, N, K) - self.mv_mma_tiler = (D, C, C) # (M, N, K) + self.vp_mma_tiler = (DV_TILE, C, C) # (M, N, K) + self.mv_mma_tiler = (DV_TILE, C, C) # (M, N, K) # (D, D, C) - self.kv_mma_tiler = (D, D, C) # (M, N, K) + self.kv_mma_tiler = (DV_TILE, D, C) # (M, N, K) # (D, C, D) # State as operand A since it's in TMEM # Q now as operand B - self.sq_mma_tiler = (D, C, D) # (M, N, K) - self.ks_mma_tiler = (D, C, D) # (M, N, K) + self.sq_mma_tiler = (DV_TILE, C, D) # (M, N, K) + self.ks_mma_tiler = (DV_TILE, C, D) # (M, N, K) # subchunk MMA SC, BK_SC = (Constant.SC, Constant.BK_SC) @@ -326,8 +357,8 @@ def _setup_attributes(self): self.k_stage = 2 self.v_stage = 1 self.o_stage = 2 - self.g_stage = 2 # Single stage for g (CUDA warp processes immediately) - self.beta_stage = 1 # TODO: two stage ? + self.g_stage = 2 + self.beta_stage = 2 self.q_k_scaled_stage = 1 # only single stage here due to smem limitation self.epi_stage = 2 @@ -348,7 +379,7 @@ def _compute_grid( # cute.ceil_div(o_shape[0], chunk_size), # For Loop to tile over chunk size, # TODO: varlen will make parallelism good enough - 1, + self.num_value_tiles, # H cute.size(o_shape[2][0]), # B @@ -368,7 +399,7 @@ def __call__( final_state_iter: cute.Pointer, # Final state [B, H, D, D], float32 or nullptr cu_seqlens_iter: cute.Pointer, # Cumulative seq lengths [num_seqs+1], int32 (varlen) workspace_iter: cute.Pointer, # Workspace buffer for TMA descriptor modification - problem_size: tuple[Int32, Int32, Int32, Int32], # (B/num_seqs, S/total_tokens, H, D) + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], # (B/num_seqs, S/total_tokens, H, HV, D) stream: cuda.CUstream, options=None, # compile options ): @@ -388,11 +419,11 @@ def __call__( final_state_iter: Final state [N, H, D, D] or nullptr cu_seqlens_iter: Cumulative seq lengths [num_seqs+1], int32 (varlen only) workspace_iter: Workspace buffer for TMA descriptor modification (varlen tail tiles) - problem_size: (N, S, H, D) where N=B or num_seqs, S=seq_len or total_tokens + problem_size: (N, S, H, HV, D) where N=B or num_seqs, S=seq_len or total_tokens stream: CUDA stream options: compile options for the kernel """ - B, S, H, D = problem_size + B, S, H, HV, D = problem_size # Setup attributes self._setup_attributes() @@ -421,43 +452,46 @@ def __call__( stride=(D * H, 1, (D, D * H * S)), ) k = cute.make_tensor(k_iter, k_layout) - kt_layout = cute.make_layout( - (D, S, (H, data_B)), - stride=(1, D * H, (D, D * H * S)), - ) - kt = cute.make_tensor(k_iter, kt_layout) # v v_layout = cute.make_layout( - (D, S, (H, data_B)), - stride=(1, D * H, (D, D * H * S)), + (D, S, (HV, data_B)), + stride=(1, D * HV, (D, D * HV * S)), ) v = cute.make_tensor(v_iter, v_layout) - # g (gate) - NEW for KDA, same layout as Q/K - g_layout = cute.make_layout( - (S, D, (H, data_B)), - stride=(D * H, 1, (D, D * H * S)), - ) + # GDN uses one scalar gate per token/value-head. Keep a singleton + # second mode so the scalar specialization can use a regular 2-D TMA + # tile while preserving the generic vector-gate layout unchanged. + if cutlass.const_expr(self.scalar_gate): + g_layout = cute.make_layout( + (S, 1, (HV, data_B)), + stride=(HV, 0, (1, HV * S)), + ) + else: + g_layout = cute.make_layout( + (S, D, (HV, data_B)), + stride=(D * HV, 1, (D, D * HV * S)), + ) g = cute.make_tensor(g_iter, g_layout) # beta - NEW for KDA, shape (B, S, H) or (1, total_tokens, H) for varlen beta_layout = cute.make_layout( - (S, (H, data_B)), - stride=(H, (1, H * S)), + (S, (HV, data_B)), + stride=(HV, (1, HV * S)), ) beta = cute.make_tensor(beta_iter, beta_layout) o_layout = cute.make_layout( - (D, S, (H, data_B)), - stride=(1, D * H, (D, D * H * S)), + (D, S, (HV, data_B)), + stride=(1, D * HV, (D, D * HV * S)), ) o = cute.make_tensor(o_iter, o_layout) # Initial state / final state: [N, H, D, D] stored as row-major # N = B (non-varlen) or num_seqs (varlen). Always uses B from problem_size. fstate_layout = cute.make_layout( - (D, D, (H, B)), - stride=(1, D, (D * D, D * D * H)), + (D, D, (HV, B)), + stride=(1, D, (D * D, D * D * HV)), ) initial_state = cute.make_tensor(initial_state_iter, fstate_layout) final_state = cute.make_tensor(final_state_iter, fstate_layout) @@ -477,7 +511,12 @@ def __call__( self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() - self.g_major_mode = utils.LayoutEnum.from_tensor(g).mma_major_mode() # NEW for KDA + # Scalar G is not an MMA operand. The value is only needed for the + # generic vector-gate TMA construction below. + if cutlass.const_expr(self.scalar_gate): + self.g_major_mode = self.q_major_mode + else: + self.g_major_mode = utils.LayoutEnum.from_tensor(g).mma_major_mode() self.k_major_mode_kv = tcgen05.OperandMajorMode.MN # For V^T*K, S dimension coalesced # TMEM register output results as (D, C) self.o_layout = utils.LayoutEnum.from_tensor(o) @@ -624,16 +663,30 @@ def __call__( self.k_dtype, self.k_stage, ) - # G (gate) - NEW for KDA - # Use same layout as Q since g has same shape and memory layout as Q - # This ensures TMA compatibility - # ((MMA_ATOM_M, MMA_ATOM_K), MMA_M, MMA_K, STAGES) - g_smem_layout_staged = sm100_utils.make_smem_layout_a( - qk_tiled_mma, - self.qk_mma_tiler, - self.g_dtype, - self.g_stage, - ) + # The scalar specialization stages only C values. Its storage is + # subsequently reused as BF16 K*exp(-g), so SharedStorage below keeps + # enough physical bytes for that tensor without retaining the 64 KiB + # FP32 vector-gate allocation. + if cutlass.const_expr(self.scalar_gate): + scalar_g_stage_elements = Constant.C * ( + 4 + Constant.C // Constant.SC + ) + (Constant.C // Constant.SC) ** 2 + g_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), + 0, + cute.make_layout( + (Constant.C, 1, self.g_stage), + stride=(1, Constant.C, scalar_g_stage_elements), + ), + ) + else: + # Generic vector gate: same MMA-compatible layout as Q. + g_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.g_dtype, + self.g_stage, + ) # V^T*P p_smem_layout_staged = sm100_utils.make_smem_layout_b( vp_tiled_mma, @@ -688,15 +741,6 @@ def __call__( qk_tiled_mma, cluster_layout_vmnk.shape, ) - kv_k_smem_layout = cute.select(kv_k_smem_layout_staged, mode=[0, 1, 2]) - tma_atom_kt, tma_tensor_kt = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, - kt, - kv_k_smem_layout, - self.kv_mma_tiler, - kv_tiled_mma, - cluster_layout_vmnk.shape, - ) # TMA load for V v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_A( @@ -707,17 +751,26 @@ def __call__( vp_tiled_mma, cluster_layout_vmnk.shape, ) - # TMA load for G (gate) - NEW for KDA - # Use same TMA atom as Q since g has same layout as Q - g_smem_layout = cute.select(g_smem_layout_staged, mode=[0, 1, 2]) - tma_atom_g, tma_tensor_g = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, - g, - g_smem_layout, - self.qk_mma_tiler, - qk_tiled_mma, - cluster_layout_vmnk.shape, - ) + # TMA load for G. Scalar G uses a (C, 1) epilogue-style tile; vector + # G retains the original MMA-operand descriptor. + if cutlass.const_expr(self.scalar_gate): + g_smem_layout = cute.select(g_smem_layout_staged, mode=[0, 1]) + tma_atom_g, tma_tensor_g = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_op, + g, + g_smem_layout, + (Constant.C, 1), + ) + else: + g_smem_layout = cute.select(g_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_g, tma_tensor_g = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + g, + g_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + cluster_layout_vmnk.shape, + ) # NOTE: G's last row will be extracted from sG in CUDA warp after TMA load # No separate TMA needed for G last row - we extract it from the full G tile @@ -731,16 +784,26 @@ def __call__( q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) - g_copy_size = cute.size_in_bytes(self.g_dtype, g_smem_layout) # NEW for KDA + v_copy_size = cute.size_in_bytes(self.v_dtype, v_smem_layout) + g_copy_size = cute.size_in_bytes(self.g_dtype, g_smem_layout) self.tma_copy_q_bytes = q_copy_size self.tma_copy_k_bytes = k_copy_size - # self.tma_copy_v_bytes = v_copy_size - self.tma_copy_v_bytes = k_copy_size + self.tma_copy_v_bytes = v_copy_size self.tma_copy_g_bytes = g_copy_size # NEW for KDA beta_layout = cute.make_layout((Constant.C, self.beta_stage), stride=(1, Constant.C)) g_last_layout = cute.make_layout((Constant.D, self.g_stage), stride=(1, Constant.D)) + # Per stage: raw G[C], row factors[C], and four sets of column + # factors[4,C]. This is 1.5 KiB/stage, still far below vector G. + if cutlass.const_expr(self.scalar_gate): + scalar_g_stage_elements = Constant.C * ( + 4 + Constant.C // Constant.SC + ) + (Constant.C // Constant.SC) ** 2 + g_storage_elements = scalar_g_stage_elements * self.g_stage + else: + g_storage_elements = cute.cosize(g_smem_layout_staged) + @cute.struct class SharedStorage: # Pipeline barriers @@ -776,8 +839,6 @@ class SharedStorage: ks_mbar_ptr: cute.struct.MemRange[Int64, self.ks_stage * 2] # type: ignore o_inter_mbar_ptr: cute.struct.MemRange[Int64, 1 * 2] # type: ignore smem_o_mbar_ptr: cute.struct.MemRange[Int64, self.acc_stage * 2] # type: ignore - kv_decay_mbar_ptr: cute.struct.MemRange[Int64, 1 * 2] # type: ignore - kv_decay_mbar_ptr: cute.struct.MemRange[Int64, 1 * 2] # type: ignore # Tmem holding buffer tmem_holding_buf: Int32 # Smem tensors @@ -804,7 +865,7 @@ class SharedStorage: ] # G (gate) - NEW for KDA sG: cute.struct.Align[ - cute.struct.MemRange[self.g_dtype, cute.cosize(g_smem_layout_staged)], # type: ignore + cute.struct.MemRange[self.g_dtype, g_storage_elements], # type: ignore self.buffer_align_bytes, ] # Store QK @@ -832,7 +893,7 @@ class SharedStorage: self.shared_storage = SharedStorage if cutlass.const_expr(self.is_varlen): - self.grid = (1, H, B) + self.grid = (1, HV, B) # TensorMapManager for TMA descriptor modification in varlen tail tiles self._tensormap_mgr = utils.TensorMapManager(utils.TensorMapUpdateMode.GMEM, 128) else: @@ -853,12 +914,11 @@ class SharedStorage: tma_tensor_q, tma_atom_k, tma_tensor_k, - tma_atom_kt, - tma_tensor_kt, tma_atom_v, tma_tensor_v, tma_atom_g, # NEW for KDA tma_tensor_g, # NEW for KDA + g, tma_atom_o, tma_tensor_o, beta, # NEW for KDA @@ -901,12 +961,11 @@ def kernel( tma_tensor_q: cute.Tensor, tma_atom_k: cute.CopyAtom, tma_tensor_k: cute.Tensor, - tma_atom_kt: cute.CopyAtom, - tma_tensor_kt: cute.Tensor, tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, tma_atom_g: cute.CopyAtom, # NEW for KDA tma_tensor_g: cute.Tensor, # NEW for KDA + g: cute.Tensor, tma_atom_o: cute.CopyAtom, tma_tensor_o: cute.Tensor, beta: cute.Tensor, # NEW for KDA - shape (S, (H, B)) @@ -926,7 +985,7 @@ def kernel( cu_seqlens: cute.Tensor, # int32 tensor for varlen o_gmem: cute.Tensor, # raw GMEM output tensor (D, S, (H, data_B)) for tail tile handling workspace_iter: cute.Pointer, # workspace buffer for TMA descriptor modification - problem_size: tuple[Int32, Int32, Int32, Int32], # (B, S, H, D) + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], # (B, S, H, HV, D) ): """ KDA Kernel - Step 1: Gate processing @@ -947,7 +1006,8 @@ def kernel( cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_g) # NEW for KDA + if cutlass.const_expr(not self.scalar_gate): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_g) cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_o) # Allocate shared memory @@ -1022,16 +1082,27 @@ def kernel( consumer_group=make_thread_cooperative_group(32 * len(self.cuda_warp_ids)), barrier_storage=storage.end_v_mbar_ptr.data_ptr(), ).make_participants() - # G (gate/g_cumsum) - NEW for KDA - load_g_producer, load_g_consumer = pipeline.PipelineTmaAsync.create( - num_stages=self.g_stage, - producer_group=make_thread_cooperative_group(len([self.load_warp_id])), - consumer_group=make_thread_cooperative_group( - len([*self.cuda_warp_ids, *self.cuda_subchunk_warp_ids]) - ), # CUDA cores will consume - tx_count=self.tma_copy_g_bytes, - barrier_storage=storage.load_g_mbar_ptr.data_ptr(), - ).make_participants() + # Scalar G is copied by the load warp (two values per lane); vector G + # retains its TMA pipeline. + if cutlass.const_expr(self.scalar_gate): + load_g_producer, load_g_consumer = pipeline.PipelineAsync.create( + num_stages=self.g_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len([*self.cuda_warp_ids, *self.cuda_subchunk_warp_ids]) + ), + barrier_storage=storage.load_g_mbar_ptr.data_ptr(), + ).make_participants() + else: + load_g_producer, load_g_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.g_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group( + len([*self.cuda_warp_ids, *self.cuda_subchunk_warp_ids]) + ), + tx_count=self.tma_copy_g_bytes, + barrier_storage=storage.load_g_mbar_ptr.data_ptr(), + ).make_participants() load_beta_producer, load_beta_consumer = pipeline.PipelineAsync.create( num_stages=self.beta_stage, producer_group=make_thread_cooperative_group(self.threads_per_warp * len([self.load_beta_warp_id])), @@ -1139,14 +1210,6 @@ def kernel( consumer_group=make_thread_cooperative_group(self.threads_per_warp * len([self.epilogue_warp_id])), barrier_storage=storage.smem_o_mbar_ptr.data_ptr(), ).make_participants() - # T2R & R2T sync in S decay - kv_decay_producer, kv_decay_consumer = pipeline.PipelineAsyncUmma.create( - num_stages=1, - producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.cuda_warp_ids)), - consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), - barrier_storage=storage.kv_decay_mbar_ptr.data_ptr(), - ).make_participants() - # TMEM tmem_alloc_barrier = pipeline.NamedBarrier( barrier_id=1, @@ -1198,54 +1261,111 @@ def kernel( sK_g = storage.sK.get_tensor(k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner) sK_kv = storage.sQ_K_scaled.get_tensor(kv_k_smem_mma_layout_staged.outer, swizzle=kv_k_smem_mma_layout_staged.inner) sK_ks = storage.sQ_K_scaled.get_tensor(k_smem_mma_layout_staged.outer, swizzle=k_smem_mma_layout_staged.inner) - # NOTE: reuse same smem as sG - sK_neg_g_f32 = storage.sG.get_tensor( - # kv_k_smem_layout_staged.outer, swizzle=kv_k_smem_layout_staged.inner - # NOTE: same swizzle atom (k-major) as k_smem_layout_staged - k_smem_layout_staged.outer, - swizzle=k_smem_layout_staged.inner, - ) - # NOTE: recast as bf16 since operand B is BF16 - # CRITICAL FIX: sK_neg_g's stage stride must match sG's byte stride - # sG (F32) has stage stride = 8192 elements = 32768 bytes - # sK_neg_g (BF16) must have stage stride = 32768 bytes = 16384 BF16 elements - # Original bug: using sK_g.layout which has stage stride = 8192 BF16 elements = 16384 bytes - # This caused sK_neg_g stage 1 to overlap with sG stage 0's second half! - - # Get the base layout from sK_g but double the stage stride - sK_g_outer = sK_g.layout - # Create new layout with corrected stage stride (16384 BF16 elements instead of 8192) - # sK_g layout is: ((64,16),1,(4,2),2):((64,1),0,(16,4096),8192) - # We need: ((64,16),1,(4,2),2):((64,1),0,(16,4096),16384) - sK_neg_g_layout = cute.make_layout( - sK_g_outer.shape, - stride=(*sK_g_outer.stride[:-1], sK_g_outer.stride[-1] * 2), # Double the stage stride - ) - sK_neg_g = cute.make_tensor( - cute.recast_ptr(sK_neg_g_f32.iterator, swizzle_=k_smem_layout_staged.inner, dtype=self.io_dtype), - layout=sK_neg_g_layout, - ) - - # Same fix for sK_neg_g_b - sK_neg_g_b_outer = kv_k_smem_layout_staged.outer - sK_neg_g_b_layout = cute.make_layout( - sK_neg_g_b_outer.shape, - stride=(*sK_neg_g_b_outer.stride[:-1], sK_neg_g_b_outer.stride[-1] * 2), # Double the stage stride - ) + # Gate staging view and the BF16 K*exp(-g) view which reuses the same + # physical buffer after all gate consumers finish. + if cutlass.const_expr(self.scalar_gate): + sG_tma = storage.sG.get_tensor(g_smem_layout_staged.outer, swizzle=g_smem_layout_staged.inner) + scalar_g_stage_elements = Constant.C * ( + 4 + Constant.C // Constant.SC + ) + (Constant.C // Constant.SC) ** 2 + sG = cute.make_tensor( + sG_tma.iterator, + layout=cute.make_layout( + (Constant.C, Constant.D, self.g_stage), + stride=(1, 0, scalar_g_stage_elements), + ), + ) + sG_factor_a = cute.make_tensor( + sG_tma.iterator + Constant.C, + layout=cute.make_layout( + (Constant.C, self.g_stage), + stride=(1, scalar_g_stage_elements), + ), + ) + sG_factor_b = cute.make_tensor( + sG_tma.iterator + 2 * Constant.C, + layout=cute.make_layout( + (Constant.C // Constant.SC, Constant.C, self.g_stage), + stride=(Constant.C, 1, scalar_g_stage_elements), + ), + ) + sG_factor_base = cute.make_tensor( + sG_tma.iterator + + Constant.C * (2 + Constant.C // Constant.SC), + layout=cute.make_layout( + ( + Constant.C // Constant.SC, + Constant.C // Constant.SC, + self.g_stage, + ), + stride=( + Constant.C // Constant.SC, + 1, + scalar_g_stage_elements, + ), + ), + ) + scalar_base_end = ( + Constant.C * (2 + Constant.C // Constant.SC) + + (Constant.C // Constant.SC) ** 2 + ) + sG_main_q = cute.make_tensor( + sG_tma.iterator + scalar_base_end, + layout=cute.make_layout( + (Constant.C, self.g_stage), + stride=(1, scalar_g_stage_elements), + ), + ) + sG_main_k = cute.make_tensor( + sG_tma.iterator + scalar_base_end + Constant.C, + layout=cute.make_layout( + (Constant.C, self.g_stage), + stride=(1, scalar_g_stage_elements), + ), + ) + # The BF16 scratch views below are compile-time dead in the + # supported safe-gate scalar specialization. + sK_neg_g_layout = sK_g.layout + sK_neg_g_b_layout = kv_k_smem_layout_staged.outer + sK_neg_g_ptr = cute.recast_ptr( + sG_tma.iterator, + swizzle_=k_smem_layout_staged.inner, + dtype=self.io_dtype, + ) + else: + sG_tma = storage.sG.get_tensor(g_smem_layout_staged.outer, swizzle=g_smem_layout_staged.inner) + sG = sG_tma + # Compile-time unused by the generic vector-gate branch. + sG_factor_a = sG_tma + sG_factor_b = sG_tma + sG_factor_base = sG_tma + sG_main_q = sG_tma + sG_main_k = sG_tma + # Vector G occupies FP32 stages, so BF16 scratch must preserve the + # corresponding byte stride when it is overlaid on that storage. + sK_g_outer = sK_g.layout + sK_neg_g_layout = cute.make_layout( + sK_g_outer.shape, + stride=(*sK_g_outer.stride[:-1], sK_g_outer.stride[-1] * 2), + ) + sK_neg_g_b_outer = kv_k_smem_layout_staged.outer + sK_neg_g_b_layout = cute.make_layout( + sK_neg_g_b_outer.shape, + stride=(*sK_neg_g_b_outer.stride[:-1], sK_neg_g_b_outer.stride[-1] * 2), + ) + sK_neg_g_ptr = cute.recast_ptr( + sG_tma.iterator, + swizzle_=k_smem_layout_staged.inner, + dtype=self.io_dtype, + ) + + sK_neg_g = cute.make_tensor(sK_neg_g_ptr, layout=sK_neg_g_layout) sK_neg_g_b = cute.make_tensor( - cute.recast_ptr(sK_neg_g_f32.iterator, swizzle_=kv_k_smem_layout_staged.inner, dtype=self.io_dtype), + cute.recast_ptr(sG_tma.iterator, swizzle_=kv_k_smem_layout_staged.inner, dtype=self.io_dtype), layout=sK_neg_g_b_layout, ) - # sK_neg_g = cute.make_tensor( - # cute.recast_ptr( - # sK_neg_g_f32.iterator, - # swizzle_=k_smem_layout_staged.inner, - # dtype=self.io_dtype), - # layout=sK_neg_g_f32.layout) # (((64,2),16),1,4,2):(((1,4096),64),0,1024,8192)> sV = storage.sV.get_tensor(v_smem_layout_staged.outer, swizzle=v_smem_layout_staged.inner) - # G (gate/g_cumsum) - NEW for KDA - sG = storage.sG.get_tensor(g_smem_layout_staged.outer, swizzle=g_smem_layout_staged.inner) # No swizzling for last row of exp(G) sG_last = self.get_smem_tensor_sG_last(storage, g_last_layout) @@ -1301,7 +1421,7 @@ def kernel( self.v_dtype, # utils.LayoutEnum.ROW_MAJOR, utils.LayoutEnum.COL_MAJOR, - (Constant.D, Constant.C), + (self.value_tile_size, Constant.C), self.v_stage, ) v_smem_layout_coalesce = cute.coalesce( @@ -1360,23 +1480,26 @@ def kernel( ) sK_flat_s2r = storage.sK.get_tensor(k_smem_layout_coalesce.outer, swizzle=k_smem_layout_coalesce.inner) - sG_flat_s2r_f32_fake = storage.sG.get_tensor(k_smem_layout_coalesce.outer, swizzle=k_smem_layout_coalesce.inner) - # CRITICAL FIX: When recasting F32 to BF16, we must double the stage stride - # so that the byte offset remains the same. - # F32 stage stride = 8192 elements = 32768 bytes - # BF16 stage stride should = 16384 elements = 32768 bytes k_smem_layout_bf16_outer = k_smem_layout_coalesce.outer - k_smem_layout_bf16_fixed = cute.make_layout( - k_smem_layout_bf16_outer.shape, - stride=(*k_smem_layout_bf16_outer.stride[:-1], k_smem_layout_bf16_outer.stride[-1] * 2), - ) + if cutlass.const_expr(self.scalar_gate): + # Scalar G uses dense BF16 scratch stages. + k_smem_layout_bf16_fixed = k_smem_layout_bf16_outer + else: + # Vector G has 32 KiB FP32 stages; retain their byte stride in the + # BF16 overlay. + k_smem_layout_bf16_fixed = cute.make_layout( + k_smem_layout_bf16_outer.shape, + stride=(*k_smem_layout_bf16_outer.stride[:-1], k_smem_layout_bf16_outer.stride[-1] * 2), + ) sG_flat_bf16 = cute.make_tensor( - cute.recast_ptr(sG_flat_s2r_f32_fake.iterator, swizzle_=k_smem_layout_coalesce.inner, dtype=self.io_dtype), + cute.recast_ptr(sG_tma.iterator, swizzle_=k_smem_layout_coalesce.inner, dtype=self.io_dtype), layout=k_smem_layout_bf16_fixed, ) - (_, hidx, bidx) = cute.arch.block_idx() - B, S, H, D = problem_size + (value_tile_idx, hidx, bidx) = cute.arch.block_idx() + value_tile_base = value_tile_idx * self.value_tile_size + B, S, H, HV, D = problem_size + qk_hidx = hidx // (HV // H) C = self.chunk_size # Varlen: compute per-CTA sequence boundary and domain offsets @@ -1499,20 +1622,22 @@ def kernel( # ------------------------------------------------------- # ((SWIZZLE_ATOM_M, REST_M), (SWIZZLE_ATOM_N, REST_N), (1, STAGES)) - g_smem_layout_epi = sm100_utils.make_smem_layout_epi( - self.g_dtype, - utils.LayoutEnum.ROW_MAJOR, - # G SMEM has the shape of - (Constant.C, Constant.D), - self.g_stage, - ) - # (C, (SWIZZLE_ATOM_N, REST_N), STAGES) - g_smem_layout_coalesce = cute.coalesce( - g_smem_layout_epi, - target_profile=(1, 1, 1), - ) - # ROW MAJOR - sG_flat = storage.sG.get_tensor(g_smem_layout_coalesce.outer, swizzle=g_smem_layout_coalesce.inner) + if cutlass.const_expr(self.scalar_gate): + # Logical vector view used by the existing safe-gate arithmetic. + # The D mode has zero stride, so no vector gate is materialized. + sG_flat = sG + else: + g_smem_layout_epi = sm100_utils.make_smem_layout_epi( + self.g_dtype, + utils.LayoutEnum.ROW_MAJOR, + (Constant.C, Constant.D), + self.g_stage, + ) + g_smem_layout_coalesce = cute.coalesce( + g_smem_layout_epi, + target_profile=(1, 1, 1), + ) + sG_flat = storage.sG.get_tensor(g_smem_layout_coalesce.outer, swizzle=g_smem_layout_coalesce.inner) # /////////////////////////////////////////////////////////////////////////////// # LOAD WARP # /////////////////////////////////////////////////////////////////////////////// @@ -1541,6 +1666,7 @@ def kernel( operand_mode="A", debug_name="Q", batch_idx=data_bidx, + head_idx=qk_hidx, ) tKsK, tKgK = self.tma_partition_for_mma_operand( @@ -1552,6 +1678,7 @@ def kernel( operand_mode="B", debug_name="K", batch_idx=data_bidx, + head_idx=qk_hidx, ) tVsV, tVgV = self.tma_partition_for_mma_operand( @@ -1565,17 +1692,17 @@ def kernel( batch_idx=data_bidx, ) - # G (gate) - NEW for KDA - tGsG, tGgG = self.tma_partition_for_mma_operand( - tma_atom_g, - tma_tensor_g_v, - sG, - self.qk_mma_tiler, # Same as Q - qk_tiled_mma, - operand_mode="A", - debug_name="G", - batch_idx=data_bidx, - ) + if cutlass.const_expr(not self.scalar_gate): + tGsG, tGgG = self.tma_partition_for_mma_operand( + tma_atom_g, + tma_tensor_g_v, + sG_tma, + self.qk_mma_tiler, + qk_tiled_mma, + operand_mode="A", + debug_name="G", + batch_idx=data_bidx, + ) if cutlass.const_expr(PRINT_DEBUG): print(f"tKsK={tKsK}") @@ -1588,29 +1715,25 @@ def kernel( idx = chunk_start // C should_debug = PRINT_DEBUG and tidx == warp_idx * 32 and hidx == 0 and bidx == 0 - # Gi (gate/g_cumsum) - NEW for KDA - g_handle = load_g_producer.acquire_and_advance() - cute.copy( - atom=tma_atom_g, - src=tGgG[None, idx, 0], - dst=tGsG[None, g_handle.index], - tma_bar_ptr=g_handle.barrier, - ) + # Vector G remains a TMA operand on the load warp. Scalar G + # is produced independently by load_beta_warp below, allowing + # this warp to issue Q/K/V immediately. + if cutlass.const_expr(not self.scalar_gate): + g_handle = load_g_producer.acquire_and_advance() + cute.copy( + atom=tma_atom_g, + src=tGgG[None, idx, 0], + dst=tGsG[None, g_handle.index], + tma_bar_ptr=g_handle.barrier, + ) - # Qi - # SRC: ((ATOM_V, REST_V), TILES_M, TILES_K) - # DST: ((ATOM_V, REST_V), INPUT_STAGE) q_handle = load_q_producer.acquire_and_advance() cute.copy( atom=tma_atom_q, - src=tQgQ[None, idx, 0], # source - dst=tQsQ[None, q_handle.index], # which stage + src=tQgQ[None, idx, 0], + dst=tQsQ[None, q_handle.index], tma_bar_ptr=q_handle.barrier, ) - - # Ki - # SRC: ((ATOM_V, REST_V), TILES_N, TILES_K) - # DST: ((ATOM_V, REST_V), INPUT_STAGE) k_handle = load_k_producer.acquire_and_advance() cute.copy( atom=tma_atom_k, @@ -1618,16 +1741,10 @@ def kernel( dst=tKsK[None, k_handle.index], tma_bar_ptr=k_handle.barrier, ) - - # Vi - # SRC: ((ATOM_V, REST_V), TILES_M, TILES_K) - # DST: ((ATOM_V, REST_V), INPUT_STAGE) v_handle = load_v_producer.acquire_and_advance() - if cutlass.const_expr(PRINT_DEBUG) and should_debug: - cute.printf("TMA v producer idx={}, v_handle={}", idx, v_handle.index) cute.copy( atom=tma_atom_v, - src=tVgV[None, 0, idx], + src=tVgV[None, value_tile_idx, idx], dst=tVsV[None, v_handle.index], tma_bar_ptr=v_handle.barrier, ) @@ -1661,7 +1778,7 @@ def kernel( tCtAcc=tCtAccSQ, tCrA=tCrState, tCrB=tCrQ_sq, - a_stage_idx=0, + a_stage_idx=kv16_handle.index, b_stage_idx=q_scaled_handle.index, acc_stage_idx=0, ) @@ -1734,7 +1851,6 @@ def kernel( # wait for sQ_K_scaled and decay(S) ready k_scaled2_handle = load_k_scaled2_consumer.wait_and_advance() - kv_decay_handle = kv_decay_consumer.wait_and_advance() kv_handle = kv_producer.acquire_and_advance() # launch S=K^T@NewV MMA @@ -1752,7 +1868,6 @@ def kernel( ) k_scaled2_handle.release() - kv_decay_handle.release() kv_handle.commit() # NOTE: Add a signal to notify the end of v consumption. @@ -1830,7 +1945,7 @@ def kernel( tCtAcc=tCtAccSQ, tCrA=tCrState, tCrB=tCrQ_sq, - a_stage_idx=0, + a_stage_idx=kv16_handle.index, b_stage_idx=q_handle.index, acc_stage_idx=0, ) @@ -2118,7 +2233,7 @@ def kernel( tCtAccKV_slice = tCtAccKV[((None, None), 0, 0, None)] ( tiled_copy_t2r_kv, - _, # thr_t2r + thr_t2r_kv, tTR_tKV, tTR_rKV, ) = self.tmem_load_partition_kv( @@ -2126,6 +2241,9 @@ def kernel( tState=tCtAccKV_slice, local_tidx=local_tidx, ) + state_tile = cute.dice(self.kv_mma_tiler, (1, 1, None)) + cM_state = cute.make_identity_tensor(state_tile) + tTR_cState = thr_t2r_kv.partition_D(cM_state) ############################################################ ( @@ -2138,15 +2256,35 @@ def kernel( ) tmem_store_rKV = cute.make_tensor(tTR_rKV.iterator, layout=tmem_store_rAccKV_f32.layout) - ( - tmem_store_kv, - tmem_store_tAccKV, - tmem_store_rAccKV, - ) = self.tmem_store_and_partition_acc( - local_tidx, - tCtAcc=tCtStateAsF32, - ) - tmem_store_rAccKVAsBF16 = cute.recast_tensor(tmem_store_rAccKV, dtype=self.io_dtype) + if cutlass.const_expr(self.split_value_tiles): + # For M=64, publish the BF16 state through the native operand-A + # TMEM layout. This is the same mapping used by the proven + # chunk_delta_h SM100 kernel; the original M=128 packed-FP32 + # alias has a different TV ownership and scrambles state reads + # in the following chunk. + state_store_atom = cute.make_copy_atom( + tcgen05.St16x128bOp(tcgen05.Repetition(16), tcgen05.Unpack.NONE), + self.io_dtype, + ) + tmem_store_kv = tcgen05.make_tmem_copy(state_store_atom, tCrState) + state_store_thr = tmem_store_kv.get_slice(local_tidx) + state_store_shape = cute.slice_( + state_store_thr.partition_S(tCrState).shape, + (None, None, None, None, 0), + ) + tmem_store_tAccKV = state_store_thr.partition_D(tCrState) + tmem_store_rAccKV = cute.make_rmem_tensor(state_store_shape, self.io_dtype) + tmem_store_rAccKVAsBF16 = tmem_store_rAccKV + else: + ( + tmem_store_kv, + tmem_store_tAccKV, + tmem_store_rAccKV, + ) = self.tmem_store_and_partition_acc( + local_tidx, + tCtAcc=tCtStateAsF32, + ) + tmem_store_rAccKVAsBF16 = cute.recast_tensor(tmem_store_rAccKV, dtype=self.io_dtype) ############################################################ if cutlass.const_expr(PRINT_DEBUG): @@ -2239,7 +2377,7 @@ def kernel( # ------------------------------------------------------- # V s2r partitions - NEW for KDA elementwise processing # shape_v = (Constant.C, Constant.D) - shape_v = (Constant.D, Constant.C) + shape_v = (self.value_tile_size, Constant.C) ( tiled_s2r_v, thr_s2r_v, @@ -2266,8 +2404,18 @@ def kernel( thr_mma_epi_half = tiled_mma_epi_half.get_slice(local_tidx) copy_op_qk_s2r = cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4) copy_op_qk_r2s = cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=False, num_matrices=4) - # FIXME: only 2 FP32 elements (64 bits) compatible with ldmatrix, how to change to 128? - copy_atom_g = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.g_dtype, num_bits_per_copy=64) + # A scalar-gate broadcast has a zero-stride D mode; use scalar + # copies so every logical element observes that stride. The + # generic vector path retains its 64-bit copy atom. + if cutlass.const_expr(self.scalar_gate): + gate_copy_bits = 32 + else: + gate_copy_bits = 64 + copy_atom_g = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.g_dtype, + num_bits_per_copy=gate_copy_bits, + ) # Half-size tiled copies for partitioned S2R tiled_load_g_half = cute.make_tiled_copy_A(copy_atom_g, tiled_mma_epi_half) thr_load_g_half = tiled_load_g_half.get_slice(local_tidx) @@ -2315,19 +2463,29 @@ def kernel( sK_flat_h0 = cute.make_tensor(sK_flat_s2r.iterator, layout=k_half_outer) sK_flat_h1 = cute.make_tensor(sK_flat_s2r.iterator + HALF_SMEM_ELEMS, layout=k_half_outer) - # G half SMEM views (FP32, from sG_flat iterator) - g_sml_epi_half = sm100_utils.make_smem_layout_epi( - self.g_dtype, - utils.LayoutEnum.ROW_MAJOR, - (Constant.C, Constant.HALF_D), - self.g_stage, - ) - g_sml_half = cute.coalesce(g_sml_epi_half, target_profile=(1, 1, 1)) - g_half_outer = cute.make_layout( - g_sml_half.outer.shape, stride=(*g_sml_half.outer.stride[:-1], g_sml_half.outer.stride[-1] * 2) - ) - sG_flat_h0 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) - sG_flat_h1 = cute.make_tensor(sG_flat.iterator + HALF_SMEM_ELEMS, layout=g_half_outer) + # G half SMEM views (FP32, from sG_flat iterator). + if cutlass.const_expr(self.scalar_gate): + g_half_outer = cute.make_layout( + (Constant.C, Constant.HALF_D, self.g_stage), + stride=(1, 0, Constant.C), + ) + sG_flat_h0 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) + # Both D halves broadcast the same per-token scalar. + sG_flat_h1 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) + else: + g_sml_epi_half = sm100_utils.make_smem_layout_epi( + self.g_dtype, + utils.LayoutEnum.ROW_MAJOR, + (Constant.C, Constant.HALF_D), + self.g_stage, + ) + g_sml_half = cute.coalesce(g_sml_epi_half, target_profile=(1, 1, 1)) + g_half_outer = cute.make_layout( + g_sml_half.outer.shape, + stride=(*g_sml_half.outer.stride[:-1], g_sml_half.outer.stride[-1] * 2), + ) + sG_flat_h0 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) + sG_flat_h1 = cute.make_tensor(sG_flat.iterator + HALF_SMEM_ELEMS, layout=g_half_outer) # Q_K_scaled half SMEM views (from sQ_K_scaled_flat iterator) qks_sml_epi_half = sm100_utils.make_smem_layout_epi( @@ -2388,12 +2546,26 @@ def index_transform_half(index_q, index_k): # State shape: (D, D) per (H, B), stored as FP32 if cutlass.const_expr(self.has_initial_state): # Load initial state from GMEM to RMEM respecting TMEM partition. - # TMEM stores S^T (transposed), so flat[i] = state[local_tidx, i] - # Each thread owns key position local_tidx, D elements cover value positions. init_state_chunk = initial_state[None, None, (hidx, bidx)] - init_flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) - for init_i in cutlass.range(0, Constant.D, unroll=0): - init_flat[init_i] = init_state_chunk[local_tidx, init_i] + if cutlass.const_expr(self.split_value_tiles): + # M=64 uses 16 TMEM datapaths/warp, so thread id is no + # longer the K coordinate. Follow the actual T2R TV map. + for init_i in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + value_coord, key_coord = tTR_cState[init_i] + tTR_rKV[init_i] = init_state_chunk[ + key_coord, + value_tile_base + value_coord, + ] + else: + init_flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) + for init_i in cutlass.range(0, self.value_tile_size, unroll=0): + init_flat[init_i] = init_state_chunk[ + local_tidx, + value_tile_base + init_i, + ] # Store FP32 state to TMEM for accumulation (tCtAccKV) init_tmem_store_tKVi = tmem_store_tAccKV_f32[None, None, None, None, 0] @@ -2413,9 +2585,8 @@ def index_transform_half(index_q, index_k): # safe_gate version for chunk_start in cutlass.range(0, seq_len, C, unroll=0): idx = chunk_start // C - if cutlass.const_expr(self.is_varlen): - valid_len_chunk = seq_len - chunk_start - else: + valid_len_chunk = seq_len - chunk_start + if valid_len_chunk > C: valid_len_chunk = C # ============================================================ @@ -2445,6 +2616,14 @@ def index_transform_half(index_q, index_k): for half_idx in cutlass.range_constexpr(2): tQrG_persists.append(cute.make_fragment_like(tQrQ_half_0, dtype=self.g_dtype)) + if cutlass.const_expr(self.scalar_gate): + if local_tidx == 0: + sG_last[0, g_stage_idx] = sG_tma[ + valid_len_chunk - 1, + 0, + g_stage_idx, + ] + # Merged g_last + Q gating path: single G half-load per half if idx != 0 or cutlass.const_expr(self.has_initial_state): q_stage_idx = q_handle.index @@ -2456,38 +2635,50 @@ def index_transform_half(index_q, index_k): # S2R G half into persistent fragment tQrG_half_cv = thr_load_g_half.retile(tQrG_persists[half_idx]) - cute.copy(tiled_load_g_half, tQsG_h[half_idx][None, None, None, g_stage_idx], tQrG_half_cv) - - # Write g_last half (before exp transforms g values) - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if cutlass.const_expr(self.is_varlen): + if cutlass.const_expr(self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, _ = index_transform_half(*tQcMq_half[i]) + tQrG_persists[half_idx][i] = sG_main_q[index_q, g_stage_idx] + else: + cute.copy( + tiled_load_g_half, + tQsG_h[half_idx][None, None, None, g_stage_idx], + tQrG_half_cv, + ) + + # Vector G needs one g_last value per feature; + # scalar G wrote its single value above. + if cutlass.const_expr(not self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) if valid_len_chunk < C: if index_q == valid_len_chunk - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] else: if index_q == Constant.C - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] - else: - if index_q == Constant.C - 1: - sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] - # exp(g) half in-place — persists for K gating reuse - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - tQrG_persists[half_idx][i] = cute.exp2(tQrG_persists[half_idx][i], fastmath=self.use_fast_math) + # exp(g) half in-place — persists for K gating reuse. + # Scalar G was already exponentiated once/token by + # the producer warp, rather than once/feature here. + if cutlass.const_expr(not self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + tQrG_persists[half_idx][i] = cute.exp2( + tQrG_persists[half_idx][i], + fastmath=self.use_fast_math, + ) # S2R Q half tQrQ_half = cute.make_fragment_like(tQrQ_half_0, self.q_dtype) tQrQ_half_cv = thr_load_qk_half.retile(tQrQ_half) cute.copy(tiled_load_qk_half, tQsQ_h[half_idx][None, None, None, q_stage_idx], tQrQ_half_cv) - # Zero Q for invalid positions (varlen only) - if cutlass.const_expr(self.is_varlen): + # Zero Q for any partial tail (varlen or fixed-B). + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) if valid_len_chunk < C: - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if index_q >= valid_len_chunk: - tQrQ_half[i] = self.q_dtype(0.0) + if index_q >= valid_len_chunk: + tQrQ_half[i] = self.q_dtype(0.0) # Q gating: Q' = Q * exp(g) * scale for i in cutlass.range_constexpr(cute.size(tQcMq_half)): @@ -2510,19 +2701,21 @@ def index_transform_half(index_q, index_k): for half_idx in cutlass.range_constexpr(2): k_offset = half_idx * Constant.HALF_D tQrG_half_cv = thr_load_g_half.retile(tQrG_persists[half_idx]) - cute.copy(tiled_load_g_half, tQsG_h[half_idx][None, None, None, g_stage_idx], tQrG_half_cv) - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if cutlass.const_expr(self.is_varlen): + if cutlass.const_expr(not self.scalar_gate): + cute.copy( + tiled_load_g_half, + tQsG_h[half_idx][None, None, None, g_stage_idx], + tQrG_half_cv, + ) + if cutlass.const_expr(not self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) if valid_len_chunk < C: if index_q == valid_len_chunk - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] else: if index_q == Constant.C - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] - else: - if index_q == Constant.C - 1: - sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] # ==================================================== # Partitioned S2R: K gating (2 half-passes) @@ -2547,13 +2740,12 @@ def index_transform_half(index_q, index_k): tQrK_half_cv = thr_load_qk_half.retile(tQrK_half) cute.copy(tiled_load_qk_half, tQsK_h[half_idx][None, None, None, k_stage_idx], tQrK_half_cv) - # Zero K for invalid positions (varlen only) - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C: - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if index_q >= valid_len_chunk: - tQrK_half[i] = self.q_dtype(0.0) + # Zero K for any partial tail (varlen or fixed-B). + if valid_len_chunk < C: + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) + if index_q >= valid_len_chunk: + tQrK_half[i] = self.q_dtype(0.0) # K gating: K' = K * exp(g) — reuse persisted exp2(g) for i in cutlass.range_constexpr(cute.size(tQcMq_half)): @@ -2676,7 +2868,6 @@ def index_transform_half(index_q, index_k): # decay S, S=S*g_last # FIXME: currently do not support initial state, # so only decay S after first block K^T@NewV - kv_decay_handle = kv_decay_producer.acquire_and_advance() if idx != 0 or cutlass.const_expr(self.has_initial_state): # NOTE: TMEM S is always ready here # T2R S @@ -2685,7 +2876,10 @@ def index_transform_half(index_q, index_k): cute.arch.fence_view_async_tmem_load() # decay S - flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) + flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) self.scale_state(flat, sG_last[None, g_stage_idx]) # R2T S @@ -2694,8 +2888,6 @@ def index_transform_half(index_q, index_k): cute.copy(tmem_store_kv_f32, tmem_store_rKV, tmem_store_tKVi) cute.arch.fence_view_async_tmem_store() - kv_decay_handle.commit() - # ==================================================== # Partitioned S2R: K^T gating — exp(g_last-g)*K # ==================================================== @@ -2708,20 +2900,28 @@ def index_transform_half(index_q, index_k): # S2R G half tQrG_half = cute.make_fragment_like(tQrQ_half_0, dtype=self.g_dtype) tQrG_half_cv = thr_load_g_half.retile(tQrG_half) - cute.copy(tiled_load_g_half, tQsG_h[half_idx][None, None, None, g_stage_idx], tQrG_half_cv) + if cutlass.const_expr(self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, _ = index_transform_half(*tQcMq_half[i]) + tQrG_half[i] = sG_main_k[index_q, g_stage_idx] + else: + cute.copy( + tiled_load_g_half, + tQsG_h[half_idx][None, None, None, g_stage_idx], + tQrG_half_cv, + ) # S2R K half tQrK_half = cute.make_fragment_like(tQrQ_half_0, dtype=self.k_dtype) tQrK_half_cv = thr_load_qk_half.retile(tQrK_half) cute.copy(tiled_load_qk_half, tQsK_h[half_idx][None, None, None, k_stage_idx], tQrK_half_cv) - # Zero K half for invalid positions (varlen only) - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C: - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if index_q >= valid_len_chunk: - tQrK_half[i] = self.k_dtype(0.0) + # Zero K half for any partial tail. + if valid_len_chunk < C: + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) + if index_q >= valid_len_chunk: + tQrK_half[i] = self.k_dtype(0.0) # K^T gating: exp(g_last - g) * K for i in cutlass.range_constexpr(cute.size(tQcMq_half)): @@ -2729,7 +2929,14 @@ def index_transform_half(index_q, index_k): g_last_val = sG_last[index_k + k_offset, g_stage_idx] k_i = tQrK_half[i].to(cutlass.Float32) g_i = tQrG_half[i] - tQrK_half[i] = (cute.exp2(g_last_val - g_i, fastmath=self.use_fast_math) * k_i).to(self.k_dtype) + if cutlass.const_expr(self.scalar_gate): + gate_k_i = g_i + else: + gate_k_i = cute.exp2( + g_last_val - g_i, + fastmath=self.use_fast_math, + ) + tQrK_half[i] = (gate_k_i * k_i).to(self.k_dtype) # R2S K^T half to sQ_K_scaled tQrK_half_cv_src = thr_store_qk_half.retile(tQrK_half) @@ -2798,6 +3005,7 @@ def index_transform_half(index_q, index_k): cute.copy(tiled_copy_t2r_kv, tTR_tKVi, tTR_rKV) cute.arch.fence_view_async_tmem_load() + if idx != final_blk: # Store as a separated BF16 state for QS and KS MMA before decay # tmem_store_rAccKVAsBF16 point to the same rmem as tmem_store_rKV @@ -2819,11 +3027,28 @@ def index_transform_half(index_q, index_k): # idx == final_blk: output final state immediately to minimize tTR_rKV lifetime if cutlass.const_expr(self.output_final_state): # Write FP32 state from RMEM to GMEM - # TMEM stores S^T (transposed), so flat[i] = state[local_tidx, i] state_out = final_state[None, None, (hidx, bidx)] - out_flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) - for out_i in cutlass.range(0, Constant.D, unroll=0): - state_out[local_tidx, out_i] = out_flat[out_i] + if cutlass.const_expr(self.split_value_tiles): + # Core ABI: h0 is contiguous [V,K], while ht is + # returned contiguous [K,V]. With fstate's + # CuTe stride (1,D), storing (V,K) below creates + # the required row-major [K,V] result. + for out_i in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + value_coord, key_coord = tTR_cState[out_i] + state_out[ + value_tile_base + value_coord, + key_coord, + ] = tTR_rKV[out_i] + else: + out_flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) + for out_i in cutlass.range(0, self.value_tile_size, unroll=0): + state_out[ + local_tidx, + value_tile_base + out_i, + ] = out_flat[out_i] # release KV kv_handle.release() @@ -2838,9 +3063,8 @@ def index_transform_half(index_q, index_k): else: for chunk_start in cutlass.range(0, seq_len, C, unroll=0): idx = chunk_start // C - if cutlass.const_expr(self.is_varlen): - valid_len_chunk = seq_len - chunk_start - else: + valid_len_chunk = seq_len - chunk_start + if valid_len_chunk > C: valid_len_chunk = C # ============================================================ @@ -2907,19 +3131,10 @@ def index_transform_half(index_q, index_k): # Element-wise processing avoids bulk .load()/.to() creating # ~300+ register SSA vectors from G, Q, K simultaneously for _zr in cutlass.range(0, Constant.C, unroll_full=True): - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C and _zr >= valid_len_chunk: - tRS_rQ[0, _zr, 0] = self.io_dtype(0.0) - tRS_rK[0, _zr, 0] = self.io_dtype(0.0) - tRS_rG_bf16[0, _zr, 0] = self.io_dtype(0.0) - else: - g_i = tRS_rG[0, _zr, 0] - exp_g_i = cute.exp2(g_i, fastmath=self.use_fast_math) - q_i = tRS_rQ[0, _zr, 0].to(cutlass.Float32) - tRS_rQ[0, _zr, 0] = (q_i * exp_g_i * self.scale).to(self.io_dtype) - k_i = tRS_rK[0, _zr, 0].to(cutlass.Float32) - tRS_rK[0, _zr, 0] = (k_i * exp_g_i).to(self.io_dtype) - tRS_rG_bf16[0, _zr, 0] = (k_i * cute.exp2(-g_i, fastmath=self.use_fast_math)).to(self.io_dtype) + if valid_len_chunk < C and _zr >= valid_len_chunk: + tRS_rQ[0, _zr, 0] = self.io_dtype(0.0) + tRS_rK[0, _zr, 0] = self.io_dtype(0.0) + tRS_rG_bf16[0, _zr, 0] = self.io_dtype(0.0) else: g_i = tRS_rG[0, _zr, 0] exp_g_i = cute.exp2(g_i, fastmath=self.use_fast_math) @@ -2993,13 +3208,13 @@ def index_transform_half(index_q, index_k): # ------------------------------------------------------------ # NOTE: Save exp(g) of last VALID row to rG_last for state update in next chunk # For full chunks, directly use C-1; only loop for partial chunks (varlen only) - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C: - rG_last = exp_g[valid_len_chunk - 1] - else: - rG_last = exp_g[Constant.C - 1] + rG_last = cutlass.Float32(1.0) + if valid_len_chunk < C: + for _zr in cutlass.range(0, Constant.C, unroll_full=True): + if _zr == valid_len_chunk - 1: + rG_last = cute.exp2(tRS_rG[0, _zr, 0], fastmath=self.use_fast_math) else: - rG_last = exp_g[Constant.C - 1] + rG_last = cute.exp2(tRS_rG[0, Constant.C - 1, 0], fastmath=self.use_fast_math) # NOTE: each thread save one element sG_last[local_tidx, g_stage_idx] = rG_last @@ -3298,7 +3513,10 @@ def index_transform_half(index_q, index_k): cute.print_tensor(tTR_rKV) self.cuda_wg_sync_barrier.arrive_and_wait() - flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) + flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) # FIXME self.cuda_wg_sync_barrier.arrive_and_wait() @@ -3340,9 +3558,25 @@ def index_transform_half(index_q, index_k): # idx == final: output final state immediately to minimize tTR_rKV lifetime if cutlass.const_expr(self.output_final_state): state_out = final_state[None, None, (hidx, bidx)] - out_flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) - for out_i in cutlass.range(0, Constant.D, unroll=0): - state_out[local_tidx, out_i] = out_flat[out_i] + if cutlass.const_expr(self.split_value_tiles): + # See the safe-gate branch above for the state + # ABI and CuTe-to-row-major axis convention. + for out_i in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + value_coord, key_coord = tTR_cState[out_i] + state_out[ + value_tile_base + value_coord, + key_coord, + ] = tTR_rKV[out_i] + else: + out_flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) + for out_i in cutlass.range(0, self.value_tile_size, unroll=0): + state_out[ + local_tidx, + value_tile_base + out_i, + ] = out_flat[out_i] # NOTE: only release v after PV and State=KV has been consumed end_v_handle = end_v_consumer.wait_and_advance() @@ -3371,8 +3605,15 @@ def index_transform_half(index_q, index_k): copy_op_A_s2r = cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4) copy_op_B_s2r = cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4) copy_op_r2s = cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2) - # FIXME: only 2 FP32 elements (64 bits) compatible with ldmatrix, how to change to 128? - copy_g_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.g_dtype, num_bits_per_copy=64) + if cutlass.const_expr(self.scalar_gate): + gate_copy_bits = 32 + else: + gate_copy_bits = 64 + copy_g_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.g_dtype, + num_bits_per_copy=gate_copy_bits, + ) G_Q_tiled_copy = cute.make_tiled_copy_A(copy_g_atom, tiled_mma_subchunk) G_Kt_tiled_copy = cute.make_tiled_copy_B(copy_g_atom, tiled_mma_subchunk) Q_tiled_copy = cute.make_tiled_copy_A(cute.make_copy_atom(copy_op_A_s2r, self.q_dtype), tiled_mma_subchunk) @@ -3391,6 +3632,9 @@ def index_transform_half(index_q, index_k): # index tensor cMqk_subchunk = cute.make_identity_tensor(self.qk_kk_subchunk_mma_tiler[:2]) tQKcMqk_subchunk = thr_mma_subchunk.partition_C(cMqk_subchunk) + cA_subchunk = cute.make_identity_tensor((Constant.SC, Constant.BK_SC)) + tAcA_subchunk = thr_mma_subchunk.partition_A(cA_subchunk) + tBcB_subchunk = thr_mma_subchunk.partition_B(cA_subchunk) def index_transform(index_q, index_k): return ( @@ -3439,13 +3683,39 @@ def index_transform(index_q, index_k): sQqk_curr = sQ_flat[None, None, load_q_consumer._PipelineConsumer__state.index] sKqk_curr = sK_flat[None, None, load_k_consumer._PipelineConsumer__state.index] sGqkq_curr = sG_flat[None, None, load_g_consumer._PipelineConsumer__state.index] + if cutlass.const_expr(self.scalar_gate): + sG_scalar_curr = sG_tma[None, 0, load_g_consumer._PipelineConsumer__state.index] + sG_factor_a_curr = sG_factor_a[None, load_g_consumer._PipelineConsumer__state.index] + sG_factor_b_curr = sG_factor_b[None, None, load_g_consumer._PipelineConsumer__state.index] + else: + # The helper signatures are shared with the scalar + # specialization, but this argument is compile-time + # unused by the vector-gate branch. Keep a congruent + # tensor here instead of indexing vector sG as scalar. + sG_scalar_curr = sGqkq_curr + sG_factor_a_curr = sGqkq_curr + sG_factor_b_curr = sGqkq_curr sBeta_curr = sBeta[None, load_beta_consumer._PipelineConsumer__state.index] # (_16,(_32,_2),_4,(_1,_2)):(_32,(_1,_2048),_512,(_0,_4096)) sQqk_slice = cute.flat_divide(sQqk_curr, tiler_subchunk_qk) sKqk_slice = cute.flat_divide(sKqk_curr, tiler_subchunk_qk) # (_16,(_64,_1),_4,(_1,_2)):(_64,(_1,_0),_1024,(_0,_4096)) - sGqkq_slice = cute.flat_divide(sGqkq_curr, tiler_subchunk_g) + if cutlass.const_expr(self.scalar_gate): + # `flat_divide` requires an injective source layout and + # therefore cannot divide the zero-stride broadcast + # view. Build the same logical subchunk coordinates + # directly: token = row + 16*subchunk, while every D + # coordinate aliases that token's scalar gate. + sGqkq_slice = cute.make_tensor( + sGqkq_curr.iterator, + layout=cute.make_layout( + (16, (64, 1), 4, (1, 2)), + stride=(1, (0, 0), 16, (0, 0)), + ), + ) + else: + sGqkq_slice = cute.flat_divide(sGqkq_curr, tiler_subchunk_g) sBeta_slice = cute.flat_divide(sBeta_curr, tiler_subchunk_beta) # Acc results @@ -3503,8 +3773,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3515,7 +3788,10 @@ def index_transform(index_q, index_k): tGsGfirst_0_j_kt = G_Kt_thr_copy.partition_S(sG_first_0_j) tGrGfirst_0_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_0_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_0_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_0_j_kt, tGrGfirst_0_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_0_j_kt.fill(sG_scalar_curr[0]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_0_j_kt, tGrGfirst_0_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3525,7 +3801,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 0, sKqk_slice, tGrGfirst_0_j_kt, ) @@ -3595,8 +3875,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3607,7 +3890,10 @@ def index_transform(index_q, index_k): tGsGfirst_3_j_kt = G_Kt_thr_copy.partition_S(sG_first_3_j) tGrGfirst_3_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_3_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_3_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_3_j_kt, tGrGfirst_3_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_3_j_kt.fill(sG_scalar_curr[3 * Constant.SC]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_3_j_kt, tGrGfirst_3_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3617,7 +3903,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3634,7 +3924,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3651,7 +3945,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3668,7 +3966,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3754,8 +4056,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3766,7 +4071,10 @@ def index_transform(index_q, index_k): tGsGfirst_1_j_kt = G_Kt_thr_copy.partition_S(sG_first_1_j) tGrGfirst_1_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_1_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_1_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_1_j_kt, tGrGfirst_1_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_1_j_kt.fill(sG_scalar_curr[Constant.SC]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_1_j_kt, tGrGfirst_1_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3776,7 +4084,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 1, sKqk_slice, tGrGfirst_1_j_kt, ) @@ -3793,7 +4105,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 1, sKqk_slice, tGrGfirst_1_j_kt, ) @@ -3865,8 +4181,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3877,7 +4196,10 @@ def index_transform(index_q, index_k): tGsGfirst_2_j_kt = G_Kt_thr_copy.partition_S(sG_first_2_j) tGrGfirst_2_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_2_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_2_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_2_j_kt, tGrGfirst_2_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_2_j_kt.fill(sG_scalar_curr[2 * Constant.SC]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_2_j_kt, tGrGfirst_2_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3887,7 +4209,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 2, sKqk_slice, tGrGfirst_2_j_kt, ) @@ -3904,7 +4230,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 2, sKqk_slice, tGrGfirst_2_j_kt, ) @@ -3921,7 +4251,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 2, sKqk_slice, tGrGfirst_2_j_kt, ) @@ -3977,9 +4311,8 @@ def index_transform(index_q, index_k): cute.copy(tiled_load_qk, thr_load_qk.partition_S(sQK_curr), tQKrQK) cute.copy(tiled_load_kk, thr_load_kk.partition_S(sKK_inv_curr), tKKrKK) # triangular mask and boundary mask - if cutlass.const_expr(self.is_varlen): - valid_len_chunk = seq_len - chunk_start - else: + valid_len_chunk = seq_len - chunk_start + if valid_len_chunk > C: valid_len_chunk = C self.apply_qk_kk_mask(tQKcMqk, tQKrQK, tKKrKK, valid_len_chunk) # R2S QK/KK @@ -4007,7 +4340,7 @@ def index_transform(index_q, index_k): # S2R, scale with beta, convert to BF16, store back to smem `sM` # TODO: make a repro for the cutedsl team - self.scale_M_inverse_with_beta(local_tidx, sBeta, curr_sM_f16, curr_sM) + self.scale_M_inverse_with_beta(local_tidx, sBeta_curr, curr_sM_f16, curr_sM) cute.arch.fence_proxy( cute.arch.ProxyKind.async_shared, @@ -4090,8 +4423,8 @@ def index_transform(index_q, index_k): o_tail = cute.make_tensor( o_gmem.iterator, cute.make_layout( - (D, new_S, (H, Int32(1))), - stride=(Int32(1), D * H, (D, D * H * S)), + (D, new_S, (HV, Int32(1))), + stride=(Int32(1), D * HV, (D, D * HV * S)), ), ) # Initialize: copy original TMA descriptor to workspace @@ -4118,13 +4451,21 @@ def index_transform(index_q, index_k): cute.copy( tma_atom_o, bSG_sO[None, smem_o_handle.index], - bSG_gO[(None, 0, 0, 0, idx)], + bSG_gO[(None, 0, 0, value_tile_idx, idx)], tma_desc_ptr=self._tensormap_mgr.get_tensormap_ptr(ws_desc_ptr, cute.AddressSpace.generic), ) else: - cute.copy(tma_atom_o, bSG_sO[None, smem_o_handle.index], bSG_gO[(None, 0, 0, 0, idx)]) + cute.copy( + tma_atom_o, + bSG_sO[None, smem_o_handle.index], + bSG_gO[(None, 0, 0, value_tile_idx, idx)], + ) else: - cute.copy(tma_atom_o, bSG_sO[None, smem_o_handle.index], bSG_gO[(None, 0, 0, 0, idx)]) + cute.copy( + tma_atom_o, + bSG_sO[None, smem_o_handle.index], + bSG_gO[(None, 0, 0, value_tile_idx, idx)], + ) # Ensure smem_o has been released. cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0, read=True) @@ -4135,6 +4476,30 @@ def index_transform(index_q, index_k): local_tidx = tidx % (self.threads_per_warp * len([self.load_beta_warp_id])) for chunk_start in cutlass.range(0, seq_len, C, unroll=0): idx = chunk_start // C + if cutlass.const_expr(self.scalar_gate): + g_handle = load_g_producer.acquire_and_advance() + self.produce_scalar_gate_chunk( + g, + sG_tma, + sG_factor_a, + sG_factor_b, + sG_factor_base, + sG_main_q, + sG_main_k, + chunk_start, + seq_len, + tok_offset, + hidx, + data_bidx, + local_tidx, + g_handle.index, + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + g_handle.commit() + # Load beta into smem beta_handle = load_beta_producer.acquire_and_advance() # Fence due to normal load @@ -4149,8 +4514,8 @@ def index_transform(index_q, index_k): if cutlass.const_expr(self.is_varlen): beta_v = cute.domain_offset((tok_offset, (0, 0)), beta) beta_chunk = beta_v[(None, (hidx, data_bidx))] - beta_chunk_layout = cute.make_layout((C, 1), stride=(H, 0)) - beta_chunk = cute.make_tensor(beta_chunk.iterator + s_idx * H, layout=beta_chunk_layout) + beta_chunk_layout = cute.make_layout((C, 1), stride=(HV, 0)) + beta_chunk = cute.make_tensor(beta_chunk.iterator + s_idx * HV, layout=beta_chunk_layout) if cutlass.const_expr(PRINT_DEBUG): print(f"sBeta: {sBeta}") @@ -4160,12 +4525,12 @@ def index_transform(index_q, index_k): valid_len_beta = seq_len - chunk_start for data_idx in cutlass.range(local_tidx, Constant.C, self.threads_per_warp): if data_idx < valid_len_beta: - sBeta[data_idx, 0] = beta_chunk[data_idx, 0] + sBeta[data_idx, beta_handle.index] = beta_chunk[data_idx, 0] else: - sBeta[data_idx, 0] = cutlass.Float32(0.0) + sBeta[data_idx, beta_handle.index] = cutlass.Float32(0.0) else: for data_idx in cutlass.range(local_tidx, Constant.C, self.threads_per_warp): - sBeta[data_idx, 0] = beta_chunk[data_idx, 0] + sBeta[data_idx, beta_handle.index] = beta_chunk[data_idx, 0] # Fence cute.arch.fence_proxy( @@ -4190,12 +4555,19 @@ def scale_state(self, flat: cute.Tensor, sG_last: cute.Tensor) -> cute.Tensor: kv_f32 = flat if cutlass.const_expr(PRINT_DEBUG): print(f"kv_f32: {kv_f32}") - for i in cutlass.range(0, Constant.D, unroll_full=True): - if not cutlass.const_expr(self.safe_gate): - kv_f32[i] = kv_f32[i] * sG_last[i] - else: - # NOTE: when safe_gate=True, sG_last stores the original G values - kv_f32[i] = kv_f32[i] * cute.exp2(sG_last[i], fastmath=self.use_fast_math) + if cutlass.const_expr(self.safe_gate and self.scalar_gate): + # Qwen GDN broadcasts one scalar decay over the full K dimension. + # Compute exp2 once per thread instead of repeating it D times. + decay = cute.exp2(sG_last[0], fastmath=self.use_fast_math) + for i in cutlass.range(0, self.value_tile_size, unroll_full=True): + kv_f32[i] = kv_f32[i] * decay + else: + for i in cutlass.range(0, self.value_tile_size, unroll_full=True): + if not cutlass.const_expr(self.safe_gate): + kv_f32[i] = kv_f32[i] * sG_last[i] + else: + # NOTE: when safe_gate=True, sG_last stores the original G values + kv_f32[i] = kv_f32[i] * cute.exp2(sG_last[i], fastmath=self.use_fast_math) return kv_f32 def tmem_load_kv16(self, local_tidx, tState): @@ -4320,12 +4692,23 @@ def tmem_load_partition_kv(self, mma_tiler, tState, local_tidx): # use_2cta_instrs=False, # ) - # In KDA, we need to make tv-layout row-wise to perform diagonal op. - copy_atom_t2r = cute.make_copy_atom( - # 32b x 32, TODO: ADJUST RMEM PEAK - tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), - self.acc_dtype, - ) + # In KDA, we need a row-wise TV layout for the state operations. The + # original M=128 tile has 32 datapaths per warp, while the split M=64 + # tile has 16; select the matching 16dp load atom for the latter. + if cutlass.const_expr(self.split_value_tiles): + copy_atom_t2r = sm100_utils.get_tmem_load_op( + mma_tiler, + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.acc_dtype, + mma_tiler[:2], + use_2cta_instrs=False, + ) + else: + copy_atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), + self.acc_dtype, + ) fake_sState = cute.make_tensor( cute.make_ptr(self.io_dtype, 0, cute.AddressSpace.smem), cute.dice(self.kv_mma_tiler, (1, 1, None)), @@ -4352,10 +4735,17 @@ def make_tmem_load_and_partition(self, copy_atom_t2r, tmem_tensor, tmem_tile_coo def tmem_store_and_partition_acc(self, local_tidx, tCtAcc): dtype = tCtAcc.element_type - copy_atom_r2t = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition(32), tcgen05.Unpack.NONE), - dtype, - ) + if cutlass.const_expr(self.split_value_tiles): + # Inverse of the split state's 16dp Ld16x256b x16 mapping. + copy_atom_r2t = cute.make_copy_atom( + tcgen05.St16x256bOp(tcgen05.Repetition(16), tcgen05.Unpack.NONE), + dtype, + ) + else: + copy_atom_r2t = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32), tcgen05.Unpack.NONE), + dtype, + ) tiled_r2t = tcgen05.make_tmem_copy(copy_atom_r2t, tCtAcc) thr_r2t = tiled_r2t.get_slice(local_tidx) @@ -4596,22 +4986,25 @@ def epilog_tmem_copy_and_partition( """ # Make tiledCopy for tensor memory load epitile = mma_tiler[:2] - assert epitile[0] == 128 + assert epitile[0] == self.value_tile_size # TODO: 32dp ease DEBUGGING - copy_atom_t2r = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), - self.acc_dtype, - ) - # copy_atom_t2r = sm100_utils.get_tmem_load_op( - # mma_tiler, - # # self.o_layout, - # # TODO - # utils.LayoutEnum.ROW_MAJOR, - # self.io_dtype, - # self.acc_dtype, - # epitile, - # use_2cta_instrs, - # ) + if cutlass.const_expr(self.split_value_tiles): + # M=64 has only 16 TMEM datapaths per warp. Let CUTLASS select the + # matching 16dp load atom (currently Ld16x256b) instead of forcing + # the 32dp atom used by the original M=128 epilogue. + copy_atom_t2r = sm100_utils.get_tmem_load_op( + mma_tiler, + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.acc_dtype, + epitile, + use_2cta_instrs, + ) + else: + copy_atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), + self.acc_dtype, + ) # (EPI_TILE_M, EPI_TILE_N, 1, 1, STAGE) tAcc_epi = cute.flat_divide( # ((EPI_TILE_M, EPI_TILE_N), EPI_M, EPI_N, STAGE) @@ -5505,10 +5898,13 @@ def local_tile_partition_for_mma_operand( debug_name=None, no_cta_coord=False, batch_idx=None, + head_idx=None, ): _, hidx, bidx = cute.arch.block_idx() if batch_idx is not None: bidx = batch_idx + if cutlass.const_expr(head_idx is not None): + hidx = head_idx # Local_tile partition global tensors # x: (0,0,0,0) o (M,K,(H,B)):(1@1,1@0,(1@2,1@3)) # (MMATile_M, MMATile_K, TILES_M, TILES_K, (H, B)) @@ -5562,6 +5958,7 @@ def tma_partition_for_mma_operand( operand_mode, debug_name=None, batch_idx=None, + head_idx=None, ): tCgX = self.local_tile_partition_for_mma_operand( tensor_x=tma_tensor_x, @@ -5570,6 +5967,7 @@ def tma_partition_for_mma_operand( operand_mode=operand_mode, debug_name=debug_name, batch_idx=batch_idx, + head_idx=head_idx, ) # Partition shared tensor with regard to TMA # ((ATOM_V, REST_V), INPUT_STAGE) @@ -5585,6 +5983,191 @@ def tma_partition_for_mma_operand( # =========== # Utility functions for Ampere-style mma.sync, used for subchunk computation + @cute.jit + def produce_scalar_gate_chunk( + self, + g: cute.Tensor, + sG_tma: cute.Tensor, + sG_factor_a: cute.Tensor, + sG_factor_b: cute.Tensor, + sG_factor_base: cute.Tensor, + sG_main_q: cute.Tensor, + sG_main_k: cute.Tensor, + chunk_start, + seq_len, + tok_offset, + hidx, + data_bidx, + lane_idx, + stage_idx, + ): + """Produce one scalar-G stage with a single 32-thread warp.""" + for lane_slot in cutlass.range_constexpr(2): + token_in_chunk = lane_idx + lane_slot * self.threads_per_warp + token_in_seq = chunk_start + token_in_chunk + if token_in_seq < seq_len: + token_in_data = tok_offset + token_in_seq + sG_tma[token_in_chunk, 0, stage_idx] = g[ + token_in_data, + 0, + (hidx, data_bidx), + ] + else: + sG_tma[token_in_chunk, 0, stage_idx] = cutlass.Float32(0.0) + cute.arch.sync_warp() + + if cutlass.const_expr(self.fuse_scalar_cumsum): + # Inclusive chunk-local cumsum in log2 space. The producer warp + # owns exactly two 32-token halves. + g_lo = sG_tma[lane_idx, 0, stage_idx] * Constant.RCP_LN2 + g_hi = ( + sG_tma[lane_idx + self.threads_per_warp, 0, stage_idx] + * Constant.RCP_LN2 + ) + for scan_step in cutlass.range_constexpr(5): + scan_offset = 1 << scan_step + add_lo = cute.arch.shuffle_sync_up( + g_lo, + scan_offset, + mask_and_clamp=0, + ) + add_hi = cute.arch.shuffle_sync_up( + g_hi, + scan_offset, + mask_and_clamp=0, + ) + if lane_idx >= scan_offset: + g_lo += add_lo + g_hi += add_hi + g_lo_total = cute.arch.shuffle_sync(g_lo, 31) + g_hi += g_lo_total + sG_tma[lane_idx, 0, stage_idx] = g_lo + sG_tma[lane_idx + self.threads_per_warp, 0, stage_idx] = g_hi + cute.arch.sync_warp() + + # Factor pairwise gates as: + # A[t] = exp2(g[t]-g[key_subchunk_base]) + # Base[q,k] = exp2(g[q_base]-g[k_base]) + # B[q,t] = Base[q,key_subchunk(t)] / A[t] + valid_len_g = seq_len - chunk_start + if valid_len_g > Constant.C: + valid_len_g = Constant.C + g_last_scalar = sG_tma[valid_len_g - 1, 0, stage_idx] + for lane_slot in cutlass.range_constexpr(2): + token_in_chunk = lane_idx + lane_slot * self.threads_per_warp + if token_in_chunk < valid_len_g: + g_i = sG_tma[token_in_chunk, 0, stage_idx] + own_base = (token_in_chunk // Constant.SC) * Constant.SC + g_own_base = sG_tma[own_base, 0, stage_idx] + sG_factor_a[token_in_chunk, stage_idx] = cute.exp2( + g_i - g_own_base, + fastmath=self.use_fast_math, + ) + sG_main_q[token_in_chunk, stage_idx] = cute.exp2( + g_i, + fastmath=self.use_fast_math, + ) + sG_main_k[token_in_chunk, stage_idx] = cute.exp2( + g_last_scalar - g_i, + fastmath=self.use_fast_math, + ) + else: + sG_factor_a[token_in_chunk, stage_idx] = cutlass.Float32(0.0) + sG_main_q[token_in_chunk, stage_idx] = cutlass.Float32(0.0) + sG_main_k[token_in_chunk, stage_idx] = cutlass.Float32(0.0) + + num_subchunks = Constant.C // Constant.SC + if lane_idx < num_subchunks * num_subchunks: + query_subchunk = lane_idx // num_subchunks + key_subchunk = lane_idx % num_subchunks + query_base = query_subchunk * Constant.SC + key_base = key_subchunk * Constant.SC + if ( + key_subchunk <= query_subchunk + and query_base < valid_len_g + and key_base < valid_len_g + ): + sG_factor_base[ + query_subchunk, + key_subchunk, + stage_idx, + ] = cute.exp2( + sG_tma[query_base, 0, stage_idx] + - sG_tma[key_base, 0, stage_idx], + fastmath=self.use_fast_math, + ) + else: + sG_factor_base[ + query_subchunk, + key_subchunk, + stage_idx, + ] = cutlass.Float32(0.0) + cute.arch.sync_warp() + + for lane_slot in cutlass.range_constexpr(2): + token_in_chunk = lane_idx + lane_slot * self.threads_per_warp + if token_in_chunk < valid_len_g: + key_subchunk = token_in_chunk // Constant.SC + inv_a = cutlass.Float32(1.0) / sG_factor_a[ + token_in_chunk, + stage_idx, + ] + for query_subchunk in cutlass.range_constexpr(num_subchunks): + if key_subchunk <= query_subchunk: + sG_factor_b[ + query_subchunk, + token_in_chunk, + stage_idx, + ] = ( + sG_factor_base[ + query_subchunk, + key_subchunk, + stage_idx, + ] + * inv_a + ) + else: + sG_factor_b[ + query_subchunk, + token_in_chunk, + stage_idx, + ] = cutlass.Float32(0.0) + else: + for query_subchunk in cutlass.range_constexpr(num_subchunks): + sG_factor_b[ + query_subchunk, + token_in_chunk, + stage_idx, + ] = cutlass.Float32(0.0) + + @cute.jit + def apply_scalar_gate_to_subchunk_acc( + self, + qk_acc: cute.Tensor, + kk_acc: cute.Tensor, + acc_coords: cute.Tensor, + scalar_g: cute.Tensor, + row_subchunk: cutlass.Constexpr[int], + col_subchunk: cutlass.Constexpr[int], + ): + """Apply exp2(g[row]-g[col]) to FP32 QK/KK accumulators. + + The vector-gate implementation factors this term into separately + scaled BF16 Q and K operands. For a native scalar gate, applying the + exact token-pair factor to the FP32 accumulator is both cheaper and + avoids relying on operand-fragment coordinates to recover token ids. + """ + row_base = row_subchunk * Constant.SC + col_base = col_subchunk * Constant.SC + for i in cutlass.range_constexpr(cute.size(acc_coords)): + row, col = acc_coords[i] + gate = cute.exp2( + scalar_g[row_base + row] - scalar_g[col_base + col], + fastmath=self.use_fast_math, + ) + qk_acc[i] *= gate + kk_acc[i] *= gate + @cute.jit def mma_sync_partition_c( self, tiled_mma: cute.atom.TiledMma, tile_shape_mnk: cute.Shape, zero_fill: cutlass.Constexpr[bool] = True @@ -5605,31 +6188,42 @@ def s2r_compute_subchunk_operand_A( q_k_tiled_copy: cute.atom.TiledCopy, q_k_thr_copy: cute.atom.ThrCopy, tv_layout_mma_A: cute.Layout, + scalar_coords: cute.Tensor, layout_g_first: cute.Layout, # for make g_first tensor sG_slice: cute.Tensor, + sG_scalar: cute.Tensor, + sG_factor_a: cute.Tensor, sQ_slice: cute.Tensor, sK_slice: cute.Tensor, ): - # S2R g, g_first - sG = sG_slice[None, None, subchunk_idx, (0, nk)] - tQKsG = g_thr_copy.partition_S(sG) tQKrG = cute.make_fragment_like(tv_layout_mma_A, dtype=self.g_dtype) - tQKrG_cv = g_thr_copy.retile(tQKrG) - cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) - - # TODO: do register shuffle g to get g_first, reduce smem load - sG_first = cute.make_tensor(sG.iterator, layout=layout_g_first) - tQKsGfirst = g_thr_copy.partition_S(sG_first) - tQKrGfirst = cute.make_fragment_like(tv_layout_mma_A, dtype=self.g_dtype) - tQKrGfirst_cv = g_thr_copy.retile(tQKrGfirst) - cute.copy(g_tiled_copy, tQKsGfirst, tQKrGfirst_cv) + if cutlass.const_expr(self.scalar_gate): + # Coordinates come from the exact MMA-A partition, matching the + # flattened register fragment one-for-one. + for i in cutlass.range_constexpr(cute.size(scalar_coords)): + index_q, _ = scalar_coords[i] + tQKrG[i] = sG_factor_a[subchunk_idx * Constant.SC + index_q] + else: + # S2R g, g_first + sG = sG_slice[None, None, subchunk_idx, (0, nk)] + tQKsG = g_thr_copy.partition_S(sG) + tQKrG_cv = g_thr_copy.retile(tQKrG) + cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) + + # TODO: do register shuffle g to get g_first, reduce smem load + sG_first = cute.make_tensor(sG.iterator, layout=layout_g_first) + tQKsGfirst = g_thr_copy.partition_S(sG_first) + tQKrGfirst = cute.make_fragment_like(tv_layout_mma_A, dtype=self.g_dtype) + tQKrGfirst_cv = g_thr_copy.retile(tQKrGfirst) + cute.copy(g_tiled_copy, tQKsGfirst, tQKrGfirst_cv) + + # gqn = exp2(g - g_first[None, :]), reuse g + g_val = tQKrG.load() + g_first_val = tQKrGfirst.load() + g_val = cute.exp2(g_val - g_first_val, fastmath=self.use_fast_math) + tQKrG.store(g_val) - # gqn = exp2(g - g_first[None, :]), reuse g g_val = tQKrG.load() - g_first_val = tQKrGfirst.load() - g_val = cute.exp2(g_val - g_first_val, fastmath=self.use_fast_math) - tQKrG.store(g_val) - # S2R q, k sQ = sQ_slice[None, None, subchunk_idx, (0, nk)] sK = sK_slice[None, None, subchunk_idx, (0, nk)] @@ -5664,23 +6258,36 @@ def s2r_compute_subchunk_operand_B( kt_tiled_copy: cute.atom.TiledCopy, kt_thr_copy: cute.atom.ThrCopy, tv_layout_mma_B: cute.Layout, + scalar_coords: cute.Tensor, sG_slice: cute.Tensor, + sG_scalar: cute.Tensor, + sG_factor_b: cute.Tensor, + query_subchunk_idx: cutlass.Constexpr[int], sK_slice: cute.Tensor, rG_first: cute.Tensor, ): - # S2R g - sG = sG_slice[None, None, subchunk_idx, (0, nk)] - tQKsG = g_thr_copy.partition_S(sG) tQKrG = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) - tQKrG_cv = g_thr_copy.retile(tQKrG) - cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) + if cutlass.const_expr(self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(scalar_coords)): + index_k, _ = scalar_coords[i] + tQKrG[i] = sG_factor_b[ + query_subchunk_idx, + subchunk_idx * Constant.SC + index_k, + ] + else: + # S2R g + sG = sG_slice[None, None, subchunk_idx, (0, nk)] + tQKsG = g_thr_copy.partition_S(sG) + tQKrG_cv = g_thr_copy.retile(tQKrG) + cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) + + # compute gktn = exp2(g_first - g), reuse g + g_val = tQKrG.load() + g_first_val = rG_first.load() + g_val = cute.exp2(g_first_val - g_val, fastmath=self.use_fast_math) + tQKrG.store(g_val) - # compute gktn = exp2(g_first - g), reuse g g_val = tQKrG.load() - g_first_val = rG_first.load() - g_val = cute.exp2(g_first_val - g_val, fastmath=self.use_fast_math) - tQKrG.store(g_val) - # S2R k sK = sK_slice[None, None, subchunk_idx, (0, nk)] tQKrKt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.k_dtype) @@ -5766,7 +6373,7 @@ def make_s2r_partitions_v( # num_bits_per_copy=dtype.width * 8, num_bits_per_copy=dtype.width * 1, ) - num_elements_per_thread = Constant.C + num_elements_per_thread = Constant.C // self.num_value_tiles num_threads_per_row = shape_x[1] // num_elements_per_thread # NOTE: Assume 128 cuda core threads num_threads_per_col = 128 // num_threads_per_row diff --git a/cula/ops/kda/experimental/sm100_fused/wrapper.py b/cula/ops/kda/experimental/sm100_fused/wrapper.py index 1896d2c4..91383cef 100644 --- a/cula/ops/kda/experimental/sm100_fused/wrapper.py +++ b/cula/ops/kda/experimental/sm100_fused/wrapper.py @@ -12,22 +12,80 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""[experimental] Unwired SM100 fully fused KDA prefill dead-path wrapper;arch=SM100""" - +import pathlib +import sys +import types import warnings +import torch +import torch.nn.functional as F + +# This module now lives under the experimental backend package; imports must +# resolve through the normal cuLA package rather than the old flat layout. + import cutlass import cutlass.cute as cute import cutlass.torch as cutlass_torch -import torch from cutlass.cute.runtime import from_dlpack -from fla.modules.l2norm import l2norm_fwd # from fla.ops.kda.chunk_inter import chunk_kda_bwd_dqkwg -from fla.ops.kda.gate import kda_gate_fwd -from fla.ops.utils import chunk_local_cumsum -from fla.ops.utils.constant import RCP_LN2 -from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +try: + from fla.modules.l2norm import l2norm_fwd + from fla.ops.kda.gate import kda_gate_fwd + from fla.ops.utils import chunk_local_cumsum + from fla.ops.utils.constant import RCP_LN2 + from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + from cula.ops.l2norm_triton import l2norm_fwd, l2norm_qk_fwd + +except ImportError: + RCP_LN2 = 1.4426950408889634 + + def input_guard(fn): + return fn + + def autocast_custom_fwd(fn): + return fn + + def autocast_custom_bwd(fn): + return fn + + def l2norm_fwd(x: torch.Tensor): + rstd = torch.rsqrt(x.float().square().sum(dim=-1, keepdim=True).clamp_min(1.0e-12)) + return (x.float() * rstd).to(x.dtype), rstd + + def l2norm_qk_fwd(q: torch.Tensor, k: torch.Tensor): + q_out, q_rstd = l2norm_fwd(q) + k_out, k_rstd = l2norm_fwd(k) + return q_out, k_out, q_rstd, k_rstd + + def kda_gate_fwd(*args, **kwargs): + raise ImportError("fla is required for use_gate_in_kernel=True in blackwell_fused_fwd") + + def chunk_local_cumsum( + g: torch.Tensor, + chunk_size: int, + scale: float = 1.0, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + if chunk_indices is not None: + raise ImportError("fla is required for chunk_indices support in blackwell_fused_fwd") + if cu_seqlens is not None: + if g.shape[0] != 1: + raise ValueError("cu_seqlens mode expects flattened g with batch size 1") + out = torch.empty_like(g.float()) + for seq_idx in range(cu_seqlens.numel() - 1): + start = int(cu_seqlens[seq_idx].item()) + end = int(cu_seqlens[seq_idx + 1].item()) + for chunk_start in range(start, end, chunk_size): + chunk_end = min(chunk_start + chunk_size, end) + out[:, chunk_start:chunk_end] = g[:, chunk_start:chunk_end].float().cumsum(dim=1) * scale + return out + chunks = [] + for chunk_start in range(0, g.shape[1], chunk_size): + chunk = g[:, chunk_start : chunk_start + chunk_size].float().cumsum(dim=1) * scale + chunks.append(chunk) + return torch.cat(chunks, dim=1).contiguous() from cula.ops.kda.experimental.sm100_fused.kda_fully_fused_wip import KDAChunkwise from cula.utils import USE_FAST_MATH, assert_blackwell @@ -41,7 +99,7 @@ _dummy_cache = {} -class BlackwellFusedKDAFunction(torch.autograd.Function): +class ChunkKDAFunction(torch.autograd.Function): @staticmethod @input_guard @autocast_custom_fwd @@ -61,21 +119,35 @@ def forward( use_gate_in_kernel: bool = False, safe_gate: bool = False, lower_bound: float | None = None, + g_is_cumsum: bool = False, + scalar_gate: bool = False, cu_seqlens: torch.IntTensor | None = None, chunk_indices: torch.IntTensor | None = None, ): chunk_size = 64 - assert q.shape[-2] == v.shape[-2] == k.shape[-2], "Number of heads must be the same for q, k, v." + assert q.shape == k.shape, "q and k must have the same shape." global compiled_kernel_cache B, S, H, D = q.shape + HV = v.shape[-2] + assert HV % H == 0, f"HV ({HV}) must be a multiple of H ({H})." is_varlen = cu_seqlens is not None if is_varlen: assert B == 1, "For varlen, batch size must be 1. Flatten variable-length inputs first." num_seqs = cu_seqlens.shape[0] - 1 else: num_seqs = B + split_value_tiles = ( + scalar_gate + and safe_gate + and not is_varlen + and B == 1 + and H == 16 + and HV == 48 + and D == 128 + and v.shape[-1] == 128 + ) g_org = None if use_gate_in_kernel: @@ -103,14 +175,22 @@ def forward( A_log=A_log, dt_bias=dt_bias, ) - if not (safe_gate and use_gate_in_kernel): + fuse_scalar_cumsum = ( + scalar_gate + and not g_is_cumsum + and not use_gate_in_kernel + ) + if ( + not g_is_cumsum + and not (safe_gate and use_gate_in_kernel) + and not fuse_scalar_cumsum + ): g = chunk_local_cumsum( g=g, chunk_size=chunk_size, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices ) q_rstd, k_rstd = None, None if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) + q, k, q_rstd, k_rstd = l2norm_qk_fwd(q, k) q_cute = from_dlpack(q.detach()) k_cute = from_dlpack(k.detach()) @@ -118,13 +198,30 @@ def forward( g_cute = from_dlpack(g.detach()) beta_cute = from_dlpack(beta.detach()) - o = torch.empty_like(q) + o = torch.empty_like(v) o_cute = from_dlpack(o.detach()) stream = cutlass_torch.default_stream() has_initial_state = initial_state is not None - cache_key = (has_initial_state, output_final_state, safe_gate, is_varlen, scale, chunk_size, D, USE_FAST_MATH) + # H/HV affect the compiled tensor layouts and GVA head mapping. Without + # them, a process exercising multiple Qwen model/TP shapes can reuse a + # kernel compiled for a different head configuration. + cache_key = ( + has_initial_state, + output_final_state, + safe_gate, + scalar_gate, + fuse_scalar_cumsum, + split_value_tiles, + is_varlen, + scale, + chunk_size, + H, + HV, + D, + USE_FAST_MATH, + ) if is_varlen: cu_seqlens_i32 = cu_seqlens.to(torch.int32).contiguous() @@ -177,13 +274,15 @@ def forward( initial_state_cute = _dummy_cache[q.device]["state_cute"] if output_final_state: - final_state_f32 = torch.zeros(num_seqs, H, D, D, dtype=torch.float32, device=q.device) + # The CuTe kernel overwrites every final-state element. Avoid a + # redundant ~3 MiB memset for Qwen3.5-27B (HV=48) on every call. + final_state_f32 = torch.empty(num_seqs, HV, D, D, dtype=torch.float32, device=q.device) final_state_cute = from_dlpack(final_state_f32.detach()) else: final_state_f32 = None final_state_cute = _dummy_cache[q.device]["state_cute"] - problem_size = (num_seqs, S, H, D) + problem_size = (num_seqs, S, H, HV, D) if cache_key in compiled_kernel_cache: compiled_kernel = compiled_kernel_cache[cache_key] @@ -195,6 +294,9 @@ def forward( io_dtype=cutlass.BFloat16, scale=scale, safe_gate=safe_gate, + scalar_gate=scalar_gate, + fuse_scalar_cumsum=fuse_scalar_cumsum, + split_value_tiles=split_value_tiles, has_initial_state=has_initial_state, output_final_state=output_final_state, is_varlen=is_varlen, @@ -264,11 +366,14 @@ def flash_kda_prefill( use_gate_in_kernel: bool = False, safe_gate: bool = False, lower_bound: float | None = None, + g_is_cumsum: bool = False, + scalar_gate: bool = False, cu_seqlens: torch.IntTensor | None = None, chunk_indices: torch.IntTensor | None = None, **kwargs, ): assert_blackwell() + assert cu_seqlens is None or q.shape[0] == 1, "For varlen, batch size must be 1. Flatten sequences first." if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -302,12 +407,13 @@ def flash_kda_prefill( assert HV % H == 0, ( f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by num_qk_heads (H={H}), but got HV % H = {HV % H}" ) - assert g.shape == (B, T, HV, K), f"g must have shape [B, T, HV, K]={[B, T, HV, K]}, got {list(g.shape)}" + expected_g_shape = (B, T, HV) if scalar_gate else (B, T, HV, K) + assert g.shape == expected_g_shape, f"g must have shape {expected_g_shape}, got {list(g.shape)}" assert beta.shape == (B, T, HV), f"beta must have shape [B, T, HV]={[B, T, HV]}, got {list(beta.shape)}" if scale is None: scale = k.shape[-1] ** -0.5 - o, final_state = BlackwellFusedKDAFunction.apply( + forward_args = ( q, k, v, @@ -322,7 +428,19 @@ def flash_kda_prefill( use_gate_in_kernel, safe_gate, lower_bound, + g_is_cumsum, + scalar_gate, cu_seqlens, chunk_indices, ) + if torch.is_grad_enabled() and any( + tensor is not None and tensor.requires_grad + for tensor in (q, k, v, g, beta, initial_state) + ): + o, final_state = ChunkKDAFunction.apply(*forward_args) + else: + # The op has no backward implementation. SGLang always reaches this + # inference branch, so avoid the measurable autograd.Function.apply + # dispatch cost while retaining AMP/input-guard behavior. + o, final_state = ChunkKDAFunction.forward(types.SimpleNamespace(), *forward_args) return o, final_state diff --git a/cula/ops/qwen35_conv1d_decode.py b/cula/ops/qwen35_conv1d_decode.py new file mode 100644 index 00000000..bd0af8f2 --- /dev/null +++ b/cula/ops/qwen35_conv1d_decode.py @@ -0,0 +1,119 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Qwen3.5 single-token conv-state update wrapper.""" + +from __future__ import annotations + +import torch + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def qwen35_conv1d_decode_reference( + x_t: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pure torch reference for Qwen3.5 single-token depthwise conv decode.""" + if weight.ndim == 3: + weight = weight.squeeze(1) + + state_tail = conv_state[..., 1:].to(torch.float32) + x_last = x_t.unsqueeze(-1).to(torch.float32) + window = torch.cat([state_tail, x_last], dim=-1) + conv = (window * weight.to(torch.float32).unsqueeze(0)).sum(dim=-1) + y = torch.nn.functional.silu(conv).to(dtype=x_t.dtype) + + conv_state_out = conv_state.clone() + conv_state_out[..., 0] = conv_state[..., 1] + conv_state_out[..., 1] = conv_state[..., 2] + conv_state_out[..., 2] = conv_state[..., 3] + conv_state_out[..., 3] = x_t + return y, conv_state_out + + +def qwen35_conv1d_decode_update( + x_t: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + *, + activation: str = "silu", + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-token depthwise causal conv1d update. + + Expected shapes: + - x_t: [B, C] + - conv_state: [B, C, 4] + - weight: [C, 1, 4] or [C, 4] + """ + + if activation != "silu": + raise ValueError(f"Only silu activation is currently supported, got {activation}") + if x_t.ndim != 2: + raise ValueError(f"x_t must be 2D [batch, channels], got {tuple(x_t.shape)}") + if conv_state.ndim != 3: + raise ValueError(f"conv_state must be 3D [batch, channels, kernel], got {tuple(conv_state.shape)}") + if conv_state.shape[:2] != x_t.shape: + raise ValueError(f"conv_state batch/channel dims must match x_t, got x_t={tuple(x_t.shape)} conv_state={tuple(conv_state.shape)}") + kernel_size = conv_state.shape[-1] + if kernel_size != 4: + raise ValueError(f"Expected kernel_size=4 for Qwen3.5, got {kernel_size}") + + if weight.ndim == 3: + if weight.shape[1] != 1 or weight.shape[2] != kernel_size: + raise ValueError(f"weight must be [channels,1,{kernel_size}], got {tuple(weight.shape)}") + weight_2d = weight.squeeze(1) + elif weight.ndim == 2: + if weight.shape[1] != kernel_size: + raise ValueError(f"weight must be [channels,{kernel_size}], got {tuple(weight.shape)}") + weight_2d = weight + else: + raise ValueError(f"weight must be 2D or 3D, got {tuple(weight.shape)}") + + if weight_2d.shape[0] != x_t.shape[1]: + raise ValueError(f"weight channels must match x_t channels, got weight={tuple(weight_2d.shape)} x_t={tuple(x_t.shape)}") + + x_t = x_t.contiguous() + conv_state = conv_state.contiguous() + weight_2d = weight_2d.contiguous() + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_conv1d_decode") + and x_t.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_conv1d_decode is not available.") + + if use_cudac: + mixed_qkv_3d = x_t.unsqueeze(1).contiguous() + out_3d = torch.empty_like(mixed_qkv_3d) + conv_state_out = conv_state.clone() + cula_cuda.qwen35_conv1d_decode( + mixed_qkv_3d, + conv_state_out, + weight_2d, + out_3d, + ) + return out_3d.squeeze(1), conv_state_out + + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + return qwen35_conv1d_decode_reference(x_t, conv_state, weight_2d) diff --git a/cula/ops/qwen35_conv1d_prefill.py b/cula/ops/qwen35_conv1d_prefill.py new file mode 100644 index 00000000..3a32b1e0 --- /dev/null +++ b/cula/ops/qwen35_conv1d_prefill.py @@ -0,0 +1,99 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Qwen3.5 depthwise causal conv1d prefill wrapper.""" + +from __future__ import annotations + +import torch + + +def qwen35_conv1d_prefill( + x: torch.Tensor, + weight: torch.Tensor, + *, + activation: str = "silu", + cu_seqlens: torch.Tensor | None = None, + output_final_state: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Depthwise causal conv1d over a full sequence. + + Expected shapes: + - x: [B, C, S] or flattened [T, C] + - weight: [C, 1, 4] or [C, 4] + """ + + if activation != "silu": + raise ValueError(f"Unsupported activation={activation}") + if weight.ndim == 3: + if weight.shape[1] != 1: + raise ValueError(f"weight must be [C,1,K] or [C,K], got {tuple(weight.shape)}") + weight_2d = weight[:, 0, :] + elif weight.ndim == 2: + weight_2d = weight + else: + raise ValueError(f"weight must be [C,1,K] or [C,K], got {tuple(weight.shape)}") + + kernel_size = weight_2d.shape[1] + + def _conv_one(seq: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # seq: [S, C] + if seq.ndim != 2 or seq.shape[1] != weight_2d.shape[0]: + raise ValueError(f"sequence must be [S,C={weight_2d.shape[0]}], got {tuple(seq.shape)}") + seq_f = seq.float() + weight_f = weight_2d.float() + out = torch.empty_like(seq) + for t in range(seq.shape[0]): + acc = torch.zeros(seq.shape[1], device=seq.device, dtype=torch.float32) + for kk in range(kernel_size): + src_t = t - (kernel_size - 1 - kk) + if src_t >= 0: + acc = acc + seq_f[src_t] * weight_f[:, kk] + out[t] = torch.nn.functional.silu(acc).to(seq.dtype) + + state = torch.zeros(seq.shape[1], kernel_size, device=seq.device, dtype=seq.dtype) + take = min(kernel_size, seq.shape[0]) + if take > 0: + state[:, kernel_size - take :] = seq[-take:].transpose(0, 1) + return out, state + + if x.ndim == 3: + # Public op shape follows the Qwen conv convention [B, C, S]. + if x.shape[1] != weight_2d.shape[0]: + raise ValueError(f"x channel dim must match weight, got x={tuple(x.shape)} weight={tuple(weight_2d.shape)}") + y = torch.empty_like(x) + states = torch.empty(x.shape[0], x.shape[1], kernel_size, device=x.device, dtype=x.dtype) + for bidx in range(x.shape[0]): + y_b, state_b = _conv_one(x[bidx].transpose(0, 1).contiguous()) + y[bidx] = y_b.transpose(0, 1).contiguous() + states[bidx] = state_b + return (y, states) if output_final_state else y + + if x.ndim == 2: + if cu_seqlens is None: + y, state = _conv_one(x) + return (y, state.unsqueeze(0)) if output_final_state else y + if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + y = torch.empty_like(x) + states = torch.empty(cu_seqlens.numel() - 1, x.shape[1], kernel_size, device=x.device, dtype=x.dtype) + for sidx in range(cu_seqlens.numel() - 1): + start = int(cu_seqlens[sidx].item()) + end = int(cu_seqlens[sidx + 1].item()) + y_s, state_s = _conv_one(x[start:end]) + y[start:end] = y_s + states[sidx] = state_s + return (y, states) if output_final_state else y + + raise ValueError(f"x must be [B,C,S] or [T,C], got {tuple(x.shape)}") diff --git a/cula/ops/qwen35_fused_kda_prefill.py b/cula/ops/qwen35_fused_kda_prefill.py new file mode 100644 index 00000000..39f32024 --- /dev/null +++ b/cula/ops/qwen35_fused_kda_prefill.py @@ -0,0 +1,152 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Qwen3.5 adapter for the native-GVA fully-fused CuTe prefill core.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def _resolve_fused_kda_prefill(device: torch.device | str | int | None = None): + try: + from cula.utils import get_kda_fused_fwd + except Exception as exc: # pragma: no cover - depends on optional runtime deps + raise RuntimeError(f"Cannot import fused KDA selector: {exc}") from exc + + try: + return get_kda_fused_fwd(device) + except Exception as exc: + raise RuntimeError(f"Cannot resolve fused KDA prefill for device={device}: {exc}") from exc + + +def has_qwen35_fused_kda_prefill(device: torch.device | str | int | None = None) -> bool: + try: + _resolve_fused_kda_prefill(device) + except Exception: + return False + return True + + +def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> tuple[int, int, int, int, torch.Tensor, torch.Tensor]: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError( + f"q/k/v must be 4D [B,T,H,D], got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}" + ) + if q.shape != k.shape: + raise ValueError(f"q and k must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)}") + B, T, H, K = q.shape + if v.shape[:2] != (B, T): + raise ValueError(f"v must match q/k batch and token dimensions, got q={tuple(q.shape)} v={tuple(v.shape)}") + HV, V = v.shape[2:] + if K != 128 or V != 128: + raise ValueError(f"Qwen3.5 fused prefill expects K=V=128, got K={K} V={V}") + if HV % H != 0: + raise ValueError(f"Qwen3.5 GVA expects HV to be divisible by H, got H={H} HV={HV}") + if a.ndim == 2: + a = a.unsqueeze(0) + if b.ndim == 2: + b = b.unsqueeze(0) + if a.shape != (B, T, HV) or b.shape != (B, T, HV): + raise ValueError(f"a/b must be [B,T,HV], got a={tuple(a.shape)} b={tuple(b.shape)} expected={(B, T, HV)}") + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + if cu_seqlens is not None: + if B != 1: + raise ValueError("cu_seqlens mode expects flattened q/k/v with batch size 1") + if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + if initial_state is not None and initial_state.shape != (state_count, HV, K, V): + raise ValueError(f"initial_state must be [{state_count},{HV},{K},{V}], got {tuple(initial_state.shape)}") + return B, T, HV, K, a, b + + +def qwen35_fused_kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + output_final_state: bool = True, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run Qwen3.5 scalar-gated KDA prefill through the fully-fused CuTe core. + + Qwen3.5 uses native grouped value attention: Q/K have H heads while V and + the scalar GDN gate have HV heads (globally H=16 and HV=48). The scalar + gate is passed to the CuTe specialization without materializing a D=128 + broadcast; Q/K likewise remain in their native, non-repeated layout. + + State is exposed in Qwen layout [N, HV, K, V]. The fused core consumes the + transposed initial state and returns final state in Qwen layout. + """ + + if not q.is_cuda: + raise RuntimeError("qwen35_fused_kda_prefill requires CUDA tensors.") + B, T, HV, K, a, b = _validate_inputs(q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens) + + fused_kda_prefill = _resolve_fused_kda_prefill(q.device) + + log_gate_scalar = -torch.exp(A_log.float()).view(1, 1, HV, 1) * F.softplus( + a.float().unsqueeze(-1) + dt_bias.float().view(1, 1, HV, 1) + ) + log_gate = log_gate_scalar.squeeze(-1).contiguous() + beta = torch.sigmoid(b.float()).contiguous() + + # A single [0, T] sequence is an ordinary equal-length prefill, not a + # variable-length batch. Keep the fast non-varlen SM100 launch in this + # common Qwen inference case; the varlen path carries extra indirection + # and workspace overhead. + kernel_cu_seqlens = cu_seqlens + if cu_seqlens is not None and q.shape[0] == 1 and cu_seqlens.numel() == 2: + kernel_cu_seqlens = None + + initial_state_vk = None + if initial_state is not None: + initial_state_vk = initial_state.float().transpose(-1, -2).contiguous() + + out, final_state_vk = fused_kda_prefill( + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + g=log_gate, + beta=beta, + scale=K**-0.5, + initial_state=initial_state_vk, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=False, + cu_seqlens=kernel_cu_seqlens, + safe_gate=True, + lower_bound=-5.0, + scalar_gate=True, + ) + final_state = None if final_state_vk is None else final_state_vk.contiguous() + return out, final_state diff --git a/cula/ops/qwen35_layout_decode.py b/cula/ops/qwen35_layout_decode.py new file mode 100644 index 00000000..eec5f975 --- /dev/null +++ b/cula/ops/qwen35_layout_decode.py @@ -0,0 +1,112 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Qwen3.5 layout decode wrapper.""" + +from __future__ import annotations + +import torch + +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig, infer_local_config + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def qwen35_layout_decode_reference( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + local_key_dim, _, local_num_k_heads = infer_local_config( + mixed_qkv_conv.shape[1], + local_num_v_heads, + config=config, + ) + + q_end = local_key_dim + k_end = q_end + local_key_dim + q_flat = mixed_qkv_conv[:, :q_end] + k_flat = mixed_qkv_conv[:, q_end:k_end] + v_flat = mixed_qkv_conv[:, k_end:] + + q = q_flat.view(tokens, local_num_k_heads, config.head_k_dim) + k = k_flat.view(tokens, local_num_k_heads, config.head_k_dim) + v = v_flat.view(tokens, local_num_v_heads, config.head_v_dim) + + repeat_factor = local_num_v_heads // local_num_k_heads + q_rep = q.repeat_interleave(repeat_factor, dim=1).contiguous() + k_rep = k.repeat_interleave(repeat_factor, dim=1).contiguous() + return q_rep, k_rep, v.contiguous(), a.contiguous(), b.contiguous() + + +def qwen35_layout_decode( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_layout_decode") + and mixed_qkv_conv.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_decode is not available.") + + if use_cudac: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + infer_local_config(mixed_qkv_conv.shape[1], local_num_v_heads, config=config) + q_rep = torch.empty( + tokens, + local_num_v_heads, + config.head_k_dim, + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + k_rep = torch.empty_like(q_rep) + v = torch.empty( + tokens, + local_num_v_heads, + config.head_v_dim, + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + a_kernel = torch.empty_like(a) + b_kernel = torch.empty_like(b) + cula_cuda.qwen35_layout_decode( + mixed_qkv_conv.contiguous(), + a.contiguous(), + b.contiguous(), + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + ) + return q_rep, k_rep, v, a_kernel, b_kernel + + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + return qwen35_layout_decode_reference(mixed_qkv_conv, a, b, config=config) diff --git a/cula/ops/qwen35_layout_prefill.py b/cula/ops/qwen35_layout_prefill.py new file mode 100644 index 00000000..4a6a2867 --- /dev/null +++ b/cula/ops/qwen35_layout_prefill.py @@ -0,0 +1,96 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Qwen3.5 layout prefill wrapper.""" + +from __future__ import annotations + +import torch + +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig, infer_local_config + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def qwen35_layout_prefill_reference( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + local_key_dim, _, local_num_k_heads = infer_local_config( + mixed_qkv_conv.shape[1], + local_num_v_heads, + config=config, + ) + + q_end = local_key_dim + k_end = q_end + local_key_dim + q = mixed_qkv_conv[:, :q_end].view(tokens, local_num_k_heads, config.head_k_dim) + k = mixed_qkv_conv[:, q_end:k_end].view(tokens, local_num_k_heads, config.head_k_dim) + v = mixed_qkv_conv[:, k_end:].view(tokens, local_num_v_heads, config.head_v_dim) + + repeat_factor = local_num_v_heads // local_num_k_heads + q_rep = q.repeat_interleave(repeat_factor, dim=1).contiguous() + k_rep = k.repeat_interleave(repeat_factor, dim=1).contiguous() + return q_rep, k_rep, v.contiguous(), a.contiguous(), b.contiguous() + + +def qwen35_layout_prefill( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_layout_prefill") + and mixed_qkv_conv.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_prefill is not available.") + + if use_cudac: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + infer_local_config(mixed_qkv_conv.shape[1], local_num_v_heads, config=config) + q_rep = torch.empty(tokens, local_num_v_heads, config.head_k_dim, device=mixed_qkv_conv.device, dtype=mixed_qkv_conv.dtype) + k_rep = torch.empty_like(q_rep) + v = torch.empty(tokens, local_num_v_heads, config.head_v_dim, device=mixed_qkv_conv.device, dtype=mixed_qkv_conv.dtype) + a_kernel = torch.empty_like(a) + b_kernel = torch.empty_like(b) + cula_cuda.qwen35_layout_prefill( + mixed_qkv_conv.contiguous(), + a.contiguous(), + b.contiguous(), + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + ) + return q_rep, k_rep, v, a_kernel, b_kernel + + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + return qwen35_layout_prefill_reference(mixed_qkv_conv, a, b, config=config) diff --git a/cula/ops/qwen35_scalar_kda_decode.py b/cula/ops/qwen35_scalar_kda_decode.py new file mode 100644 index 00000000..45d2c6a6 --- /dev/null +++ b/cula/ops/qwen35_scalar_kda_decode.py @@ -0,0 +1,227 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""CuTe DSL placeholder for Qwen3.5 scalar-gated KDA decode.""" + +from __future__ import annotations + +import torch + +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def has_qwen35_layout_scalar_kda_decode_cudac() -> bool: + return cula_cuda is not None and hasattr(cula_cuda, "qwen35_layout_scalar_kda_decode") + + +def _validate_cudac_state_indices(state_indices: torch.Tensor, *, rows: int, pool_size: int) -> None: + if state_indices.ndim != 1 or state_indices.numel() != rows: + raise ValueError(f"state_indices must be 1D with {rows} entries, got {tuple(state_indices.shape)}") + if rows == 0: + return + min_idx = int(state_indices.min().item()) + max_idx = int(state_indices.max().item()) + if min_idx < 0 or max_idx >= pool_size: + raise ValueError(f"state_indices must be in [0, {pool_size}), got min={min_idx} max={max_idx}") + if torch.unique(state_indices).numel() != rows: + raise ValueError( + "backend='cudac' requires unique state_indices within one decode launch; " + "duplicate rows need a sequential decode path." + ) + + +def qwen35_scalar_kda_decode( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + *, + state_indices: torch.Tensor | None = None, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-token scalar-gated delta-rule decode for Qwen3.5.""" + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError(f"q/k/v must be 4D, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + if q.shape != k.shape: + raise ValueError(f"q and k must have the same shape, got q={tuple(q.shape)} vs k={tuple(k.shape)}") + if q.shape[1] != 1 or v.shape[1] != 1: + raise ValueError(f"Decode expects single-token sequence dim, got q={tuple(q.shape)} v={tuple(v.shape)}") + + N, _, HV, K = q.shape + if a.ndim == 2: + a = a.unsqueeze(1) + if b.ndim == 2: + b = b.unsqueeze(1) + if a.shape != (N, 1, HV) or b.shape != (N, 1, HV): + raise ValueError(f"a/b must be [N,1,HV], got a={tuple(a.shape)} b={tuple(b.shape)}") + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + + state_indices = ( + torch.arange(N, device=q.device, dtype=torch.int32) + if state_indices is None + else state_indices.to(device=q.device, dtype=torch.int32) + ) + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_scalar_kda_decode") + and q.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_decode is not available.") + + if use_cudac: + _validate_cudac_state_indices(state_indices, rows=N, pool_size=recurrent_state.shape[0]) + q_rep = q.squeeze(1).contiguous() + k_rep = k.squeeze(1).contiguous() + v_rep = v.squeeze(1).contiguous() + a_kernel = a.squeeze(1).contiguous() + b_kernel = b.squeeze(1).contiguous() + out = torch.empty_like(v_rep) + recurrent_state_out = recurrent_state.clone() + cula_cuda.qwen35_scalar_kda_decode( + q_rep, + k_rep, + v_rep, + a_kernel, + b_kernel, + A_log.contiguous(), + dt_bias.contiguous(), + recurrent_state_out, + state_indices, + out, + ) + return out.unsqueeze(1), recurrent_state_out + + if backend not in ("auto", "generic_kda"): + raise ValueError(f"Unsupported backend={backend}") + + from cula.ops.kda_decode import kda_decode + + a_expanded = a.unsqueeze(-1).expand(N, 1, HV, K) + dt_bias_expanded = dt_bias[:, None].expand(HV, K).contiguous() + o = kda_decode( + A_log=A_log.contiguous(), + dt_bias=dt_bias_expanded, + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + a=a_expanded.contiguous(), + b=b.contiguous(), + initial_state_source=recurrent_state, + initial_state_indices=state_indices, + scale=K**-0.5, + use_qk_l2norm_in_kernel=True, + state_layout="kv", + ) + return o, recurrent_state + + +def qwen35_layout_scalar_kda_decode( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + *, + state_indices: torch.Tensor | None = None, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused Qwen3.5 layout decode + scalar-gated KDA decode.""" + if mixed_qkv_conv.ndim != 2: + raise ValueError(f"mixed_qkv_conv must be 2D, got {tuple(mixed_qkv_conv.shape)}") + if a.ndim == 3: + if a.shape[1] != 1: + raise ValueError(f"a sequence dim must be 1 for decode, got {tuple(a.shape)}") + a = a.squeeze(1) + if b.ndim == 3: + if b.shape[1] != 1: + raise ValueError(f"b sequence dim must be 1 for decode, got {tuple(b.shape)}") + b = b.squeeze(1) + + N = mixed_qkv_conv.shape[0] + if a.ndim != 2 or b.ndim != 2 or a.shape != b.shape or a.shape[0] != N: + raise ValueError(f"a/b must be [N, HV], got a={tuple(a.shape)} b={tuple(b.shape)}") + HV = a.shape[1] + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + + state_indices = ( + torch.arange(N, device=mixed_qkv_conv.device, dtype=torch.int32) + if state_indices is None + else state_indices.to(device=mixed_qkv_conv.device, dtype=torch.int32) + ) + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_layout_scalar_kda_decode") + and mixed_qkv_conv.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_scalar_kda_decode is not available.") + + if use_cudac: + _validate_cudac_state_indices(state_indices, rows=N, pool_size=recurrent_state.shape[0]) + out = torch.empty( + N, + HV, + recurrent_state.shape[-1], + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + recurrent_state_out = recurrent_state.clone() + cula_cuda.qwen35_layout_scalar_kda_decode( + mixed_qkv_conv.contiguous(), + a.contiguous(), + b.contiguous(), + A_log.contiguous(), + dt_bias.contiguous(), + recurrent_state_out, + state_indices, + out, + ) + return out.unsqueeze(1), recurrent_state_out + + if backend not in ("auto", "generic_kda"): + raise ValueError(f"Unsupported backend={backend}") + + from cula.ops.qwen35_layout_decode import qwen35_layout_decode_reference + + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode_reference(mixed_qkv_conv, a, b, config=config) + return qwen35_scalar_kda_decode( + q=q_rep.unsqueeze(1).contiguous(), + k=k_rep.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a_kernel, + b=b_kernel, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="generic_kda", + ) diff --git a/cula/ops/qwen35_scalar_kda_prefill.py b/cula/ops/qwen35_scalar_kda_prefill.py new file mode 100644 index 00000000..53d5c486 --- /dev/null +++ b/cula/ops/qwen35_scalar_kda_prefill.py @@ -0,0 +1,228 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Qwen3.5 scalar-gated KDA prefill wrapper.""" + +from __future__ import annotations + +import torch + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def qwen35_scalar_kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Chunked scalar-gated delta-rule prefill for Qwen3.5.""" + + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError(f"q/k/v must be 4D, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + if q.shape != k.shape: + raise ValueError(f"q/k must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)}") + B, T, H, K = q.shape + HV = v.shape[2] + if K != 128 or v.shape[-1] != 128 or v.shape[:2] != q.shape[:2]: + raise ValueError(f"Qwen3.5 prefill expects q/k=[B,T,H,128], v=[B,T,HV,128], got q={tuple(q.shape)} v={tuple(v.shape)}") + if H <= 0 or HV <= 0 or HV % H: + raise ValueError(f"local V heads must be divisible by local Q/K heads, got H={H} HV={HV}") + if a.ndim == 2: + a = a.unsqueeze(0) + if b.ndim == 2: + b = b.unsqueeze(0) + if a.shape != (B, T, HV) or b.shape != (B, T, HV): + raise ValueError(f"a/b must be [B,T,HV], got a={tuple(a.shape)} b={tuple(b.shape)} expected={(B, T, HV)}") + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + if cu_seqlens is not None: + if B != 1: + raise ValueError("cu_seqlens mode expects flattened q/k/v with batch size 1") + if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + if initial_state is not None and initial_state.shape[1:] != (HV, K, K): + raise ValueError(f"initial_state must be [N,HV,128,128], got {tuple(initial_state.shape)}") + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_scalar_kda_prefill") + and q.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_prefill is not available.") + + if use_cudac: + supported_hv = (64, 48, 32, 24, 16, 12, 8, 6, 4, 2) + if HV not in supported_hv: + raise ValueError(f"backend='cudac' supports Qwen local HV in {supported_hv}, got {HV}") + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + out = torch.empty_like(v) + final_state = torch.empty(state_count, HV, K, K, device=q.device, dtype=torch.float32) + initial_state_arg = ( + torch.empty(0, device=q.device, dtype=torch.float32) + if initial_state is None + else initial_state.contiguous() + ) + cu_seqlens_arg = ( + torch.empty(0, device=q.device, dtype=torch.int32) + if cu_seqlens is None + else cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + ) + cula_cuda.qwen35_scalar_kda_prefill( + q.contiguous(), + k.contiguous(), + v.contiguous(), + a.contiguous(), + b.contiguous(), + A_log.contiguous(), + dt_bias.contiguous(), + initial_state_arg, + cu_seqlens_arg, + out, + final_state, + ) + return out, final_state + + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + state = ( + torch.zeros(state_count, HV, K, K, device=q.device, dtype=torch.float32) + if initial_state is None + else initial_state.float().clone() + ) + out = torch.empty_like(v) + q_f = torch.nn.functional.normalize(q.float(), dim=-1) * (K**-0.5) + k_f = torch.nn.functional.normalize(k.float(), dim=-1) + v_f = v.float() + a_f = a.float() + b_f = b.float() + A_log_f = A_log.float() + dt_bias_f = dt_bias.float() + + def _run_sequence(batch_idx: int, state_idx: int, start: int, end: int) -> None: + repeat = HV // H + for t in range(start, end): + for hv in range(HV): + qk_h = hv // repeat + state_kv = state[state_idx, hv] + decay = torch.exp(-torch.exp(A_log_f[hv]) * torch.nn.functional.softplus(a_f[batch_idx, t, hv] + dt_bias_f[hv])) + beta = torch.sigmoid(b_f[batch_idx, t, hv]) + k_vec = k_f[batch_idx, t, qk_h] + q_vec = q_f[batch_idx, t, qk_h] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v_f[batch_idx, t, hv] - proj) + state_kv_new = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[batch_idx, t, hv] = (state_kv_new.transpose(0, 1) @ q_vec).to(out.dtype) + state[state_idx, hv] = state_kv_new + + if cu_seqlens is None: + for bidx in range(B): + _run_sequence(bidx, bidx, 0, T) + else: + for sidx in range(state_count): + _run_sequence(0, sidx, int(cu_seqlens[sidx].item()), int(cu_seqlens[sidx + 1].item())) + return out, state + + +def qwen35_scalar_kda_prefill_core( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the Qwen GDN calculation after scalar gate/beta preprocessing. + + ``g`` is the natural-log per-token gate before the chunk-local prefix + scan, matching the tensors passed to SGLang's ``TritonGDNKernel.extend``. + The CUDA core still performs Q/K normalization and the prefix scan, while + the raw ``A_log/a/b/dt_bias`` conversion is intentionally outside timing. + """ + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4 or q.shape != k.shape: + raise ValueError("q/k/v must be 4D and q/k must have identical shapes") + B, T, H, K = q.shape + HV = v.shape[2] + if K != 128 or v.shape[:2] != q.shape[:2] or v.shape[-1] != 128 or HV % H: + raise ValueError(f"invalid native-GVA shapes q={tuple(q.shape)} v={tuple(v.shape)}") + if g.ndim == 2: + g = g.unsqueeze(0) + if beta.ndim == 2: + beta = beta.unsqueeze(0) + if g.shape != (B, T, HV) or beta.shape != g.shape: + raise ValueError(f"g/beta must be [B,T,HV], got {tuple(g.shape)} {tuple(beta.shape)}") + if g.dtype != torch.float32 or beta.dtype != torch.float32: + raise ValueError("g and beta must be float32") + if cu_seqlens is not None: + if B != 1 or cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError("cu_seqlens must be 1D int32 with B=1") + if initial_state is not None and initial_state.shape[1:] != (HV, K, K): + raise ValueError(f"initial_state must be [N,HV,128,128], got {tuple(initial_state.shape)}") + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core") + and q.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_prefill_core is unavailable") + if not use_cudac: + raise ValueError("qwen35_scalar_kda_prefill_core currently requires the CUDA backend") + + supported_hv = (64, 48, 32, 24, 16, 12, 8, 6, 4, 2) + if HV not in supported_hv: + raise ValueError(f"backend='cudac' supports local HV in {supported_hv}, got {HV}") + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + out = torch.empty_like(v) + final_state = torch.empty(state_count, HV, K, K, device=q.device, dtype=torch.float32) + initial_state_arg = ( + torch.empty(0, device=q.device, dtype=torch.float32) + if initial_state is None + else initial_state.contiguous() + ) + cu_seqlens_arg = ( + torch.empty(0, device=q.device, dtype=torch.int32) + if cu_seqlens is None + else cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + ) + cula_cuda.qwen35_scalar_kda_prefill_core( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta.contiguous(), + initial_state_arg, + cu_seqlens_arg, + out, + final_state, + ) + return out, final_state diff --git a/cula/qwen35/__init__.py b/cula/qwen35/__init__.py new file mode 100644 index 00000000..4cb88c3d --- /dev/null +++ b/cula/qwen35/__init__.py @@ -0,0 +1,28 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Qwen3.5-specific linear attention support built on top of cuLA primitives.""" + +from cula.qwen35.common import Qwen35LinearAttentionConfig + +try: + from cula.qwen35.runtime import ( + qwen35_linear_attention_decode, + qwen35_linear_attention_prefill, + ) +except Exception: # pragma: no cover - optional runtime dependency during partial imports + qwen35_linear_attention_decode = None + qwen35_linear_attention_prefill = None + +__all__ = ["Qwen35LinearAttentionConfig", "qwen35_linear_attention_prefill", "qwen35_linear_attention_decode"] diff --git a/cula/qwen35/common.py b/cula/qwen35/common.py new file mode 100644 index 00000000..63299c06 --- /dev/null +++ b/cula/qwen35/common.py @@ -0,0 +1,135 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Shared constants and validation helpers for Qwen3.5 linear attention.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class Qwen35LinearAttentionConfig: + """Minimal runtime config for Qwen3.5 linear-attention kernels.""" + + hidden_size: int = 5120 + conv_kernel_size: int = 4 + num_k_heads: int = 16 + num_v_heads: int = 48 + head_k_dim: int = 128 + head_v_dim: int = 128 + qkv_dtype: torch.dtype = torch.bfloat16 + state_dtype: torch.dtype = torch.float32 + + @property + def key_dim(self) -> int: + return self.num_k_heads * self.head_k_dim + + @property + def value_dim(self) -> int: + return self.num_v_heads * self.head_v_dim + + @property + def conv_dim(self) -> int: + return self.key_dim * 2 + self.value_dim + + @property + def qk_repeat_factor(self) -> int: + assert self.num_v_heads % self.num_k_heads == 0 + return self.num_v_heads // self.num_k_heads + + +DEFAULT_QWEN35_LINEAR_ATTN_CONFIG = Qwen35LinearAttentionConfig() + + +def validate_mixed_qkv( + mixed_qkv: torch.Tensor, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> None: + if mixed_qkv.dtype != config.qkv_dtype: + raise TypeError(f"mixed_qkv must be {config.qkv_dtype}, got {mixed_qkv.dtype}") + if mixed_qkv.ndim != 2: + raise ValueError(f"mixed_qkv must be 2D [tokens, conv_dim_local], got {tuple(mixed_qkv.shape)}") + if mixed_qkv.shape[-1] <= 0: + raise ValueError("mixed_qkv must have a non-zero channel dimension") + if mixed_qkv.shape[-1] % config.conv_dim != 0 and mixed_qkv.shape[-1] != config.conv_dim: + # In TP mode this is expected to be a local shard, so only require alignment + # with the Qwen3.5 packed layout ratio. + local_dim = mixed_qkv.shape[-1] + expected_splits = (config.key_dim, config.key_dim, config.value_dim) + if local_dim % sum(expected_splits) != 0: + raise ValueError(f"mixed_qkv last dim must match packed local conv dim, got {local_dim}") + + +def validate_scalar_gate_inputs( + a: torch.Tensor, + b: torch.Tensor, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> None: + if a.shape != b.shape: + raise ValueError(f"a and b must have the same shape, got a={tuple(a.shape)} vs b={tuple(b.shape)}") + if a.ndim != 2: + raise ValueError(f"a and b must be 2D [tokens, num_v_heads_local], got {tuple(a.shape)}") + if a.dtype != config.qkv_dtype or b.dtype != config.qkv_dtype: + raise TypeError(f"a and b must be {config.qkv_dtype}, got a={a.dtype}, b={b.dtype}") + + +def validate_state_tensors( + conv_state: torch.Tensor | None, + recurrent_state: torch.Tensor | None, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> None: + if conv_state is not None: + if conv_state.ndim != 3: + raise ValueError(f"conv_state must be 3D [batch, channels, {config.conv_kernel_size}], got {tuple(conv_state.shape)}") + if conv_state.dtype != config.qkv_dtype: + raise TypeError(f"conv_state must be {config.qkv_dtype}, got {conv_state.dtype}") + if recurrent_state is not None: + if recurrent_state.ndim != 4: + raise ValueError(f"recurrent_state must be 4D [batch, hv, k, v], got {tuple(recurrent_state.shape)}") + if recurrent_state.dtype != config.state_dtype: + raise TypeError(f"recurrent_state must be {config.state_dtype}, got {recurrent_state.dtype}") + + +def infer_local_config( + mixed_qkv_dim: int, + local_num_v_heads: int, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> tuple[int, int, int]: + """Infer local packed dims from runtime shard sizes. + + Returns: + - local_key_dim + - local_value_dim + - local_num_k_heads + """ + + local_value_dim = local_num_v_heads * config.head_v_dim + remaining = mixed_qkv_dim - local_value_dim + if remaining <= 0 or remaining % 2 != 0: + raise ValueError( + f"Cannot infer local q/k dims from mixed_qkv_dim={mixed_qkv_dim}, local_num_v_heads={local_num_v_heads}" + ) + local_key_dim = remaining // 2 + if local_key_dim % config.head_k_dim != 0: + raise ValueError(f"Local key dim must be divisible by head_k_dim={config.head_k_dim}, got {local_key_dim}") + local_num_k_heads = local_key_dim // config.head_k_dim + if local_num_v_heads % local_num_k_heads != 0: + raise ValueError( + f"Local num_v_heads={local_num_v_heads} must be divisible by local num_k_heads={local_num_k_heads}" + ) + return local_key_dim, local_value_dim, local_num_k_heads diff --git a/cula/qwen35/runtime.py b/cula/qwen35/runtime.py new file mode 100644 index 00000000..f3f55091 --- /dev/null +++ b/cula/qwen35/runtime.py @@ -0,0 +1,368 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""Runtime dispatch for Qwen3.5 linear-attention kernels.""" + +from __future__ import annotations + +import torch + +try: + import cuda.bindings.driver as cuda +except ImportError: # pragma: no cover - optional runtime dependency + cuda = None + +from cula.qwen35.common import ( + DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + Qwen35LinearAttentionConfig, + infer_local_config, + validate_mixed_qkv, + validate_scalar_gate_inputs, + validate_state_tensors, +) +from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_update +from cula.ops.qwen35_conv1d_prefill import qwen35_conv1d_prefill +from cula.ops.qwen35_layout_decode import qwen35_layout_decode +from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill +from cula.ops.qwen35_scalar_kda_decode import ( + has_qwen35_layout_scalar_kda_decode_cudac, + qwen35_layout_scalar_kda_decode, + qwen35_scalar_kda_decode, +) +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill + +_stream_cache: dict[tuple[str, int], object] = {} + + +def _get_cached_stream(device: torch.device) -> object: + if cuda is None: + raise RuntimeError("cuda.bindings.driver is not available in this environment.") + stream_id = int(torch.cuda.current_stream(device=device).cuda_stream) + cache_key = (str(device), stream_id) + if cache_key not in _stream_cache: + _stream_cache[cache_key] = cuda.CUstream(stream_id) + return _stream_cache[cache_key] + + +def _torch_qwen35_scalar_kda_decode_reference( + q_rep: torch.Tensor, + k_rep: torch.Tensor, + v: torch.Tensor, + a_kernel: torch.Tensor, + b_kernel: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pure torch reference for Qwen3.5 scalar-gated decode.""" + tokens, num_v_heads, head_k_dim = q_rep.shape + head_v_dim = v.shape[-1] + state_out = recurrent_state.clone() + out = torch.empty(tokens, num_v_heads, head_v_dim, device=q_rep.device, dtype=q_rep.dtype) + + scale = q_rep.shape[-1] ** -0.5 + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) * scale + k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) + v_f = v.float() + a_f = a_kernel.float() + b_f = b_kernel.float() + + for token_idx in range(tokens): + pool_idx = int(state_indices[token_idx].item()) + for hv in range(num_v_heads): + state_kv = state_out[pool_idx, hv] + state_vk = state_kv.transpose(0, 1).contiguous() + + decay_pre = a_f[token_idx, hv] + dt_bias[hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(decay_pre)) + beta = torch.sigmoid(b_f[token_idx, hv]) + + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + + proj = decay * (state_vk @ k_vec) + v_new = beta * (v_f[token_idx, hv] - proj) + state_vk_new = decay * state_vk + v_new.unsqueeze(1) * k_vec.unsqueeze(0) + out[token_idx, hv] = (state_vk_new @ q_vec).to(out.dtype) + state_out[pool_idx, hv] = state_vk_new.transpose(0, 1).contiguous() + + return out, state_out + + +def qwen35_linear_attention_decode_reference( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + conv_state: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure torch reference for the full Qwen3.5 decode chain.""" + tokens = mixed_qkv.shape[0] + if state_indices is None: + state_indices = torch.arange(tokens, device=mixed_qkv.device, dtype=torch.int32) + else: + state_indices = state_indices.to(device=mixed_qkv.device, dtype=torch.int32) + + conv_out, conv_state_out = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + backend="reference", + ) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_prefill( + conv_out, + a, + b, + config=config, + backend="reference", + ) + core_attn_out, recurrent_state_out = _torch_qwen35_scalar_kda_decode_reference( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log.float(), + dt_bias.float(), + recurrent_state.float(), + state_indices, + ) + return core_attn_out.reshape(tokens, -1), conv_state_out, recurrent_state_out + + +def qwen35_linear_attention_prefill( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + cu_seqlens: torch.Tensor | None = None, + recurrent_state: torch.Tensor | None = None, + conv_state: torch.Tensor | None = None, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Qwen3.5 prefill wrapper. + + Args: + mixed_qkv: flattened [tokens, local_conv_dim] + a, b: [tokens, local_num_v_heads] + conv_weight: [local_conv_dim, 1, 4] or [local_conv_dim, 4] + A_log, dt_bias: [local_num_v_heads] + cu_seqlens: optional int32 sequence offsets for flattened input + recurrent_state: optional initial recurrent state [num_sequences, HV, 128, 128] + + Returns: + - core_attn_out_flat: [tokens, local_value_dim] + - final conv_state: [num_sequences, local_conv_dim, 4] + - final recurrent_state: [num_sequences, local_num_v_heads, 128, 128] + """ + + validate_mixed_qkv(mixed_qkv, config) + validate_scalar_gate_inputs(a, b, config) + validate_state_tensors(conv_state, recurrent_state, config) + if mixed_qkv.is_cuda: + _get_cached_stream(mixed_qkv.device) + if conv_state is not None: + raise NotImplementedError("Qwen3.5 prefill with non-empty conv_state is not implemented yet.") + if mixed_qkv.shape[0] != a.shape[0]: + raise ValueError(f"Token dimension mismatch, got mixed_qkv={tuple(mixed_qkv.shape)} a={tuple(a.shape)}") + if A_log.ndim != 1 or dt_bias.ndim != 1 or A_log.shape != dt_bias.shape: + raise ValueError(f"A_log and dt_bias must be matching 1D tensors, got {tuple(A_log.shape)} and {tuple(dt_bias.shape)}") + if cu_seqlens is not None and (cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32): + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + + tokens = mixed_qkv.shape[0] + local_num_v_heads = a.shape[1] + _, local_value_dim, _ = infer_local_config( + mixed_qkv.shape[1], + local_num_v_heads, + config=config, + ) + if A_log.numel() != local_num_v_heads: + raise ValueError(f"A_log must match local_num_v_heads={local_num_v_heads}, got {A_log.numel()}") + + conv_out, conv_state_out = qwen35_conv1d_prefill( + mixed_qkv, + conv_weight, + activation="silu", + cu_seqlens=cu_seqlens, + output_final_state=True, + ) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_prefill( + conv_out, + a, + b, + config=config, + backend="reference" if backend == "reference" else "auto", + ) + core_attn_out, recurrent_state_out = qwen35_scalar_kda_prefill( + q=q_rep.unsqueeze(0).contiguous(), + k=k_rep.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a_kernel.unsqueeze(0).contiguous(), + b=b_kernel.unsqueeze(0).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + initial_state=recurrent_state, + cu_seqlens=cu_seqlens, + backend=backend, + ) + return core_attn_out.reshape(tokens, local_value_dim), conv_state_out, recurrent_state_out + + +def qwen35_linear_attention_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + conv_state: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor | None = None, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Qwen3.5 decode wrapper. + + Args: + mixed_qkv: [tokens, local_conv_dim] + a, b: [tokens, local_num_v_heads] + conv_weight: [local_conv_dim, 1, 4] or [local_conv_dim, 4] + A_log, dt_bias: [local_num_v_heads] + conv_state: [tokens, local_conv_dim, 4] + recurrent_state: [pool, local_num_v_heads, 128, 128] + + Returns: + - core_attn_out_flat: [tokens, local_value_dim] + - updated_conv_state + - updated_recurrent_state + """ + + validate_mixed_qkv(mixed_qkv, config) + validate_scalar_gate_inputs(a, b, config) + validate_state_tensors(conv_state, recurrent_state, config) + if mixed_qkv.is_cuda: + _get_cached_stream(mixed_qkv.device) + + if mixed_qkv.shape[0] != a.shape[0]: + raise ValueError(f"Token dimension mismatch, got mixed_qkv={tuple(mixed_qkv.shape)} a={tuple(a.shape)}") + if A_log.ndim != 1 or dt_bias.ndim != 1: + raise ValueError(f"A_log and dt_bias must be 1D, got {tuple(A_log.shape)} and {tuple(dt_bias.shape)}") + if A_log.shape != dt_bias.shape: + raise ValueError(f"A_log and dt_bias must have the same shape, got {tuple(A_log.shape)} vs {tuple(dt_bias.shape)}") + + tokens = mixed_qkv.shape[0] + local_num_v_heads = a.shape[1] + local_key_dim, local_value_dim, local_num_k_heads = infer_local_config( + mixed_qkv.shape[1], + local_num_v_heads, + config=config, + ) + if conv_state.shape != (tokens, mixed_qkv.shape[1], config.conv_kernel_size): + raise ValueError( + f"conv_state must be [tokens, local_conv_dim, {config.conv_kernel_size}], got {tuple(conv_state.shape)}" + ) + if recurrent_state.shape[1:] != (local_num_v_heads, config.head_k_dim, config.head_v_dim): + raise ValueError( + "recurrent_state must be [pool, local_num_v_heads, head_k_dim, head_v_dim], " + f"got {tuple(recurrent_state.shape)}" + ) + if A_log.numel() != local_num_v_heads: + raise ValueError(f"A_log must match local_num_v_heads={local_num_v_heads}, got {A_log.numel()}") + + if backend == "auto" and not mixed_qkv.is_cuda: + backend = "reference" + + if backend == "reference": + return qwen35_linear_attention_decode_reference( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + ) + + conv_out, conv_state_out = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + backend=backend, + ) + use_fused_layout_kda = ( + backend in ("auto", "cudac") + and mixed_qkv.is_cuda + and has_qwen35_layout_scalar_kda_decode_cudac() + ) + if backend == "cudac" and not use_fused_layout_kda: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_scalar_kda_decode is not available.") + + if use_fused_layout_kda: + core_attn_out, recurrent_state_out = qwen35_layout_scalar_kda_decode( + mixed_qkv_conv=conv_out, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + config=config, + backend=backend, + ) + core_attn_out = core_attn_out.reshape(tokens, local_value_dim) + return core_attn_out, conv_state_out, recurrent_state_out + + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode( + conv_out, + a, + b, + config=config, + backend=backend, + ) + q = q_rep.unsqueeze(1).contiguous() + k = k_rep.unsqueeze(1).contiguous() + v = v.unsqueeze(1).contiguous() + + core_attn_out, recurrent_state_out = qwen35_scalar_kda_decode( + q=q, + k=k, + v=v, + a=a_kernel, + b=b_kernel, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend=backend, + ) + core_attn_out = core_attn_out.reshape(tokens, local_value_dim) + return core_attn_out, conv_state_out, recurrent_state_out diff --git a/docs/qwen35_kernel_plan.md b/docs/qwen35_kernel_plan.md new file mode 100644 index 00000000..5aeaba4d --- /dev/null +++ b/docs/qwen35_kernel_plan.md @@ -0,0 +1,40 @@ +# Qwen3.5 Kernel Landing Plan Inside cuLA + +This note records the internal landing structure for Qwen3.5 linear-attention +support added directly inside `cuLA`. + +## New Python package surface + +- `cula/qwen35/__init__.py` +- `cula/qwen35/common.py` +- `cula/qwen35/runtime.py` + +## New CuTe op entry files + +- `cula/ops/qwen35_conv1d_prefill.py` +- `cula/ops/qwen35_conv1d_decode.py` +- `cula/ops/qwen35_scalar_kda_prefill.py` +- `cula/ops/qwen35_scalar_kda_decode.py` + +## Intended ownership + +- `common.py` + shared constants, local-head config, shape validation +- `runtime.py` + compile-cache, stream-cache, prefill/decode dispatch boundaries +- `qwen35_conv1d_*` + depthwise causal conv1d + silu +- `qwen35_scalar_kda_*` + scalar-gated delta-rule prefill/decode kernels + +## What should be reused from existing cuLA code + +- runtime compile-cache patterns from `cula/ops/kda_decode.py` +- device helpers from `cula/utils.py` +- operator boundary style from `cula/kda/chunk.py` + +## What should stay isolated at first + +- no direct mutation of the generic `chunk_kda` public entry +- no pybind work until Python/CuTe path is numerically correct +- no conv + kda fusion until standalone kernels are validated diff --git a/setup.py b/setup.py index a187764a..17503dfc 100644 --- a/setup.py +++ b/setup.py @@ -151,8 +151,8 @@ def get_nvcc_thread_args(): include_dirs = [ Path(this_dir) / "csrc", Path(this_dir) / "csrc" / "kerutils" / "include", - Path(this_dir) / "csrc" / "cutlass" / "include", - Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", + Path(this_dir) / "csrc" / "cutlass" / "include", + Path(this_dir) / "csrc" / "cutlass" / "tools" / "util" / "include", ] major, minor = get_nvcc_version() @@ -160,6 +160,14 @@ def get_nvcc_thread_args(): assert_blackwell_build_env() ext_modules = [] +qwen35_sources = [ + "csrc/qwen35/decode/qwen35_conv1d_decode.cu", + "csrc/qwen35/decode/qwen35_layout_decode.cu", + "csrc/qwen35/decode/qwen35_scalar_kda_decode.cu", + "csrc/qwen35/prefill/qwen35_layout_prefill.cu", + "csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu", + "csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu", +] if not DISABLE_SM100 or not DISABLE_SM103: sm100_arch_flags = [] @@ -175,7 +183,8 @@ def get_nvcc_thread_args(): "csrc/api/kda_sm100.cu", "csrc/api/pybind_sm100.cu", "csrc/kda/sm100/kda_fwd_sm100.cu", - ], + ] + + qwen35_sources, extra_compile_args={ "cxx": cxx_args + get_features_args(), "nvcc": nvcc_common_args @@ -199,7 +208,8 @@ def get_nvcc_thread_args(): "csrc/api/pybind_sm90.cu", "csrc/kda/sm90/kda_fwd_sm90.cu", "csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu", - ], + ] + + qwen35_sources, extra_compile_args={ "cxx": cxx_args + get_features_args(), "nvcc": nvcc_common_args diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py new file mode 100644 index 00000000..7f206185 --- /dev/null +++ b/tests/test_qwen35_decode.py @@ -0,0 +1,645 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.ops.qwen35_layout_decode import qwen35_layout_decode, qwen35_layout_decode_reference +from cula.ops.qwen35_scalar_kda_decode import qwen35_layout_scalar_kda_decode, qwen35_scalar_kda_decode +from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_reference, qwen35_conv1d_decode_update +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig +from cula.qwen35.runtime import qwen35_linear_attention_decode + +try: + from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as triton_fused_sigmoid_update +except ImportError: + triton_fused_sigmoid_update = None + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def _device(): + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def _has_qwen35_cudac(): + return ( + torch.cuda.is_available() + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_conv1d_decode") + and hasattr(cula_cuda, "qwen35_layout_decode") + and hasattr(cula_cuda, "qwen35_scalar_kda_decode") + ) + + +def _has_qwen35_fused_layout_kda_cudac(): + return _has_qwen35_cudac() and hasattr(cula_cuda, "qwen35_layout_scalar_kda_decode") + + +def make_inputs( + tokens: int = 2, + pool_size: int = 3, + device: torch.device | None = None, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +): + device = _device() if device is None else device + torch.manual_seed(0) + mixed_qkv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + conv_weight = torch.randn(config.conv_dim, config.conv_kernel_size, device=device, dtype=config.qkv_dtype) + conv_state = torch.randn(tokens, config.conv_dim, config.conv_kernel_size, device=device, dtype=config.qkv_dtype) + recurrent_state = torch.randn( + pool_size, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) % pool_size + return mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices + + +def manual_conv_decode(x_t: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor): + state_tail = conv_state[..., 1:].float() + window = torch.cat([state_tail, x_t.unsqueeze(-1).float()], dim=-1) + conv = (window * weight.float().unsqueeze(0)).sum(dim=-1) + y = torch.nn.functional.silu(conv).to(dtype=x_t.dtype) + state_new = conv_state.clone() + state_new[..., 0] = conv_state[..., 1] + state_new[..., 1] = conv_state[..., 2] + state_new[..., 2] = conv_state[..., 3] + state_new[..., 3] = x_t + return y, state_new + + +def manual_qwen35_decode_reference( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + conv_state: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +): + conv_out, conv_state_out = manual_conv_decode(mixed_qkv, conv_state, conv_weight) + q_end = config.key_dim + k_end = q_end + config.key_dim + q = conv_out[:, :q_end].view(mixed_qkv.shape[0], config.num_k_heads, config.head_k_dim) + k = conv_out[:, q_end:k_end].view(mixed_qkv.shape[0], config.num_k_heads, config.head_k_dim) + v = conv_out[:, k_end:].view(mixed_qkv.shape[0], config.num_v_heads, config.head_v_dim) + q_rep = q.repeat_interleave(config.qk_repeat_factor, dim=1) + k_rep = k.repeat_interleave(config.qk_repeat_factor, dim=1) + + scale = config.head_k_dim**-0.5 + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) * scale + k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) + v_f = v.float() + state_out = recurrent_state.clone() + out = torch.empty(mixed_qkv.shape[0], config.value_dim, device=mixed_qkv.device, dtype=mixed_qkv.dtype) + + for token_idx in range(mixed_qkv.shape[0]): + per_token = [] + pool_idx = int(state_indices[token_idx].item()) + for hv in range(config.num_v_heads): + state_kv = state_out[pool_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(a[token_idx, hv].float() + dt_bias[hv])) + beta = torch.sigmoid(b[token_idx, hv].float()) + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v_f[token_idx, hv] - proj) + state_new_kv = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + per_token.append((state_new_kv.transpose(0, 1) @ q_vec).to(mixed_qkv.dtype)) + state_out[pool_idx, hv] = state_new_kv + out[token_idx] = torch.cat(per_token, dim=0) + return out, conv_state_out, state_out + + +def manual_qwen35_scalar_kda_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, +): + if a.ndim == 2: + a = a.unsqueeze(1) + if b.ndim == 2: + b = b.unsqueeze(1) + N, _, HV, K = q.shape + scale = K**-0.5 + q_f = torch.nn.functional.normalize(q.squeeze(1).float(), dim=-1) * scale + k_f = torch.nn.functional.normalize(k.squeeze(1).float(), dim=-1) + v_f = v.squeeze(1).float() + state_out = recurrent_state.clone() + out = torch.empty(N, 1, HV, v.shape[-1], device=q.device, dtype=v.dtype) + + for token_idx in range(N): + pool_idx = int(state_indices[token_idx].item()) + for hv in range(HV): + state_kv = state_out[pool_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(a[token_idx, 0, hv].float() + dt_bias[hv])) + beta = torch.sigmoid(b[token_idx, 0, hv].float()) + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v_f[token_idx, hv] - proj) + state_new_kv = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[token_idx, 0, hv] = (state_new_kv.transpose(0, 1) @ q_vec).to(v.dtype) + state_out[pool_idx, hv] = state_new_kv + return out, state_out + + +def manual_qwen35_layout_scalar_kda_reference( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +): + q_rep, k_rep, v, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv_conv, a, b, config=config) + + scale = config.head_k_dim**-0.5 + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) * scale + k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) + v_f = v.float() + state_out = recurrent_state.clone() + out = torch.empty( + mixed_qkv_conv.shape[0], + q_rep.shape[1], + config.head_v_dim, + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + + for token_idx in range(mixed_qkv_conv.shape[0]): + pool_idx = int(state_indices[token_idx].item()) + for hv in range(q_rep.shape[1]): + state_kv = state_out[pool_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(a_ref[token_idx, hv].float() + dt_bias[hv])) + beta = torch.sigmoid(b_ref[token_idx, hv].float()) + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v_f[token_idx, hv] - proj) + state_new_kv = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[token_idx, hv] = (state_new_kv.transpose(0, 1) @ q_vec).to(mixed_qkv_conv.dtype) + state_out[pool_idx, hv] = state_new_kv + return out.unsqueeze(1), state_out + + +def _local_config(local_v_heads: int) -> Qwen35LinearAttentionConfig: + return Qwen35LinearAttentionConfig(num_k_heads=local_v_heads // 3, num_v_heads=local_v_heads) + + +@pytest.mark.parametrize("tokens", [1, 2]) +def test_qwen35_conv_decode_reference(tokens: int): + mixed_qkv, _, _, conv_weight, conv_state, _, _, _, _ = make_inputs(tokens=tokens) + y_ref, state_ref = manual_conv_decode(mixed_qkv, conv_state, conv_weight) + y_op, state_op = qwen35_conv1d_decode_update(mixed_qkv, conv_state, conv_weight, backend="reference") + assert torch.equal(y_ref, y_op) + assert torch.equal(state_ref, state_op) + y_ref2, state_ref2 = qwen35_conv1d_decode_reference(mixed_qkv, conv_state, conv_weight) + assert torch.equal(y_ref, y_ref2) + assert torch.equal(state_ref, state_ref2) + + +def test_qwen35_layout_decode_reference(): + mixed_qkv, a, b, _, _, _, _, _, _ = make_inputs(tokens=2) + q_rep_ref, k_rep_ref, v_ref, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv, a, b) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode(mixed_qkv, a, b, backend="reference") + assert torch.equal(q_rep_ref, q_rep) + assert torch.equal(k_rep_ref, k_rep) + assert torch.equal(v_ref, v) + assert torch.equal(a_ref, a_kernel) + assert torch.equal(b_ref, b_kernel) + + +@pytest.mark.parametrize("tokens", [1, 2]) +def test_qwen35_decode_reference_chain(tokens: int): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs(tokens=tokens) + out_ref, conv_state_ref, recurrent_state_ref = manual_qwen35_decode_reference( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state, + recurrent_state, + state_indices, + ) + out, conv_state_out, recurrent_state_out = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="reference", + ) + + assert torch.allclose(out_ref.float(), out.float(), atol=1e-5, rtol=1e-5) + assert torch.equal(conv_state_ref, conv_state_out) + assert torch.allclose(recurrent_state_ref, recurrent_state_out, atol=1e-6, rtol=1e-6) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("tokens", [1, 2, 4]) +def test_qwen35_decode_cudac_matches_reference(tokens: int): + # Decode batches represent distinct active sequences, so keep state rows unique + # to avoid intentionally racing multiple token updates against one cache row. + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=tokens, + pool_size=max(tokens, 3), + device=torch.device("cuda"), + ) + out_ref, conv_state_ref, recurrent_state_ref = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="reference", + ) + out, conv_state_out, recurrent_state_out = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + assert torch.allclose(out_ref.float(), out.float(), atol=3e-2, rtol=3e-2) + assert torch.equal(conv_state_ref, conv_state_out) + assert torch.allclose(recurrent_state_ref, recurrent_state_out, atol=3e-5, rtol=3e-5) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_conv_decode_cudac_supports_local_tp_shapes(local_v_heads: int): + config = _local_config(local_v_heads) + mixed_qkv, _, _, conv_weight, conv_state, _, _, _, _ = make_inputs( + tokens=3, + pool_size=3, + device=torch.device("cuda"), + config=config, + ) + y_ref, state_ref = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + backend="reference", + ) + y, state = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(y, y_ref) + torch.testing.assert_close(state, state_ref) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_scalar_kda_decode_cudac_supports_local_tp_shapes(local_v_heads: int): + torch.manual_seed(3) + config = _local_config(local_v_heads) + tokens = 3 + device = torch.device("cuda") + q = torch.randn(tokens, 1, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + k = torch.randn_like(q) + v = torch.randn(tokens, 1, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, 1, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, 1, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + recurrent_state = torch.randn( + tokens, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + + out_ref, state_ref = manual_qwen35_scalar_kda_reference( + q, + k, + v, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices, + ) + out, state = qwen35_scalar_kda_decode( + q, + k, + v, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state, state_ref, atol=3e-5, rtol=3e-5) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_decode_cudac_supports_local_tp_shapes(local_v_heads: int): + config = _local_config(local_v_heads) + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=3, + pool_size=3, + device=torch.device("cuda"), + config=config, + ) + out_ref, conv_state_ref, recurrent_state_ref = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="reference", + ) + out, conv_state_out, recurrent_state_out = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(conv_state_out, conv_state_ref) + torch.testing.assert_close(recurrent_state_out, recurrent_state_ref, atol=3e-5, rtol=3e-5) + + +@pytest.mark.skipif(not _has_qwen35_fused_layout_kda_cudac(), reason="Qwen3.5 fused layout+KDA CUDA backend is not available") +@pytest.mark.parametrize("tokens", [1, 2, 4]) +def test_qwen35_fused_layout_kda_cudac_matches_reference_unfused_and_triton(tokens: int): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=tokens, + pool_size=max(tokens, 3), + device=torch.device("cuda"), + ) + conv_out, _ = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + backend="cudac", + ) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode(conv_out, a, b, backend="cudac") + out_ref, state_ref = manual_qwen35_layout_scalar_kda_reference( + conv_out, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices, + ) + out_unfused, state_unfused = qwen35_scalar_kda_decode( + q=q_rep.unsqueeze(1), + k=k_rep.unsqueeze(1), + v=v.unsqueeze(1), + a=a_kernel, + b=b_kernel, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + if triton_fused_sigmoid_update is not None: + state_triton = recurrent_state.clone() + out_triton = triton_fused_sigmoid_update( + A_log=A_log, + a=a_kernel.unsqueeze(1).contiguous(), + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_rep.unsqueeze(1).contiguous(), + k=k_rep.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + b=b_kernel.unsqueeze(1).contiguous(), + initial_state_source=state_triton, + initial_state_indices=state_indices, + scale=DEFAULT_QWEN35_LINEAR_ATTN_CONFIG.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + is_kda=False, + ) + out_fused, state_fused = qwen35_layout_scalar_kda_decode( + mixed_qkv_conv=conv_out, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + config=DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend="cudac", + ) + + torch.cuda.synchronize() + assert torch.allclose(out_ref.float(), out_fused.float(), atol=3e-2, rtol=3e-2) + assert torch.allclose(state_ref, state_fused, atol=3e-5, rtol=3e-5) + assert torch.equal(out_unfused, out_fused) + assert torch.equal(state_unfused, state_fused) + if triton_fused_sigmoid_update is not None: + assert torch.allclose(out_triton.float(), out_fused.float(), atol=3e-2, rtol=3e-2) + assert torch.allclose(state_triton, state_fused, atol=3e-5, rtol=3e-5) + + +@pytest.mark.skipif( + not _has_qwen35_fused_layout_kda_cudac(), + reason="Qwen3.5 fused layout+KDA CUDA backend is not available", +) +@pytest.mark.parametrize("tokens", [64, 128]) +def test_qwen35_fused_layout_kda_cudac_long_matches_reference(tokens: int): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=tokens, + pool_size=tokens, + device=torch.device("cuda"), + ) + conv_out, _ = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + backend="cudac", + ) + out_ref, state_ref = manual_qwen35_layout_scalar_kda_reference( + conv_out, a, b, A_log, dt_bias, recurrent_state, state_indices + ) + out_fused, state_fused = qwen35_layout_scalar_kda_decode( + mixed_qkv_conv=conv_out, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + config=DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out_fused.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state_fused, state_ref, atol=3e-5, rtol=3e-5) + + +@pytest.mark.skipif(not _has_qwen35_fused_layout_kda_cudac(), reason="Qwen3.5 fused layout+KDA CUDA backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_layout_scalar_kda_cudac_supports_local_tp_shards(local_v_heads: int): + config = _local_config(local_v_heads) + mixed_qkv, a, b, _, _, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=2, + pool_size=3, + device=torch.device("cuda"), + config=config, + ) + out_ref, state_ref = manual_qwen35_layout_scalar_kda_reference( + mixed_qkv, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices, + config=config, + ) + q_rep_ref, k_rep_ref, v_ref, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv, a, b, config=config) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode(mixed_qkv, a, b, config=config, backend="cudac") + out, state = qwen35_layout_scalar_kda_decode( + mixed_qkv, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices=state_indices, + config=config, + backend="cudac", + ) + out_3d_gate, state_3d_gate = qwen35_layout_scalar_kda_decode( + mixed_qkv, + a.unsqueeze(1), + b.unsqueeze(1), + A_log, + dt_bias, + recurrent_state, + state_indices=state_indices, + config=config, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(q_rep, q_rep_ref) + torch.testing.assert_close(k_rep, k_rep_ref) + torch.testing.assert_close(v, v_ref) + torch.testing.assert_close(a_kernel, a_ref) + torch.testing.assert_close(b_kernel, b_ref) + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state, state_ref, atol=3e-5, rtol=3e-5) + torch.testing.assert_close(out_3d_gate, out) + torch.testing.assert_close(state_3d_gate, state) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +def test_qwen35_decode_cudac_rejects_duplicate_state_indices(): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, _ = make_inputs( + tokens=2, + pool_size=3, + device=torch.device("cuda"), + ) + state_indices = torch.zeros(2, device=mixed_qkv.device, dtype=torch.int32) + + with pytest.raises(ValueError, match="requires unique state_indices"): + qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) diff --git a/tests/test_qwen35_prefill.py b/tests/test_qwen35_prefill.py new file mode 100644 index 00000000..d3279eb3 --- /dev/null +++ b/tests/test_qwen35_prefill.py @@ -0,0 +1,569 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +import math +import pathlib +import sys + +import torch +import pytest + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.ops.qwen35_conv1d_prefill import qwen35_conv1d_prefill +from cula.ops.qwen35_fused_kda_prefill import has_qwen35_fused_kda_prefill, qwen35_fused_kda_prefill +from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill, qwen35_layout_prefill_reference +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill, qwen35_scalar_kda_prefill_core +from cula.qwen35.common import Qwen35LinearAttentionConfig +from cula.qwen35.runtime import qwen35_linear_attention_prefill + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def _manual_scalar_prefill(q, k, v, a, b, A_log, dt_bias, initial_state=None, cu_seqlens=None): + B, T, HV, K = q.shape + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + state = torch.zeros(state_count, HV, K, K, device=q.device, dtype=torch.float32) + if initial_state is not None: + state = initial_state.float().clone() + out = torch.empty_like(v) + q_f = torch.nn.functional.normalize(q.float(), dim=-1) * (K**-0.5) + k_f = torch.nn.functional.normalize(k.float(), dim=-1) + + def run_seq(batch_idx, state_idx, start, end): + for t in range(start, end): + for hv in range(HV): + state_kv = state[state_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv].float()) * torch.nn.functional.softplus(a[batch_idx, t, hv].float() + dt_bias[hv].float())) + beta = torch.sigmoid(b[batch_idx, t, hv].float()) + k_vec = k_f[batch_idx, t, hv] + q_vec = q_f[batch_idx, t, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v[batch_idx, t, hv].float() - proj) + state_new = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[batch_idx, t, hv] = (state_new.transpose(0, 1) @ q_vec).to(out.dtype) + state[state_idx, hv] = state_new + + if cu_seqlens is None: + for batch_idx in range(B): + run_seq(batch_idx, batch_idx, 0, T) + else: + for state_idx in range(state_count): + run_seq(0, state_idx, int(cu_seqlens[state_idx].item()), int(cu_seqlens[state_idx + 1].item())) + return out, state + + +def _local_config(local_v_heads: int) -> Qwen35LinearAttentionConfig: + return Qwen35LinearAttentionConfig(num_k_heads=local_v_heads // 3, num_v_heads=local_v_heads) + + +def test_qwen35_scalar_kda_prefill_reference_matches_manual(): + torch.manual_seed(0) + B, T, HV, K = 2, 3, 2, 128 + q = torch.randn(B, T, HV, K, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(B, T, HV, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, dtype=torch.bfloat16) + A_log = -torch.rand(HV, dtype=torch.float32) + dt_bias = torch.randn(HV, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, dtype=torch.float32) * 0.01 + + out_ref, state_ref = _manual_scalar_prefill(q, k, v, a, b, A_log, dt_bias, initial_state) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + + torch.testing.assert_close(out.float(), out_ref.float(), atol=1e-3, rtol=1e-3) + torch.testing.assert_close(state, state_ref, atol=1e-4, rtol=1e-4) + + +def test_qwen35_scalar_kda_prefill_varlen_reference_matches_manual(): + torch.manual_seed(1) + T, HV, K = 4, 2, 128 + q = torch.randn(1, T, HV, K, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(1, T, HV, dtype=torch.bfloat16) + b = torch.randn(1, T, HV, dtype=torch.bfloat16) + A_log = -torch.rand(HV, dtype=torch.float32) + dt_bias = torch.randn(HV, dtype=torch.float32) * 0.1 + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + + out_ref, state_ref = _manual_scalar_prefill(q, k, v, a, b, A_log, dt_bias, cu_seqlens=cu_seqlens) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + cu_seqlens=cu_seqlens, + backend="reference", + ) + + torch.testing.assert_close(out.float(), out_ref.float(), atol=1e-3, rtol=1e-3) + torch.testing.assert_close(state, state_ref, atol=1e-4, rtol=1e-4) + + +def test_qwen35_scalar_kda_prefill_cuda_matches_reference(): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + import pytest + + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(10) + device = torch.device("cuda") + B, T, HV, K = 1, 8, 48, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_layout_prefill_cuda_supports_local_tp_shards(local_v_heads: int): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_layout_prefill"): + pytest.skip("qwen35_layout_prefill CUDA extension is not available") + + torch.manual_seed(20 + local_v_heads) + device = torch.device("cuda") + config = _local_config(local_v_heads) + tokens = 5 + mixed_qkv = torch.randn(tokens, config.conv_dim, device=device, dtype=torch.bfloat16) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=torch.bfloat16) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=torch.bfloat16) + + ref = qwen35_layout_prefill_reference(mixed_qkv, a, b, config=config) + out = qwen35_layout_prefill(mixed_qkv, a, b, config=config, backend="cudac") + + torch.cuda.synchronize() + for out_tensor, ref_tensor in zip(out, ref, strict=True): + torch.testing.assert_close(out_tensor, ref_tensor) + + +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_scalar_kda_prefill_cuda_supports_local_tp_shards(local_v_heads: int): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(30 + local_v_heads) + device = torch.device("cuda") + B, T, HV, K = 1, 4, local_v_heads, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + +def test_qwen35_scalar_kda_prefill_cuda_long_nonzero_state_and_current_stream(): + """Exercise the tiled CTA path at the main Qwen prefill target length.""" + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(130) + device = torch.device("cuda") + B, T, H, HV, K = 1, 128, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out = torch.empty_like(v) + state = torch.empty_like(initial_state) + empty = torch.empty(0, device=device, dtype=torch.float32) + cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32) + # The extension must honor the current stream; using the default stream + # here would race this event stream in real inference. + stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(stream): + cula_cuda.qwen35_scalar_kda_prefill( + q.contiguous(), + k.contiguous(), + v.contiguous(), + a.contiguous(), + b.contiguous(), + A_log.contiguous(), + dt_bias.contiguous(), + initial_state.contiguous(), + cu_seqlens, + out, + state, + ) + stream.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("T", [32, 65, 128, 129]) +def test_qwen35_scalar_kda_prefill_core_cuda_matches_raw_reference(T: int): + """The preprocessed-gate ABI must preserve the raw scalar-kernel result.""" + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core"): + pytest.skip("qwen35_scalar_kda_prefill_core CUDA extension is not available") + + torch.manual_seed(132) + device = torch.device("cuda") + B, H, HV, K = 1, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + g = -torch.exp(A_log).view(1, 1, HV) * torch.nn.functional.softplus(a.float() + dt_bias.view(1, 1, HV)) + beta = torch.sigmoid(b.float()) + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out_core, state_core = qwen35_scalar_kda_prefill_core( + q, k, v, g, beta, initial_state=initial_state, backend="cudac" + ) + torch.cuda.synchronize() + torch.testing.assert_close(out_core.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state_core, state_ref, atol=2e-2, rtol=2e-2) + + +def test_qwen35_scalar_kda_prefill_core_sm90_falls_back_for_large_negative_gate(): + """Finite gates below the SM90 safe domain must use the exact fallback.""" + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] != 9 + or cula_cuda is None + or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core") + ): + pytest.skip("SM90 qwen35_scalar_kda_prefill_core CUDA extension is not available") + + torch.manual_seed(177) + device = torch.device("cuda") + B, T, H, HV, K = 1, 32, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.zeros(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = torch.full((HV,), math.log(10.0), device=device, dtype=torch.float32) + dt_bias = torch.zeros(HV, device=device, dtype=torch.float32) + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + g = -torch.exp(A_log).view(1, 1, HV) * torch.nn.functional.softplus(a.float()) + beta = torch.sigmoid(b.float()) + assert g.max().item() < -5.0 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out_core, state_core = qwen35_scalar_kda_prefill_core( + q, k, v, g, beta, initial_state=initial_state, backend="cudac" + ) + torch.cuda.synchronize() + assert torch.isfinite(out_core).all() + assert torch.isfinite(state_core).all() + torch.testing.assert_close(out_core.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state_core, state_ref, atol=2e-2, rtol=2e-2) + + +def test_qwen35_scalar_kda_prefill_core_sm90_tp8_hv6_exact_fallback(): + """TP8's HV=6 shape cannot use four-head TMA groups but remains correct.""" + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] != 9 + or cula_cuda is None + or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core") + ): + pytest.skip("SM90 qwen35_scalar_kda_prefill_core CUDA extension is not available") + + torch.manual_seed(181) + device = torch.device("cuda") + B, T, H, HV, K = 1, 32, 2, 6, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.zeros(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = torch.full((HV,), math.log(0.25), device=device, dtype=torch.float32) + dt_bias = torch.zeros(HV, device=device, dtype=torch.float32) + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + g = -torch.exp(A_log).view(1, 1, HV) * torch.nn.functional.softplus(a.float()) + beta = torch.sigmoid(b.float()) + g_before = g.clone() + beta_before = beta.clone() + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out_core, state_core = qwen35_scalar_kda_prefill_core( + q, k, v, g, beta, initial_state=initial_state, backend="cudac" + ) + torch.cuda.synchronize() + torch.testing.assert_close(out_core.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state_core, state_ref, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(g, g_before, atol=0.0, rtol=0.0) + torch.testing.assert_close(beta, beta_before, atol=0.0, rtol=0.0) + + +def test_qwen35_scalar_kda_prefill_cuda_varlen_multi_sequence(): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(131) + device = torch.device("cuda") + T, HV, K = 257, 12, 128 + cu_seqlens = torch.tensor([0, 65, 128, T], device=device, dtype=torch.int32) + q = torch.randn(1, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(1, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(3, HV, K, K, device=device, dtype=torch.float32) * 0.01 + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, cu_seqlens=cu_seqlens, backend="reference" + ) + out, state = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, cu_seqlens=cu_seqlens, backend="cudac" + ) + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + +def test_qwen35_chunk_qk_prefill_sm90_matches_torch(): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): + import pytest + + pytest.skip("qwen35_chunk_qk_prefill_sm90 CUDA extension is not available") + + torch.manual_seed(11) + device = torch.device("cuda") + B, T, HV, K = 1, 64, 48, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + out = torch.empty(B, HV, T, T, device=device, dtype=torch.float32) + + cula_cuda.qwen35_chunk_qk_prefill_sm90(q.contiguous(), k.contiguous(), out) + torch.cuda.synchronize() + + ref = torch.einsum("bthd,bshd->bhts", q.float(), k.float()) + torch.testing.assert_close(out, ref, atol=2e-1, rtol=2e-2) + + +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_chunk_qk_prefill_sm90_supports_local_tp_shards(local_v_heads: int): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): + pytest.skip("qwen35_chunk_qk_prefill_sm90 CUDA extension is not available") + + torch.manual_seed(40 + local_v_heads) + device = torch.device("cuda") + B, T, HV, K = 1, 32, local_v_heads, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + out = torch.empty(B, HV, T, T, device=device, dtype=torch.float32) + + cula_cuda.qwen35_chunk_qk_prefill_sm90(q.contiguous(), k.contiguous(), out) + torch.cuda.synchronize() + + ref = torch.einsum("bthd,bshd->bhts", q.float(), k.float()) + torch.testing.assert_close(out, ref, atol=2e-1, rtol=2e-2) + + +@pytest.mark.parametrize("T", [64, 128]) +def test_qwen35_fused_kda_prefill_matches_reference(T: int): + if not torch.cuda.is_available(): + import pytest + + pytest.skip("CUDA is not available") + if not has_qwen35_fused_kda_prefill(torch.device("cuda")): + import pytest + + pytest.skip("Qwen3.5 fused KDA prefill backend is not available") + + torch.manual_seed(12) + device = torch.device("cuda") + B, H, HV, K = 1, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q.repeat_interleave(HV // H, dim=2), + k.repeat_interleave(HV // H, dim=2), + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + out, state = qwen35_fused_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state, state_ref, atol=3e-2, rtol=3e-2) + + +def test_qwen35_conv1d_prefill_flattened_state(): + x = torch.arange(5 * 3, dtype=torch.bfloat16).reshape(5, 3) + weight = torch.ones(3, 4, dtype=torch.bfloat16) + cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.int32) + + y, state = qwen35_conv1d_prefill(x, weight, cu_seqlens=cu_seqlens, output_final_state=True) + + assert y.shape == x.shape + assert state.shape == (2, 3, 4) + torch.testing.assert_close(state[0, :, -2:], x[:2].transpose(0, 1)) + torch.testing.assert_close(state[1, :, -3:], x[2:5].transpose(0, 1)) + + +def test_qwen35_layout_prefill_reference(): + torch.manual_seed(2) + config = Qwen35LinearAttentionConfig(num_k_heads=1, num_v_heads=2) + tokens = 3 + mixed_qkv = torch.randn(tokens, config.conv_dim, dtype=torch.bfloat16) + a = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + b = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + + ref = qwen35_layout_prefill_reference(mixed_qkv, a, b, config=config) + out = qwen35_layout_prefill(mixed_qkv, a, b, config=config, backend="reference") + + for out_tensor, ref_tensor in zip(out, ref, strict=True): + assert torch.equal(out_tensor, ref_tensor) + + +def test_qwen35_linear_attention_prefill_reference_shapes(): + torch.manual_seed(2) + config = Qwen35LinearAttentionConfig(num_k_heads=1, num_v_heads=2) + tokens = 3 + mixed_qkv = torch.randn(tokens, config.conv_dim, dtype=torch.bfloat16) + a = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + b = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + conv_weight = torch.randn(config.conv_dim, config.conv_kernel_size, dtype=torch.bfloat16) + A_log = -torch.rand(config.num_v_heads, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, dtype=torch.float32) * 0.1 + cu_seqlens = torch.tensor([0, 2, 3], dtype=torch.int32) + + out, conv_state, recurrent_state = qwen35_linear_attention_prefill( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + cu_seqlens=cu_seqlens, + backend="reference", + ) + + assert out.shape == (tokens, config.value_dim) + assert conv_state.shape == (2, config.conv_dim, config.conv_kernel_size) + assert recurrent_state.shape == (2, config.num_v_heads, config.head_k_dim, config.head_v_dim)