From da9fb689d1e234ec7672134edad28ccaaddee79d Mon Sep 17 00:00:00 2001 From: Zheng QiHang Date: Wed, 26 Aug 2026 21:38:41 +0800 Subject: [PATCH] feat(fixtures): survey FlashInfer as an authored-HIR corpus WIP. Collects 89 FlashInfer kernels rewritten as authored HIR, plus the two scripts that produced and triaged them. 35 parse and analyze; 54 carry a `.blocked.py` suffix and record in their docstring the refusal they hit, so the corpus states what the surface cannot express yet rather than hiding it. The suffix is what keeps a blocked fixture out of packaging, and the source loader now resolves a sibling by its base stem so the rest still import. A fixture docstring is a parsed key/value record rather than prose, so the corpus joins the directories the documentation lint leaves alone. --- scripts/comment_hygiene_lint.py | 2 +- scripts/summarize_blocked.py | 176 ++ scripts/survey_flashinfer.py | 758 ++++++++ src/tilefoundry/cli/source.py | 2 +- tests/fixtures/flashinfer/__init__.py | 1 + .../attention_collective.blocked.py | 254 +++ ..._nvfp4_attention_sm120_quantize.blocked.py | 93 + .../attention_csrc_xqa_mha.blocked.py | 171 ++ .../flashinfer/attention_csrc_xqa_mha_sm90.py | 175 ++ .../attention_csrc_xqa_mla_sm120.blocked.py | 175 ++ .../attention_include_attention_batch_pod.py | 195 ++ ...nclude_attention_blackwell_plan.blocked.py | 44 + ...ention_include_attention_decode.blocked.py | 284 +++ ..._attention_decode_mla_cute_sm80.blocked.py | 101 ++ ...n_include_attention_hopper_prefill_sm90.py | 195 ++ ...ention_hopper_quantization_prefill_sm90.py | 195 ++ ...attention_include_attention_mla.blocked.py | 101 ++ ...on_include_attention_mla_hopper.blocked.py | 101 ++ ...n_include_attention_persistent_template.py | 195 ++ .../attention_include_attention_pod.py | 195 ++ .../attention_include_attention_prefill.py | 537 ++++++ ...ion_include_attention_scheduler.blocked.py | 82 + ..._mla_sm120_decode_dsv3_2_kernel.blocked.py | 101 ++ ...se_mla_sm120_decode_dsv4_kernel.blocked.py | 177 ++ ...sparse_mla_sm120_prefill_kernel.blocked.py | 329 ++++ .../attention_include_page.blocked.py | 197 +++ .../attention_py_attention_core.blocked.py | 203 +++ .../attention_py_attention_cute_dsl_fmha.py | 195 ++ ...ntion_cute_dsl_fmha_blockscaled.blocked.py | 60 + .../attention_py_cudnn_decode.blocked.py | 116 ++ .../flashinfer/attention_py_cudnn_prefill.py | 195 ++ ..._cute_dsl_attention_dsa_hca_fp8.blocked.py | 101 ++ ...tention_py_cute_dsl_attention_fmha_fmha.py | 195 ++ ...attention_fmha_fmha_blockscaled.blocked.py | 60 + ...ute_dsl_attention_fmha_quantize.blocked.py | 60 + ...ntion_py_cute_dsl_attention_fusion_mask.py | 76 + ...te_dsl_attention_fusion_variant.blocked.py | 109 ++ ...on_py_cute_dsl_attention_fusion_variant.py | 240 +++ ...y_cute_dsl_attention_gqa_decode.blocked.py | 116 ++ ..._dsl_attention_gqa_decode_paged.blocked.py | 116 ++ ...y_cute_dsl_attention_mla_decode.blocked.py | 101 ++ ...te_dsl_attention_mla_decode_fp8.blocked.py | 101 ++ ...cute_dsl_attention_mla_dispatch.blocked.py | 101 ++ ...attention_monolithic_mla_decode.blocked.py | 101 ++ ...tion_monolithic_mla_decode_fp16.blocked.py | 101 ++ ...ntion_monolithic_mla_decode_fp8.blocked.py | 101 ++ ...attention_py_cute_dsl_attention_prefill.py | 195 ++ ...attention_wrappers_batch_decode.blocked.py | 207 +++ ...sl_attention_wrappers_batch_hca.blocked.py | 101 ++ ...sl_attention_wrappers_batch_mla.blocked.py | 177 ++ ...te_dsl_attention_wrappers_batch_prefill.py | 195 ++ .../flashinfer/attention_py_decode.blocked.py | 757 ++++++++ .../attention_py_mla_core.blocked.py | 557 ++++++ ...tention_py_mla_sparse_mla_sm120.blocked.py | 177 ++ ...py_msa_ops_cute_dsl_sparse_decode_sm12x.py | 124 ++ ...y_msa_ops_cute_dsl_sparse_prefill_sm12x.py | 124 ++ .../attention_py_msa_ops_sparse_decode.py | 124 ++ .../attention_py_msa_ops_sparse_prefill.py | 124 ++ ...ention_py_nvfp4_attention_sm120.blocked.py | 95 + .../flashinfer/attention_py_page.blocked.py | 197 +++ ...arallel_attention_attention_ops.blocked.py | 44 + ...y_parallel_attention_parallel_attention.py | 195 ++ ...llel_attention_parallel_wrapper.blocked.py | 45 + tests/fixtures/flashinfer/attention_py_pod.py | 366 ++++ .../flashinfer/attention_py_prefill.py | 1563 +++++++++++++++++ .../flashinfer/attention_py_rope.blocked.py | 454 +++++ .../flashinfer/attention_py_xqa.blocked.py | 104 ++ tests/fixtures/flashinfer/attention_py_xqa.py | 197 +++ tests/fixtures/flashinfer/attention_state.py | 91 + .../fixtures/flashinfer/attention_variants.py | 88 + .../cute_dsl_add_rmsnorm_fp4quant.blocked.py | 115 ++ .../cute_dsl_rmsnorm_fp4quant.blocked.py | 86 + .../flashinfer/fp4_quantization.blocked.py | 83 + .../flashinfer/fused_dit_layernorm.py | 120 ++ .../fused_qk_rmsnorm_rope.blocked.py | 116 ++ .../flashinfer/gemm_svdquant.blocked.py | 86 + .../flashinfer/gemm_transforms.blocked.py | 151 ++ .../fixtures/flashinfer/ln_fwd_silu_kernel.py | 63 + .../flashinfer/moe_recurrent.blocked.py | 41 + .../flashinfer/moe_recurrent_spec.blocked.py | 139 ++ .../flashinfer/moe_state_hir.blocked.py | 162 ++ tests/fixtures/flashinfer/norm.py | 157 ++ tests/fixtures/flashinfer/norm_cuda.py | 65 + tests/fixtures/flashinfer/norm_cuh.py | 71 + .../norm_kernels_fused_add_rmsnorm.py | 142 ++ .../flashinfer/norm_kernels_layernorm.py | 71 + .../flashinfer/norm_kernels_rmsnorm.py | 169 ++ .../flashinfer/norm_quant_rope.blocked.py | 49 + .../flashinfer/paged_gather.blocked.py | 36 + .../paged_gather_placement.blocked.py | 30 + .../flashinfer/routing_gather.blocked.py | 81 + .../flashinfer/triton_kernels_norm.py | 86 + tests/fixtures/flashinfer/triton_norm.py | 111 ++ tests/fixtures/test_corpus.py | 126 ++ 94 files changed, 15939 insertions(+), 2 deletions(-) create mode 100644 scripts/summarize_blocked.py create mode 100644 scripts/survey_flashinfer.py create mode 100644 tests/fixtures/flashinfer/__init__.py create mode 100644 tests/fixtures/flashinfer/attention_collective.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_csrc_nvfp4_attention_sm120_quantize.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_csrc_xqa_mha.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_csrc_xqa_mha_sm90.py create mode 100644 tests/fixtures/flashinfer/attention_csrc_xqa_mla_sm120.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_batch_pod.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_blackwell_plan.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_decode.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_decode_mla_cute_sm80.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_hopper_prefill_sm90.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_hopper_quantization_prefill_sm90.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_mla.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_mla_hopper.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_persistent_template.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_pod.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_prefill.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_scheduler.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv3_2_kernel.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv4_kernel.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_prefill_kernel.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_include_page.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_attention_core.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha.py create mode 100644 tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha_blockscaled.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cudnn_decode.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cudnn_prefill.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_dsa_hca_fp8.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha_blockscaled.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_quantize.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_mask.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode_paged.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode_fp8.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_dispatch.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp16.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp8.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_prefill.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_decode.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_hca.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_mla.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_prefill.py create mode 100644 tests/fixtures/flashinfer/attention_py_decode.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_mla_core.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_mla_sparse_mla_sm120.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_decode_sm12x.py create mode 100644 tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_prefill_sm12x.py create mode 100644 tests/fixtures/flashinfer/attention_py_msa_ops_sparse_decode.py create mode 100644 tests/fixtures/flashinfer/attention_py_msa_ops_sparse_prefill.py create mode 100644 tests/fixtures/flashinfer/attention_py_nvfp4_attention_sm120.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_page.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_parallel_attention_attention_ops.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_attention.py create mode 100644 tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_wrapper.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_pod.py create mode 100644 tests/fixtures/flashinfer/attention_py_prefill.py create mode 100644 tests/fixtures/flashinfer/attention_py_rope.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_xqa.blocked.py create mode 100644 tests/fixtures/flashinfer/attention_py_xqa.py create mode 100644 tests/fixtures/flashinfer/attention_state.py create mode 100644 tests/fixtures/flashinfer/attention_variants.py create mode 100644 tests/fixtures/flashinfer/cute_dsl_add_rmsnorm_fp4quant.blocked.py create mode 100644 tests/fixtures/flashinfer/cute_dsl_rmsnorm_fp4quant.blocked.py create mode 100644 tests/fixtures/flashinfer/fp4_quantization.blocked.py create mode 100644 tests/fixtures/flashinfer/fused_dit_layernorm.py create mode 100644 tests/fixtures/flashinfer/fused_qk_rmsnorm_rope.blocked.py create mode 100644 tests/fixtures/flashinfer/gemm_svdquant.blocked.py create mode 100644 tests/fixtures/flashinfer/gemm_transforms.blocked.py create mode 100644 tests/fixtures/flashinfer/ln_fwd_silu_kernel.py create mode 100644 tests/fixtures/flashinfer/moe_recurrent.blocked.py create mode 100644 tests/fixtures/flashinfer/moe_recurrent_spec.blocked.py create mode 100644 tests/fixtures/flashinfer/moe_state_hir.blocked.py create mode 100644 tests/fixtures/flashinfer/norm.py create mode 100644 tests/fixtures/flashinfer/norm_cuda.py create mode 100644 tests/fixtures/flashinfer/norm_cuh.py create mode 100644 tests/fixtures/flashinfer/norm_kernels_fused_add_rmsnorm.py create mode 100644 tests/fixtures/flashinfer/norm_kernels_layernorm.py create mode 100644 tests/fixtures/flashinfer/norm_kernels_rmsnorm.py create mode 100644 tests/fixtures/flashinfer/norm_quant_rope.blocked.py create mode 100644 tests/fixtures/flashinfer/paged_gather.blocked.py create mode 100644 tests/fixtures/flashinfer/paged_gather_placement.blocked.py create mode 100644 tests/fixtures/flashinfer/routing_gather.blocked.py create mode 100644 tests/fixtures/flashinfer/triton_kernels_norm.py create mode 100644 tests/fixtures/flashinfer/triton_norm.py create mode 100644 tests/fixtures/test_corpus.py diff --git a/scripts/comment_hygiene_lint.py b/scripts/comment_hygiene_lint.py index 8c4e026f..99c9cba8 100644 --- a/scripts/comment_hygiene_lint.py +++ b/scripts/comment_hygiene_lint.py @@ -20,7 +20,7 @@ MAX_PROSE_LINES = 8 MAX_COLUMNS = 100 -EXEMPT_PREFIXES = ("tests/models/", "examples/") +EXEMPT_PREFIXES = ("tests/models/", "tests/fixtures/flashinfer/", "examples/") DIRECTIVE_PREFIXES = ("ruff:", "noqa", "type:", "pragma:", "mypy:", "fmt:", "isort:") PYTHON_SUFFIXES = frozenset({".py"}) C_SUFFIXES = frozenset({".h", ".hpp", ".cuh", ".cu", ".cpp", ".cc"}) diff --git a/scripts/summarize_blocked.py b/scripts/summarize_blocked.py new file mode 100644 index 00000000..866d5f10 --- /dev/null +++ b/scripts/summarize_blocked.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Summarize blocked fixtures by the first capability that prevents promotion.""" + +from __future__ import annotations + +import argparse +import ast +import re +from collections import Counter +from dataclasses import dataclass, field +from pathlib import Path + +_UNSUPPORTED_CALL = re.compile(r"^runtime_expression: unsupported call '([^']+)'") + + +@dataclass +class Reason: + """One observed blocker shared by one or more fixture files.""" + + state: str + key: str + files: list[Path] = field(default_factory=list) + diagnostics: Counter[str] = field(default_factory=Counter) + ledgers: set[str] = field(default_factory=set) + + +def _fields(path: Path) -> dict[str, str]: + document = ast.get_docstring(ast.parse(path.read_text(encoding="utf-8"))) + if document is None: + raise ValueError(f"{path}: missing module docstring") + fields = {} + for line in document.splitlines(): + key, separator, value = line.partition(":") + if separator: + fields[key] = value.strip() + return fields + + +def _ledger_ids(value: str | None) -> set[str]: + if value is None: + return set() + return {item.strip() for item in value.split(",") if item.strip()} + + +def _reason_key(fields: dict[str, str]) -> tuple[str, str]: + state = fields["blocked"] + error = fields.get("error") + if error is not None: + match = _UNSUPPORTED_CALL.match(error) + if match is not None: + return state, f"unsupported authored call `{match.group(1)}`" + return state, error + if state == "mis-analyzed": + ledger = fields.get("ledger", "untracked") + return state, f"mis-analysis tracked by {ledger}" + raise ValueError("refused fixture is missing error") + + +def _validate(path: Path, fields: dict[str, str]) -> None: + state = fields.get("blocked") + if state not in {"refused", "mis-analyzed"}: + raise ValueError(f"{path}: invalid blocked state {state!r}") + if fields.get("phase") not in {"load", "selection/analysis"}: + raise ValueError(f"{path}: invalid or missing phase") + if state == "refused" and "error" not in fields: + raise ValueError(f"{path}: refused fixture is missing error") + if state == "mis-analyzed": + missing = {"got", "expected", "why"} - fields.keys() + if missing: + raise ValueError(f"{path}: mis-analyzed fixture is missing {sorted(missing)}") + + +def _summary_table(reasons: list[Reason]) -> list[str]: + lines = [ + "| Rank | State | First blocker | Blocked files | Fixture ledger refs |", + "|---:|---|---|---:|---|", + ] + for rank, reason in enumerate(reasons, 1): + ledgers = ", ".join(f"`{item}`" for item in sorted(reason.ledgers)) or "untracked" + lines.append( + f"| {rank} | {reason.state} | {reason.key} | {len(reason.files)} | {ledgers} |" + ) + return lines + + +def _reason_details(reasons: list[Reason], root: Path) -> list[str]: + lines = [] + for rank, reason in enumerate(reasons, 1): + lines.extend( + [ + "", + f"## {rank}. {reason.key}", + "", + f"- State: `{reason.state}`", + f"- Blocked files: {len(reason.files)}", + "- Fixture ledger refs: " + + (", ".join(f"`{item}`" for item in sorted(reason.ledgers)) or "untracked"), + "- Observed diagnostics:", + "", + ] + ) + lines.extend( + f" - {count} x `{diagnostic}`" + for diagnostic, count in reason.diagnostics.most_common() + ) + lines.extend(["", "- Files:", ""]) + lines.extend(f" - `{path.relative_to(root)}`" for path in sorted(reason.files)) + return lines + + +def _non_gap_details(rows: list[tuple[Path, dict[str, str]]], root: Path) -> list[str]: + lines = [ + "", + "# Non-capability classifications", + "", + "These fixtures are excluded from repair priority counts.", + "", + ] + for path, fields in sorted(rows): + lines.extend( + [ + f"## `{path.relative_to(root)}`", + "", + f"- Classification: {fields['classification']}", + f"- Observed result: `{fields.get('error', fields.get('got', 'unknown'))}`", + "", + ] + ) + return lines + + +def summarize(root: Path) -> str: + """Return a deterministic Markdown summary for all blocked fixtures below root.""" + grouped: dict[tuple[str, str], Reason] = {} + non_gap = [] + paths = sorted(root.rglob("*.blocked.py")) + for path in paths: + fields = _fields(path) + _validate(path, fields) + if "classification" in fields: + non_gap.append((path, fields)) + continue + state, key = _reason_key(fields) + reason = grouped.setdefault((state, key), Reason(state=state, key=key)) + reason.files.append(path) + diagnostic = fields.get("error", fields.get("got", "unknown")) + reason.diagnostics[diagnostic] += 1 + reason.ledgers.update(_ledger_ids(fields.get("ledger"))) + + reasons = sorted(grouped.values(), key=lambda item: (-len(item.files), item.state, item.key)) + lines = [ + "# Blocked fixture reasons", + "", + "- Generator: `uv run python scripts/summarize_blocked.py tests/fixtures`", + f"- Blocked fixture files scanned: {len(paths)}", + f"- Capability blocker groups: {len(reasons)}", + f"- Non-capability classifications: {len(non_gap)}", + "- Priority proxy: descending files currently stopped at the first blocker; actual " + "promotions require a post-repair rerun.", + "", + *_summary_table(reasons), + *_reason_details(reasons, root), + *_non_gap_details(non_gap, root), + ] + return "\n".join(lines).rstrip() + "\n" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("fixtures", type=Path) + args = parser.parse_args() + print(summarize(args.fixtures), end="") + + +if __name__ == "__main__": + main() diff --git a/scripts/survey_flashinfer.py b/scripts/survey_flashinfer.py new file mode 100644 index 00000000..001032d3 --- /dev/null +++ b/scripts/survey_flashinfer.py @@ -0,0 +1,758 @@ +#!/usr/bin/env python3 +"""Generate the pinned FlashInfer kernel-source survey.""" + +from __future__ import annotations + +import argparse +import ast +import collections +import dataclasses +import re +import subprocess +from pathlib import Path, PurePosixPath + +FAMILIES = ( + "norm", + "attention", + "moe", + "quant", + "gemm", + "comm", + "recurrent", + "sampling", + "unclassified", +) +EXPECTED_COVERED_STRATEGIES = { + "A03", + "A07", + "G01", + "N01", + "N03", + "N05", + "N06", + "N07", + "Q01", +} +SUPPORT_ROW = re.compile(r"^\|\s*([A-Z]\d{2})\s*\|\s*([^|]+?)\s*\|") +CATALOG_ROW = re.compile(r"^\|\s*([A-Z]\d{2})\s*\|\s*([^|]+?)\s*\|\s*([^|]+?)\s*\|") +STRATEGY_ID = re.compile(r"\b[A-Z]\d{2}\b") +REVISION = re.compile(r"\b[0-9a-f]{40}\b") +SOURCE_REFERENCE = re.compile( + r"(?P(?:csrc|include/flashinfer|flashinfer|benchmarks|tests)/[^\s`:,+]+\.(?:cuh|cu|py|rst))" + r"(?::(?P[A-Za-z_][A-Za-z0-9_.]*))?" +) +GLOBAL_MARKER = re.compile(r"(?:__global__|FLASHINFER_GLOBAL)") +CALL_NAME = re.compile(r"\b([A-Za-z_][A-Za-z0-9_]*)\s*\(") +ENTRY_TOKEN = re.compile( + r"(kernel|fused|fusion|attention|prefill|decode|cascade|merge_state|norm|quant|rope|page|" + r"append|moe|expert|routing|topk|top_k|sampling|sample|gemm|matmul|mm_|_mm|allreduce|" + r"all_reduce|alltoall|all_to_all|all_gather|kda|gdn|delta_rule|mamba|ssd|state_update|" + r"mhc|silu_and_mul|packbits)", + re.IGNORECASE, +) +FUSION_TOKEN = re.compile( + r"(fused|fusion|and_mul|add_rmsnorm|rmsnorm_silu|quantize_append|all_gather_matmul|" + r"allreduce_fusion|attention|moe|kda|gdn|mamba|mhc|ssd|sampling)", + re.IGNORECASE, +) +NON_ENTRY_PREFIXES = ( + "check_", + "compile_", + "create_", + "gen_", + "get_", + "has_", + "is_", + "make_", + "set_", + "validate_", +) +NON_ENTRY_CLASS_SUFFIXES = ("Config", "Info", "Role", "Type", "Workspace") +PUBLIC_ENTRY_DECORATORS = {"custom_op", "flashinfer_api"} + + +@dataclasses.dataclass(frozen=True) +class CatalogEntry: + strategy: str + boundary: str + paths: tuple[str, ...] + symbols: tuple[str, ...] + + +@dataclasses.dataclass(frozen=True) +class SurveyRow: + path: str + family: str + entries: tuple[str, ...] + fusion: bool + strategies: tuple[str, ...] + fixtures: tuple[str, ...] + + +def _run_git(upstream: Path, *args: str) -> str: + completed = subprocess.run( + ["git", "-C", str(upstream), *args], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + raise SystemExit(completed.stderr.strip() or "git command failed") + return completed.stdout + + +def _tracked_files(upstream: Path) -> tuple[str, ...]: + return tuple(path for path in _run_git(upstream, "ls-files").splitlines() if path) + + +def _in_source_scope(path: str) -> bool: + item = PurePosixPath(path) + if path.startswith("csrc/"): + return item.suffix in {".cu", ".cuh"} + if path.startswith("include/flashinfer/"): + return item.suffix == ".cuh" + return path.startswith("flashinfer/") and item.suffix == ".py" + + +def _parse_support_matrix(path: Path) -> dict[str, str]: + strategies: dict[str, str] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + match = SUPPORT_ROW.match(line) + if match: + strategies[match.group(1)] = match.group(2).strip() + if len(strategies) != 60: + raise SystemExit(f"expected 60 SUPPORT-MATRIX strategies, found {len(strategies)}") + return strategies + + +def _parse_catalog(path: Path, support: dict[str, str]) -> tuple[CatalogEntry, ...]: + entries: list[CatalogEntry] = [] + previous_paths: tuple[str, ...] = () + previous_symbols: tuple[str, ...] = () + for line in path.read_text(encoding="utf-8").splitlines(): + match = CATALOG_ROW.match(line) + if not match or match.group(1) not in support: + continue + strategy, provenance = match.group(1), match.group(3) + references = list(SOURCE_REFERENCE.finditer(provenance)) + paths = tuple(dict.fromkeys(item.group("path") for item in references)) + symbols = tuple( + dict.fromkeys(item.group("symbol") for item in references if item.group("symbol")) + ) + if not paths and re.search(r"\bsame (?:symbol|template|API)\b", provenance, re.IGNORECASE): + paths, symbols = previous_paths, previous_symbols + if paths: + previous_paths, previous_symbols = paths, symbols + entries.append(CatalogEntry(strategy, support[strategy], paths, symbols)) + found = {entry.strategy for entry in entries} + if found != set(support): + missing = ", ".join(sorted(set(support) - found)) + raise SystemExit(f"catalog is missing SUPPORT-MATRIX strategies: {missing}") + return tuple(entries) + + +def _fixture_coverage(fixtures: Path) -> dict[str, tuple[str, ...]]: + coverage: dict[str, list[str]] = collections.defaultdict(list) + for path in sorted(fixtures.glob("*.py")): + if path.name == "__init__.py" or path.name.endswith(".blocked.py"): + continue + for strategy in sorted(set(STRATEGY_ID.findall(path.read_text(encoding="utf-8")))): + coverage[strategy].append(path.name) + found = set(coverage) + if found != EXPECTED_COVERED_STRATEGIES: + missing = ", ".join(sorted(EXPECTED_COVERED_STRATEGIES - found)) or "none" + extra = ", ".join(sorted(found - EXPECTED_COVERED_STRATEGIES)) or "none" + raise SystemExit(f"covered strategy drift: missing [{missing}], extra [{extra}]") + return {strategy: tuple(paths) for strategy, paths in coverage.items()} + + +def _catalog_indexes( + catalog: tuple[CatalogEntry, ...], +) -> tuple[dict[str, tuple[str, ...]], dict[str, tuple[str, ...]]]: + strategies: dict[str, list[str]] = collections.defaultdict(list) + symbols: dict[str, list[str]] = collections.defaultdict(list) + for entry in catalog: + for path in entry.paths: + strategies[path].append(entry.strategy) + symbols[path].extend(entry.symbols) + return ( + {path: tuple(dict.fromkeys(values)) for path, values in strategies.items()}, + {path: tuple(dict.fromkeys(values)) for path, values in symbols.items()}, + ) + + +def _family_from_strategy(strategy: str) -> str: + return { + "N": "norm", + "A": "attention", + "R": "attention", + "M": "moe", + "Q": "quant", + "G": "gemm", + "C": "comm", + "K": "recurrent", + "S": "recurrent", + "H": "recurrent", + "T": "sampling", + "L": "sampling", + }[strategy[0]] + + +def _classify_family(path: str, strategies: tuple[str, ...]) -> str | None: + if strategies: + families = {_family_from_strategy(strategy) for strategy in strategies} + if len(families) == 1: + return families.pop() + value = path.lower() + rooted_rules = ( + ("comm", ("/comm/", "include/flashinfer/comm/")), + ("moe", ("/fused_moe/", "/moe_ep/", "/grouped_mm/")), + ("recurrent", ("/kda/", "/kda_", "/gdn_", "/gdn/", "/mamba/", "/mhc")), + ("attention", ("/attention/", "/mla/", "/xqa/")), + ("norm", ("/norm/",)), + ("gemm", ("/gemm/", "/deep_gemm/")), + ("quant", ("/quantization/",)), + ("sampling", ("/sampling", "/topk", "/logits_processor/")), + ) + for family, needles in rooted_rules: + if any(needle in value for needle in needles): + return family + rules = ( + ("comm", ("/comm/", "allreduce", "all_reduce", "alltoall", "all_to_all", "nvshmem")), + ("moe", ("moe", "expert", "routing", "router", "grouped_mm")), + ("recurrent", ("kda", "gdn", "delta_rule", "mamba", "ssd", "mhc", "state_update")), + ("attention", ("attention", "/mla", "decode", "prefill", "page", "rope", "xqa", "pod")), + ("norm", ("norm", "layernorm", "rmsnorm")), + ("quant", ("quant", "fp4", "fp8", "mxfp", "activation", "packbits")), + ("gemm", ("gemm", "matmul", "tinygemm", "bgmv", "/bmm")), + ("sampling", ("sampling", "topk", "top_k", "top_p", "logits", "air_top_p")), + ) + for family, needles in rules: + if any(needle in value for needle in needles): + return family + return None + + +def _decorator_name(node: ast.expr) -> str | None: + if isinstance(node, ast.Call): + return _decorator_name(node.func) + if isinstance(node, ast.Attribute): + return node.attr + if isinstance(node, ast.Name): + return node.id + return None + + +def _python_entries( + path: Path, + catalog_symbols: tuple[str, ...], + *, + explicit_only: bool = False, +) -> tuple[str, ...]: + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except (SyntaxError, UnicodeDecodeError): + return catalog_symbols + entries: list[str] = list(catalog_symbols) + for node in tree.body: + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + explicitly_public = any( + _decorator_name(decorator) in PUBLIC_ENTRY_DECORATORS + for decorator in node.decorator_list + ) + if explicitly_public: + entries.append(node.name) + continue + if explicit_only: + continue + if node.name.startswith("_"): + continue + if node.name.startswith(NON_ENTRY_PREFIXES): + continue + if isinstance(node, ast.ClassDef) and node.name.endswith(NON_ENTRY_CLASS_SUFFIXES): + continue + if ENTRY_TOKEN.search(node.name): + entries.append(node.name) + return tuple(dict.fromkeys(entries)) + + +def _remove_balanced_call(value: str, name: str) -> str: + while True: + start = value.find(name) + if start < 0: + return value + opening = value.find("(", start + len(name)) + if opening < 0: + return value + depth = 0 + closing = -1 + for index in range(opening, len(value)): + if value[index] == "(": + depth += 1 + elif value[index] == ")": + depth -= 1 + if depth == 0: + closing = index + break + if closing < 0: + return value + value = value[:start] + value[closing + 1 :] + + +def _cpp_entries(path: Path, catalog_symbols: tuple[str, ...]) -> tuple[str, ...]: + text = path.read_text(encoding="utf-8", errors="replace") + entries: list[str] = list(catalog_symbols) + for marker in GLOBAL_MARKER.finditer(text): + line_start = text.rfind("\n", 0, marker.start()) + 1 + if text[line_start : marker.start()].lstrip().startswith("#define"): + continue + tail = text[marker.end() : marker.end() + 1600] + brace = tail.find("{") + semicolon = tail.find(";") + ends = tuple(index for index in (brace, semicolon) if index >= 0) + if not ends: + continue + end = min(ends) + header = "\n".join( + line for line in tail[:end].splitlines() if not line.lstrip().startswith("#") + ) + for attribute in ("__launch_bounds__", "__cluster_dims__", "__maxnreg__"): + header = _remove_balanced_call(header, attribute) + functions = re.findall(r"\bvoid\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(", header) + if functions: + entries.append(functions[0]) + continue + candidates = CALL_NAME.findall(header) + if candidates: + entries.append(candidates[0]) + return tuple(dict.fromkeys(entries)) + + +def _hard_exclusion(path: str) -> tuple[str, str] | None: + if path.startswith("csrc/nv_internal/"): + return ( + "vendored third-party", + "vendored TensorRT-LLM/deep-gemm implementation is not FlashInfer's own kernel surface", + ) + parts = PurePosixPath(path).parts + excluded_segments = { + "jit": ( + "JIT/build plumbing", + "generates or loads kernels; it is not a semantic kernel boundary", + ), + "testing": ( + "test support", + "test-only generators and reference helpers are not shipped kernels", + ), + "autotuner": ( + "tuning/config", + "search and configuration code does not define kernel semantics", + ), + "profiler": ( + "observability", + "profiling and tracing code records execution rather than computing it", + ), + "trace": ( + "observability", + "profiling and tracing code records execution rather than computing it", + ), + "trace_apply": ( + "observability", + "profiling and tracing code records execution rather than computing it", + ), + "tuning_configs": ( + "tuning/config", + "search and configuration code does not define kernel semantics", + ), + } + for part in parts: + if part in excluded_segments: + return excluded_segments[part] + stem = PurePosixPath(path).stem.lower() + if stem in {"testing", "test_utils"}: + return ( + "test support", + "test-only generators and reference helpers are not shipped kernels", + ) + if any(token in stem for token in ("benchmark", "reference", "validation")): + return ( + "benchmark/reference", + "benchmark and reference implementations are evidence, not kernel APIs", + ) + return None + + +def _helper_exclusion(path: str) -> tuple[str, str] | None: + parts = PurePosixPath(path).parts + if any(part in {"helpers", "roles"} for part in parts): + return ( + "helper/utility", + "shared primitives and configuration are not complete kernel boundaries", + ) + stem = PurePosixPath(path).stem.lower() + helper_tokens = ( + "backend", + "compiler", + "epilogue", + "helper", + "mainloop", + "registry", + "runner", + "scheduler", + "staging", + "utils", + ) + if any(token in stem for token in helper_tokens): + return ( + "helper/utility", + "shared primitives and configuration are not complete kernel boundaries", + ) + if stem in { + "algo_knobs", + "api", + "base", + "collective_builder", + "common", + "compile", + "config", + "configs", + "enums", + "errors", + "fusion_rules", + "layer", + "mainloop_spec", + "pipeline_topology", + "prepare", + "runner_common", + "runners", + "schedule", + "tensors", + "tuner", + "utils", + "helpers", + "weights", + "workspace_base", + }: + return ( + "helper/utility", + "shared primitives and configuration are not complete kernel boundaries", + ) + if stem in {"__main__", "aot"} or stem.endswith("_enums"): + return ( + "helper/utility", + "shared primitives and configuration are not complete kernel boundaries", + ) + return None + + +def _is_fusion( + path: str, + family: str, + strategies: tuple[str, ...], + entries: tuple[str, ...], +) -> bool: + if strategies: + return True + if family in {"attention", "moe", "recurrent"}: + return True + return bool(FUSION_TOKEN.search(" ".join((path, *entries)))) + + +def _scan( + upstream: Path, + tracked: tuple[str, ...], + catalog_strategies: dict[str, tuple[str, ...]], + catalog_symbols: dict[str, tuple[str, ...]], + fixture_coverage: dict[str, tuple[str, ...]], +) -> tuple[tuple[SurveyRow, ...], dict[tuple[str, str], list[str]]]: + rows: list[SurveyRow] = [] + excluded: dict[tuple[str, str], list[str]] = collections.defaultdict(list) + for relative in tracked: + if not _in_source_scope(relative): + continue + strategies = catalog_strategies.get(relative, ()) + hard_exclusion = _hard_exclusion(relative) + if hard_exclusion: + excluded[hard_exclusion].append(relative) + continue + path = upstream / relative + symbols = catalog_symbols.get(relative, ()) + helper_exclusion = _helper_exclusion(relative) + if path.suffix == ".py": + entries = _python_entries( + path, + symbols, + explicit_only=helper_exclusion is not None, + ) + else: + entries = _cpp_entries(path, symbols) + if not entries: + if path.suffix in {".cu", ".cuh"} and "__global__" in path.read_text( + encoding="utf-8", errors="replace" + ): + excluded[ + ( + "CUDA marker without entry", + "contains __global__ only in a macro or unparseable declaration, not a kernel entry", + ) + ].append(relative) + continue + if helper_exclusion: + excluded[helper_exclusion].append(relative) + continue + excluded[ + ("no kernel entry", "no CUDA global definition or public Python kernel/API entry") + ].append(relative) + continue + family = _classify_family(relative, strategies) or "unclassified" + fixtures = tuple( + sorted( + { + fixture + for strategy in strategies + for fixture in fixture_coverage.get(strategy, ()) + } + ) + ) + rows.append( + SurveyRow( + path=relative, + family=family, + entries=entries, + fusion=_is_fusion(relative, family, strategies, entries), + strategies=strategies, + fixtures=fixtures, + ) + ) + return tuple(sorted(rows, key=lambda row: row.path)), excluded + + +def _cell(values: tuple[str, ...], empty: str = "-") -> str: + if not values: + return empty + return "
".join(f"`{value.replace('|', '|')}`" for value in values) + + +def _render( + upstream: Path, + revision: str, + tracked: tuple[str, ...], + rows: tuple[SurveyRow, ...], + excluded: dict[tuple[str, str], list[str]], + catalog: tuple[CatalogEntry, ...], + coverage: dict[str, tuple[str, ...]], +) -> str: + scoped = tuple(path for path in tracked if _in_source_scope(path)) + outside = len(tracked) - len(scoped) + covered_rows = tuple(row for row in rows if row.fixtures) + uncovered = tuple(row for row in rows if not row.fixtures) + family_counts = collections.Counter(row.family for row in uncovered) + root_counts = collections.Counter( + PurePosixPath(path).parts[0] for path in tracked if path not in scoped + ) + global_marker_paths = { + path + for path in scoped + if PurePosixPath(path).suffix in {".cu", ".cuh"} + and "__global__" in (upstream / path).read_text(encoding="utf-8", errors="replace") + } + included_paths = {row.path for row in rows} + global_exclusions = { + reason: len(global_marker_paths.intersection(paths)) + for (reason, _), paths in excluded.items() + if global_marker_paths.intersection(paths) + } + global_accounting = [ + f"{len(global_marker_paths.intersection(included_paths))} included", + *(f"{count} {reason}" for reason, count in sorted(global_exclusions.items())), + ] + lines = [ + "# FlashInfer transcribable-kernel survey", + "", + f"Pinned upstream: `flashinfer-ai/flashinfer@{revision}`.", + "", + "This file is generated. Reproduce it from the TileFoundry worktree with:", + "", + "```bash", + "PLAN_DIR=/path/to/docs/plans/hir-frontend", + "python scripts/survey_flashinfer.py --upstream ~/flashinfer \\", + ' --evidence-dir "$PLAN_DIR/evidence/flashinfer" \\', + ' --fixtures tests/fixtures/flashinfer --output "$PLAN_DIR/SURVEY.md"', + "```", + "", + "## Result", + "", + f"- Repository accounting: **{len(tracked)} tracked files = {len(scoped)} scoped source files + {outside} outside the source roots**.", + f"- Scoped accounting: **{len(scoped)} = {len(rows)} transcribable groups + {sum(len(paths) for paths in excluded.values())} excluded source files**.", + f"- D28 output surface: **{len(rows)} `.py` groups**; **{len(covered_rows)} covered**, **{len(uncovered)} uncovered**.", + f"- Existing semantic coverage: **{len(coverage)} strategies** across **{len(covered_rows)} upstream source files**. Strategy count and D28 file count are intentionally different.", + f"- CUDA marker audit: **{len(global_marker_paths)} files containing `__global__` = " + + " + ".join(global_accounting) + + "**.", + "", + "### Uncovered groups by mechanism family", + "", + "| Family | Files |", + "| --- | ---: |", + ] + lines.extend(f"| {family} | {family_counts[family]} |" for family in FAMILIES) + lines.extend( + [ + f"| **Total** | **{len(uncovered)}** |", + "", + "## Method", + "", + "The scan uses every Git-tracked file at the pinned revision. The source roots are", + "`csrc/**/*.{cu,cuh}`, `include/flashinfer/**/*.cuh`, and `flashinfer/**/*.py`.", + "After the explicit exclusions below, a C++ file enters the D28 surface when it", + "declares or defines a CUDA global kernel, or when the 60-strategy catalog names it as", + "a semantic boundary. A Python file enters when it defines a public kernel/API entry", + "or is catalog provenance. Family assignment happens only after admission; a file that", + "does not fit the eight D29 work-splitting labels is retained as `unclassified`.", + "", + "The Python test is static and deterministic: in an operation module, a top-level", + "public Function or Class name must contain a kernel/operation term such as attention,", + "norm, quant, MoE, GEMM, collective, recurrent, or sampling. A helper/utility module", + "enters only when the catalog names it or a top-level entry is explicitly exported by", + "`@flashinfer_api` or `@custom_op`; internal `@cute.jit` primitives do not make the", + "whole helper file a D28 kernel boundary. Factory/support names and test/JIT paths are", + "excluded.", + "", + "Catalog symbols are retained even when their implementation is a C++ template method", + "rather than a CUDA global. C++ declarations and definitions both create rows; bare", + "`#define ... __global__` aliases are accounted separately and do not invent entries.", + "", + "`fusion=yes` means either the support catalog names the file as a fusion boundary or", + "the file/entry names identify a multi-operation attention, MoE, recurrent, sampling,", + "or explicitly fused implementation. It is a survey classification, not a claim that", + "TileFoundry already expresses that boundary.", + "", + "## Existing coverage", + "", + "| Strategy | Boundary | Upstream source | Fixture |", + "| --- | --- | --- | --- |", + ] + ) + catalog_by_strategy = {entry.strategy: entry for entry in catalog} + for strategy in sorted(coverage): + entry = catalog_by_strategy[strategy] + source = _cell(tuple(path for path in entry.paths if _in_source_scope(path))) + lines.append(f"| {strategy} | {entry.boundary} | {source} | {_cell(coverage[strategy])} |") + lines.extend( + [ + "", + "## D28 source groups", + "", + "One row is one prospective corpus `.py`; one row may contain multiple independent", + "authored Modules when the upstream file exposes multiple kernel entries.", + "", + "| Upstream source | Family | Public/kernel entries | Fusion | Strategies | Coverage |", + "| --- | --- | --- | :---: | --- | --- |", + ] + ) + for row in rows: + lines.append( + f"| `{row.path}` | {row.family} | {_cell(row.entries)} | " + f"{'yes' if row.fusion else 'no'} | {_cell(row.strategies)} | {_cell(row.fixtures)} |" + ) + lines.extend( + [ + "", + "## Exclusions", + "", + "Every scoped source file not listed above is assigned exactly one reason. Family", + "classification is never an exclusion. `helper/utility` is considered only after no", + "CUDA or public Python kernel entry was found:", + "", + "| Reason | Files | Why |", + "| --- | ---: | --- |", + ] + ) + for (reason, why), paths in sorted(excluded.items()): + lines.append(f"| {reason} | {len(paths)} | {why} |") + lines.extend( + [ + "", + "Tracked files outside the three source roots were still counted. They are excluded", + "because docs, tests, benchmarks, CI, examples, packaging, and vendored dependencies", + "are not upstream kernel-source grouping units:", + "", + "| Top-level path | Files |", + "| --- | ---: |", + ] + ) + lines.extend( + f"| `{root + '/' if '.' not in root else root}` | {count} |" + for root, count in sorted(root_counts.items()) + ) + lines.append("") + for (reason, why), paths in sorted(excluded.items()): + lines.extend( + [ + f"
{reason}: {len(paths)} files", + "", + f"Reason: {why}.", + "", + "```text", + *sorted(paths), + "```", + "", + "
", + "", + ] + ) + return "\n".join(lines) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream", type=Path, required=True) + parser.add_argument("--evidence-dir", type=Path, required=True) + parser.add_argument("--fixtures", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def main() -> int: + args = _parse_args() + upstream = args.upstream.resolve() + evidence = args.evidence_dir.resolve() + revision = _run_git(upstream, "rev-parse", "HEAD").strip() + dirty = _run_git(upstream, "status", "--porcelain") + if dirty: + raise SystemExit("upstream worktree must be clean") + catalog_path = evidence / "CATALOG.md" + expected_revisions = set(REVISION.findall(catalog_path.read_text(encoding="utf-8"))) + if revision not in expected_revisions: + raise SystemExit(f"upstream revision {revision} is not pinned by {catalog_path}") + support = _parse_support_matrix(evidence / "SUPPORT-MATRIX.md") + catalog = _parse_catalog(catalog_path, support) + coverage = _fixture_coverage(args.fixtures.resolve()) + catalog_strategies, catalog_symbols = _catalog_indexes(catalog) + tracked = _tracked_files(upstream) + rows, excluded = _scan( + upstream, + tracked, + catalog_strategies, + catalog_symbols, + coverage, + ) + known_catalog_paths = { + path for entry in catalog for path in entry.paths if _in_source_scope(path) + } + surveyed_paths = {row.path for row in rows} + missing_catalog_paths = known_catalog_paths - surveyed_paths + if missing_catalog_paths: + raise SystemExit( + "catalog source paths missing from survey: " + ", ".join(sorted(missing_catalog_paths)) + ) + args.output.write_text( + _render(upstream, revision, tracked, rows, excluded, catalog, coverage), + encoding="utf-8", + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/tilefoundry/cli/source.py b/src/tilefoundry/cli/source.py index 198f8cab..8ed308d0 100644 --- a/src/tilefoundry/cli/source.py +++ b/src/tilefoundry/cli/source.py @@ -270,7 +270,7 @@ def load_namespace(source: str) -> tuple[dict[str, object], str | None]: path, selector = _split_source(source) directory = str(path.parent) sibling_names = { - child.stem for child in path.parent.glob("*.py") + child.stem.partition(".")[0] for child in path.parent.glob("*.py") } | { child.name for child in path.parent.iterdir() diff --git a/tests/fixtures/flashinfer/__init__.py b/tests/fixtures/flashinfer/__init__.py new file mode 100644 index 00000000..4df40af9 --- /dev/null +++ b/tests/fixtures/flashinfer/__init__.py @@ -0,0 +1 @@ +"""Authored HIR fixtures grouped by upstream source repository.""" diff --git a/tests/fixtures/flashinfer/attention_collective.blocked.py b/tests/fixtures/flashinfer/attention_collective.blocked.py new file mode 100644 index 00000000..f86cd276 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_collective.blocked.py @@ -0,0 +1,254 @@ +"""Authored HIR sketches for FlashInfer attention and collectives. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c58fdd2392914ea05e2a8714946ac0eef6 +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: selection/analysis +error: source defines no TileFoundry Module +classification: expected-spec notation; no authored Module is declared. +This expected-spec corpus is intentionally not required to parse. +""" + +from tilefoundry.dsl.tf import * +from tilefoundry.ir.types.shard import Layout, Mesh, Topology + +CTA = Mesh((Topology("cta", 1),), Layout(shape=(1,), strides=(1,))) +THREADS = Mesh((Topology("thread", 128),), Layout(shape=(128,), strides=(1,))) +GPU = Mesh((Topology("gpu", 8),), Layout(shape=(8,), strides=(1,))) + + +def _place(x, scope, storage, tile=None): + + return place(x, scope=scope, storage=storage, tile=tile) + + +def _cta(x, tile=None): + return _place(x, CTA, "smem", tile) + + +def _thread(x, tile=None): + return _place(x, THREADS, "rmem", tile) + + +def _gpu(x, tile=None): + return _place(x, GPU, "gmem", tile) + + +# noqa +def A01_prefill(q, k, v, mask): + return normalize(online_softmax(where(mask, matmul(q, transpose(k, (-1, -2))), -inf), value=v)) + + +def A01_prefill_cta(q, k, v, mask): + return _cta(A01_prefill(q, k, v, mask), (64, 64, 64)) + + +def A01_prefill_thread(q, k, v, mask): + return _thread(A01_prefill(q, k, v, mask), (16, 16, 16)) + + +# noqa +def A02_paged_decode(q, pk, pv, indptr, indices): + return A01_prefill( + q, paged_gather(pk, indptr, indices), paged_gather(pv, indptr, indices), None + ) + + +def A02_paged_decode_cta(q, pk, pv, indptr, indices): + return _cta(A02_paged_decode(q, pk, pv, indptr, indices), (1, 64, 64)) + + +def A02_paged_decode_thread(q, pk, pv, indptr, indices): + return _thread(A02_paged_decode(q, pk, pv, indptr, indices), (1, 16, 16)) + + +# noqa +def A03_softcap(q, k, v): + s = matmul(q, transpose(k, (-1, -2))) * scale + return normalize(online_softmax(cap * tanh(s / cap), value=v)) + + +def A03_softcap_cta(q, k, v): + return _cta(A03_softcap(q, k, v), (64, 64, 64)) + + +def A03_softcap_thread(q, k, v): + return _thread(A03_softcap(q, k, v), (16, 16, 16)) + + +# noqa +def A04_alibi(q, k, v, slope): + return normalize( + online_softmax( + matmul(q, transpose(k, (-1, -2))) * scale + slope * (kv_position - q_position), value=v + ) + ) + + +def A04_alibi_cta(q, k, v, slope): + return _cta(A04_alibi(q, k, v, slope), (64, 64, 64)) + + +def A04_alibi_thread(q, k, v, slope): + return _thread(A04_alibi(q, k, v, slope), (16, 16, 16)) + + +# noqa +def A05_masked(q, k, v, bits): + return normalize( + online_softmax(masked(matmul(q, transpose(k, (-1, -2))) * scale, bits), value=v) + ) + + +def A05_masked_cta(q, k, v, bits): + return _cta(A05_masked(q, k, v, bits), (64, 64, 64)) + + +def A05_masked_thread(q, k, v, bits): + return _thread(A05_masked(q, k, v, bits), (16, 16, 16)) + + +# noqa +def A06_pod(q, k, v, mode, indptr=None, indices=None): + return dynamic_select( + mode, A01_prefill(q, k, v, None), A02_paged_decode(q, k, v, indptr, indices) + ) + + +def A06_pod_cta(q, k, v, mode): + return _cta(A06_pod(q, k, v, mode), (64, 64, 64)) + + +def A06_pod_thread(q, k, v, mode): + return _thread(A06_pod(q, k, v, mode), (16, 16, 16)) + + +# noqa +def A07_merge(o0, m0, d0, o1, m1, d1): + m = maximum(m0, m1) + d = exp(m0 - m) * d0 + exp(m1 - m) * d1 + return (exp(m0 - m) * o0 + exp(m1 - m) * o1) / d, m, d + + +def A07_merge_cta(*xs): + return _cta(A07_merge(*xs)) + + +def A07_merge_thread(*xs): + return _thread(A07_merge(*xs)) + + +# noqa +def C01(x, r): + return rms_norm(all_reduce(x) + r) + + +def C01_cta(x, r): + return _cta(C01(x, r)) + + +def C01_thread(x, r): + return _thread(C01(x, r)) + + +def C01_gpu(x, r): + return _gpu(C01(x, r)) + + +def C02(x, r): + return block_quant(rms_norm(all_reduce(x) + r), format="fp8") + + +def C02_cta(x, r): + return _cta(C02(x, r)) + + +def C02_thread(x, r): + return _thread(C02(x, r)) + + +def C02_gpu(x, r): + return _gpu(C02(x, r)) + + +def C03(x): + return rms_norm(all_reduce(x)) + + +def C03_cta(x): + return _cta(C03(x)) + + +def C03_thread(x): + return _thread(C03(x)) + + +def C03_gpu(x): + return _gpu(C03(x)) + + +def C04(expert, r): + return block_quant(rms_norm(all_reduce(expert_reduce(expert)) + r), format="fp8") + + +def C04_cta(expert, r): + return _cta(C04(expert, r)) + + +def C04_thread(expert, r): + return _thread(C04(expert, r)) + + +def C04_gpu(expert, r): + return _gpu(C04(expert, r)) + + +def C05(permuted, shared): + return rms_norm(all_reduce(expert_finalize(permuted) + shared)) + + +def C05_cta(permuted, shared): + return _cta(C05(permuted, shared)) + + +def C05_thread(permuted, shared): + return _thread(C05(permuted, shared)) + + +def C05_gpu(permuted, shared): + return _gpu(C05(permuted, shared)) + + +# noqa +def C06(x, w): + return matmul(all_gather_tiles(x), w) + + +def C06_cta(x, w): + return _cta(C06(x, w), (128, 64, 32)) + + +def C06_thread(x, w): + return _thread(C06(x, w), (16, 16, 16)) + + +def C06_gpu(x, w): + return _gpu(C06(x, w)) + + +# noqa +def C07(a, b): + return two_shot_all_reduce(matmul(a, b), phases=("reduce_scatter", "all_gather")) + + +def C07_cta(a, b): + return _cta(C07(a, b), (128, 128, 32)) + + +def C07_thread(a, b): + return _thread(C07(a, b), (16, 16, 16)) + + +def C07_gpu(a, b): + return _gpu(C07(a, b)) diff --git a/tests/fixtures/flashinfer/attention_csrc_nvfp4_attention_sm120_quantize.blocked.py b/tests/fixtures/flashinfer/attention_csrc_nvfp4_attention_sm120_quantize.blocked.py new file mode 100644 index 00000000..d392f8b9 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_csrc_nvfp4_attention_sm120_quantize.blocked.py @@ -0,0 +1,93 @@ +"""Block-scaled NVFP4 quantization for SM120 attention operands. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +The specialization keeps 16-token blocks, packed half-width payloads, and one +E4M3 scale per 16 input elements for normal and transposed layouts. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, TOKENS, HEADS, HEAD_DIM = 2, 128, 32, 128 +BLOCK, CTA_COUNT, THREAD_COUNT = 16, BATCH * HEADS * (TOKENS // 16), 128 +TARGET = CudaTarget("nvidia.b200_sxm") +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="quantize", target=TARGET, topologies=TOPOLOGIES) +class ScaledFP4QuantKernel: + """FlashInfer scaled_fp4_quant_kernel entry.""" + + @func + def quantize(x: Tensor[(BATCH, TOKENS, HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, HEADS, TOKENS // BLOCK, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + TOKENS @ (mesh.token_block, mesh.token), + HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + TOKENS @ (mesh.token_block, mesh.token), + HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales + + +@module(entry="quantize", target=TARGET, topologies=TOPOLOGIES) +class ScaledFP4QuantTransKernel: + """FlashInfer scaled_fp4_quant_trans_kernel entry.""" + + kernel = ScaledFP4QuantKernel.renamed("kernel") + + @func + def quantize(x: Tensor[(BATCH, TOKENS, HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, HEADS, TOKENS // BLOCK, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + staged = tf.reshard( + x, + ( + BATCH @ mesh.batch, + TOKENS @ (mesh.token_block, mesh.token), + HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) + packed, scales = kernel(staged) + return ( + tf.transpose(packed, perm=(0, 2, 1, 3)), + tf.transpose(scales, perm=(0, 2, 1, 3)), + ) + + +__all__ = ["ScaledFP4QuantKernel", "ScaledFP4QuantTransKernel"] diff --git a/tests/fixtures/flashinfer/attention_csrc_xqa_mha.blocked.py b/tests/fixtures/flashinfer/attention_csrc_xqa_mha.blocked.py new file mode 100644 index 00000000..0881acbc --- /dev/null +++ b/tests/fixtures/flashinfer/attention_csrc_xqa_mha.blocked.py @@ -0,0 +1,171 @@ +"""Warp-specialized XQA paged multi-head attention. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +csrc/xqa/mha.cu +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +The specialization keeps four KV head groups, eight warps per CTA, paged +context indirection, and online softmax state. +""" + +from __future__ import annotations + +from dataclasses import replace + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, Q_HEADS, KV_HEADS, HEAD_DIM, CONTEXT = 4, 32, 4, 128, 2048 +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 8 * 32 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XQAMHAImpl: + """FlashInfer kernel_mha_impl entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 8, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + 1, + Q_HEADS @ (mesh.kv_head, mesh.warp), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + k_smem = tf.reshard( + k, + ( + BATCH @ mesh.batch, + CONTEXT, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v, + ( + BATCH @ mesh.batch, + CONTEXT, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + pages_reg = tf.reshard( + page_indices, + (BATCH @ mesh.batch, CONTEXT // 16), + "rmem", + ) + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XQAMHA: + """FlashInfer kernel_mha entry.""" + + impl = replace(XQAMHAImpl, name="impl", target=None) + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 8, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_local = tf.reshard( + q, + ( + BATCH @ mesh.batch, + 1, + Q_HEADS @ (mesh.kv_head, mesh.warp), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) + return impl(q_local, k, v, page_indices) + + +__all__ = ["XQAMHAImpl", "XQAMHA"] diff --git a/tests/fixtures/flashinfer/attention_csrc_xqa_mha_sm90.py b/tests/fixtures/flashinfer/attention_csrc_xqa_mha_sm90.py new file mode 100644 index 00000000..f7855853 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_csrc_xqa_mha_sm90.py @@ -0,0 +1,175 @@ +"""SM90 GMMA XQA paged multi-head attention. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +csrc/xqa/mha_sm90.cu +license: Apache-2.0 (no upstream source is vendored) + +The specialization keeps the upstream twelve-warp CTA, GQA head grouping, +split context, and online-softmax output contract. + +The blockwise online-softmax state follows +tests/fixtures/placed/flash_split_k_decode.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, Q_HEADS, KV_HEADS, HEAD_DIM, CONTEXT = 4, 48, 4, 128, 3072 +BLOCK = 64 +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 12 * 32 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XQAMHASM90: + """FlashInfer kernel_mha entry for the SM90 path. + + predicted-ns: 66567 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 12, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + 1, + Q_HEADS @ (mesh.kv_head, mesh.warp), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + 1, + Q_HEADS @ (mesh.kv_head, mesh.warp), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + 1, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + 1, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.warp, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + 1, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.warp, + 1, + BLOCK, + ), + "smem", + ) + score = score * 0.08838834764831845 + block_max = tf.reduce(score, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(score - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out = tf.matmul(weight, value) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + 1, + Q_HEADS @ (mesh.kv_head, mesh.warp), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) + + +__all__ = ["XQAMHASM90"] diff --git a/tests/fixtures/flashinfer/attention_csrc_xqa_mla_sm120.blocked.py b/tests/fixtures/flashinfer/attention_csrc_xqa_mla_sm120.blocked.py new file mode 100644 index 00000000..500725e8 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_csrc_xqa_mla_sm120.blocked.py @@ -0,0 +1,175 @@ +"""SM120 clustered XQA multi-latent attention. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +csrc/xqa/mla_sm120.cu +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +The specialization keeps four CTAs per input token, twelve warps per CTA, +compressed 576-wide KV latent state, and the 128-head MLA output. + +The explicit score and online state follow +tests/fixtures/placed/flash_split_k_decode.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, TOKENS, HEADS, QK_DIM, KV_DIM, CONTEXT = 4, 1, 128, 192, 576, 1920 +BLOCK = 24 +CTA_COUNT, THREAD_COUNT = 4 * BATCH * TOKENS, 12 * 32 +TARGET = CudaTarget("nvidia.b200_sxm") +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XQAMLASM120: + """FlashInfer kernel_mha entry for the SM120 MLA path.""" + + @func + def run( + q: Tensor[(BATCH, TOKENS, HEADS, QK_DIM), "bf16"], + kv: Tensor[(BATCH, CONTEXT, KV_DIM), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + TOKENS, + HEADS @ mesh.head_group, + QK_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + TOKENS, + HEADS @ mesh.head_group, + QK_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=(BATCH, TOKENS, HEADS, 1, QK_DIM), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + TOKENS, + HEADS @ mesh.head_group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + latent = kv[:, start : start + BLOCK, :] + key_smem = tf.reshard( + latent[:, :, :QK_DIM], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + QK_DIM @ mesh.lane, + ), + "smem", + ) + value_smem = tf.reshard( + latent[:, :, QK_DIM : 2 * QK_DIM], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + QK_DIM @ mesh.lane, + ), + "smem", + ) + key = tf.reshape( + tf.cast(key_smem, "f32"), + new_shape=(BATCH, 1, 1, BLOCK, QK_DIM), + ) + value = tf.reshape( + tf.cast(value_smem, "f32"), + new_shape=(BATCH, 1, 1, BLOCK, QK_DIM), + ) + score_partial = tf.reduce( + query * key, axes=(-1,), keepdim=True, kind="sum" + ) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + TOKENS, + HEADS @ mesh.head_group, + BLOCK @ mesh.warp, + 1, + ), + "smem", + ) + score = score * 0.07216878364870322 + block_max = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(score - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-2,), keepdim=True, kind="sum" + ) + block_out_partial = tf.reduce( + weight * value, axes=(-2,), keepdim=False, kind="sum" + ) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + TOKENS, + HEADS @ mesh.head_group, + QK_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + tf.reshape( + block_out, + new_shape=(BATCH, TOKENS, HEADS, 1, QK_DIM), + ) + running_max = next_max + running_sum = next_sum + running_out = next_out + + return tf.reshard( + tf.cast( + tf.reshape( + running_out / running_sum, + new_shape=(BATCH, TOKENS, HEADS, QK_DIM), + ), + "bf16", + ), + ( + BATCH @ mesh.batch, + TOKENS, + HEADS @ mesh.head_group, + QK_DIM @ mesh.lane, + ), + "gmem", + ) + + +__all__ = ["XQAMLASM120"] diff --git a/tests/fixtures/flashinfer/attention_include_attention_batch_pod.py b/tests/fixtures/flashinfer/attention_include_attention_batch_pod.py new file mode 100644 index 00000000..64e035dc --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_batch_pod.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/batch_pod.cuh +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchPODWithKVCacheTensorKernelModule1: + """FlashInfer BatchPODWithKVCacheTensorKernel entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_include_attention_blackwell_plan.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_blackwell_plan.blocked.py new file mode 100644 index 00000000..9c181e06 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_blackwell_plan.blocked.py @@ -0,0 +1,44 @@ +"""Data-dependent attention planning and CTA work dispatch. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/blackwell/plan.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.dynamic_cta_dispatch' (4 positional, no keywords) +ledger: OP-12 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class PlanKernelModule1: + """FlashInfer plan_kernel entry.""" + + @func + def run( + qo_indptr: Tensor[(BATCH + 1,), "i32"], + kv_indptr: Tensor[(BATCH + 1,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 4, 32), + names=("request", "head", "warp", "lane"), + ) as mesh: + qo_reg = tf.reshard(qo_indptr, (BATCH + 1,), "rmem") + kv_smem = tf.reshard(kv_indptr, (BATCH + 1,), "smem") + return tf.dynamic_cta_dispatch(qo_reg, kv_smem, mesh.request, mesh.head) diff --git a/tests/fixtures/flashinfer/attention_include_attention_decode.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_decode.blocked.py new file mode 100644 index 00000000..67d0cfff --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_decode.blocked.py @@ -0,0 +1,284 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/decode.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py, tests/fixtures/placed/mha_decode_paged.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) +MLA_TOPOLOGIES = (Topology("cta", BATCH * 4), Topology("thread", 12 * 32)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SingleDecodeWithKVCacheKernelModule1: + """FlashInfer SingleDecodeWithKVCacheKernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeWithPagedKVCacheKernelModule2: + """FlashInfer BatchDecodeWithPagedKVCacheKernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=MLA_TOPOLOGIES) +class BatchDecodeWithPagedKVCacheKernelMLAModule3: + """FlashInfer BatchDecodeWithPagedKVCacheKernelMLA entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_include_attention_decode_mla_cute_sm80.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_decode_mla_cute_sm80.blocked.py new file mode 100644 index 00000000..4a8869f9 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_decode_mla_cute_sm80.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/decode_mla_cute_sm80.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeWithPagedKVCacheKernelMlaCuteSM80Module1: + """FlashInfer BatchDecodeWithPagedKVCacheKernelMlaCuteSM80 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_include_attention_hopper_prefill_sm90.py b/tests/fixtures/flashinfer/attention_include_attention_hopper_prefill_sm90.py new file mode 100644 index 00000000..4d2b56bc --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_hopper_prefill_sm90.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/hopper/prefill_sm90.cuh +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class PrefillWithKVCacheKernelModule1: + """FlashInfer PrefillWithKVCacheKernel entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_include_attention_hopper_quantization_prefill_sm90.py b/tests/fixtures/flashinfer/attention_include_attention_hopper_quantization_prefill_sm90.py new file mode 100644 index 00000000..92ff7334 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_hopper_quantization_prefill_sm90.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/hopper/quantization/prefill_sm90.cuh +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class FP8PrefillWithKVCacheKernelModule1: + """FlashInfer FP8PrefillWithKVCacheKernel entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_include_attention_mla.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_mla.blocked.py new file mode 100644 index 00000000..095c22e4 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_mla.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/mla.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchMLAPagedAttentionKernelModule1: + """FlashInfer BatchMLAPagedAttentionKernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_include_attention_mla_hopper.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_mla_hopper.blocked.py new file mode 100644 index 00000000..c784cb51 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_mla_hopper.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/mla_hopper.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchMLAPageAttentionHopperKernelModule1: + """FlashInfer BatchMLAPageAttentionHopperKernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_include_attention_persistent_template.py b/tests/fixtures/flashinfer/attention_include_attention_persistent_template.py new file mode 100644 index 00000000..56f8e6c4 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_persistent_template.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/persistent_template.cuh +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class PersistentKernelTemplateModule1: + """FlashInfer PersistentKernelTemplate entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_include_attention_pod.py b/tests/fixtures/flashinfer/attention_include_attention_pod.py new file mode 100644 index 00000000..ac00f9d3 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_pod.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/pod.cuh +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class PODWithKVCacheTensorKernelModule1: + """FlashInfer PODWithKVCacheTensorKernel entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_include_attention_prefill.py b/tests/fixtures/flashinfer/attention_include_attention_prefill.py new file mode 100644 index 00000000..519dc2ab --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_prefill.py @@ -0,0 +1,537 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/prefill.cuh +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SinglePrefillWithKVCacheKernelModule1: + """FlashInfer SinglePrefillWithKVCacheKernel entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchPrefillWithRaggedKVCacheKernelModule2: + """FlashInfer BatchPrefillWithRaggedKVCacheKernel entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchPrefillWithPagedKVCacheKernelModule3: + """FlashInfer BatchPrefillWithPagedKVCacheKernel entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_include_attention_scheduler.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_scheduler.blocked.py new file mode 100644 index 00000000..413e910a --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_scheduler.blocked.py @@ -0,0 +1,82 @@ +"""Data-dependent attention planning and CTA work dispatch. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/scheduler.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.dynamic_cta_dispatch' (4 positional, no keywords) +ledger: OP-12 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeWithPagedKVCacheKernelModule1: + """FlashInfer BatchDecodeWithPagedKVCacheKernel entry.""" + + @func + def run( + qo_indptr: Tensor[(BATCH + 1,), "i32"], + kv_indptr: Tensor[(BATCH + 1,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 4, 32), + names=("request", "head", "warp", "lane"), + ) as mesh: + qo_reg = tf.reshard(qo_indptr, (BATCH + 1,), "rmem") + kv_smem = tf.reshard(kv_indptr, (BATCH + 1,), "smem") + return tf.dynamic_cta_dispatch(qo_reg, kv_smem, mesh.request, mesh.head) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeWithPagedKVCacheKernelMLAModule2: + """FlashInfer BatchDecodeWithPagedKVCacheKernelMLA entry.""" + + @func + def run( + qo_indptr: Tensor[(BATCH + 1,), "i32"], + kv_indptr: Tensor[(BATCH + 1,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 4, 32), + names=("request", "head", "warp", "lane"), + ) as mesh: + qo_reg = tf.reshard(qo_indptr, (BATCH + 1,), "rmem") + kv_smem = tf.reshard(kv_indptr, (BATCH + 1,), "smem") + return tf.dynamic_cta_dispatch(qo_reg, kv_smem, mesh.request, mesh.head) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeWithPagedKVCacheKernelMlaCuteSM80Module3: + """FlashInfer BatchDecodeWithPagedKVCacheKernelMlaCuteSM80 entry.""" + + @func + def run( + qo_indptr: Tensor[(BATCH + 1,), "i32"], + kv_indptr: Tensor[(BATCH + 1,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 4, 32), + names=("request", "head", "warp", "lane"), + ) as mesh: + qo_reg = tf.reshard(qo_indptr, (BATCH + 1,), "rmem") + kv_smem = tf.reshard(kv_indptr, (BATCH + 1,), "smem") + return tf.dynamic_cta_dispatch(qo_reg, kv_smem, mesh.request, mesh.head) diff --git a/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv3_2_kernel.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv3_2_kernel.blocked.py new file mode 100644 index 00000000..381fd82a --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv3_2_kernel.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaDecodeDsv32KernelModule1: + """FlashInfer sparse_mla_decode_dsv3_2_kernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv4_kernel.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv4_kernel.blocked.py new file mode 100644 index 00000000..aa104d8e --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_decode_dsv4_kernel.blocked.py @@ -0,0 +1,177 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaDecodeDsv4KernelModule1: + """FlashInfer sparse_mla_decode_dsv4_kernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaDecodeDsv4MergeKernelModule2: + """FlashInfer sparse_mla_decode_dsv4_merge_kernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_prefill_kernel.blocked.py b/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_prefill_kernel.blocked.py new file mode 100644 index 00000000..c9e69a14 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_attention_sparse_mla_sm120_prefill_kernel.blocked.py @@ -0,0 +1,329 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaPrefillKernelModule1: + """FlashInfer sparse_mla_prefill_kernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaPrefillMgKernelModule2: + """FlashInfer sparse_mla_prefill_mg_kernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaPrefillMgDualKernelModule3: + """FlashInfer sparse_mla_prefill_mg_dual_kernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaPrefillMgDualFulltileKernelModule4: + """FlashInfer sparse_mla_prefill_mg_dual_fulltile_kernel entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_include_page.blocked.py b/tests/fixtures/flashinfer/attention_include_page.blocked.py new file mode 100644 index 00000000..01db8411 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_include_page.blocked.py @@ -0,0 +1,197 @@ +"""Indexed paged-cache append with functional cache ownership. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 include/flashinfer/page.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.scatter_update' (3 positional, keywords ['axis']) +ledger: OP-08 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AppendPagedKVCacheDecodeKernelModule1: + """FlashInfer AppendPagedKVCacheDecodeKernel entry.""" + + @func + def run( + cache: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + values: Tensor[(BATCH, 1, KV_HEADS, HEAD_DIM), "bf16"], + slots: Tensor[(BATCH,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + cache_smem = tf.reshard( + cache, + ( + BATCH @ mesh.batch, + CONTEXT @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + values_reg = tf.reshard( + values, + (BATCH @ mesh.batch, 1, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + slots_reg = tf.reshard(slots, (BATCH @ mesh.batch,), "rmem") + return tf.scatter_update(cache_smem, slots_reg, values_reg, axis=1) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AppendPagedKVCacheKernelModule2: + """FlashInfer AppendPagedKVCacheKernel entry.""" + + @func + def run( + cache: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + values: Tensor[(BATCH, 1, KV_HEADS, HEAD_DIM), "bf16"], + slots: Tensor[(BATCH,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + cache_smem = tf.reshard( + cache, + ( + BATCH @ mesh.batch, + CONTEXT @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + values_reg = tf.reshard( + values, + (BATCH @ mesh.batch, 1, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + slots_reg = tf.reshard(slots, (BATCH @ mesh.batch,), "rmem") + return tf.scatter_update(cache_smem, slots_reg, values_reg, axis=1) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class NVFP4QuantizeAppendPagedKVCacheKernelModule3: + """FlashInfer NVFP4QuantizeAppendPagedKVCacheKernel entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class NVFP4QuantizeAppendPagedKVCacheWithSlotMappingKernelModule4: + """FlashInfer NVFP4QuantizeAppendPagedKVCacheWithSlotMappingKernel entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AppendPagedKVMlaCacheKernelModule5: + """FlashInfer AppendPagedKVMlaCacheKernel entry.""" + + @func + def run( + cache: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + values: Tensor[(BATCH, 1, KV_HEADS, HEAD_DIM), "bf16"], + slots: Tensor[(BATCH,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + cache_smem = tf.reshard( + cache, + ( + BATCH @ mesh.batch, + CONTEXT @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + values_reg = tf.reshard( + values, + (BATCH @ mesh.batch, 1, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + slots_reg = tf.reshard(slots, (BATCH @ mesh.batch,), "rmem") + return tf.scatter_update(cache_smem, slots_reg, values_reg, axis=1) diff --git a/tests/fixtures/flashinfer/attention_py_attention_core.blocked.py b/tests/fixtures/flashinfer/attention_py_attention_core.blocked.py new file mode 100644 index 00000000..462c47d0 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_attention_core.blocked.py @@ -0,0 +1,203 @@ +"""Holistic mixed prefill/decode attention over a paged KV cache. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/attention/_core.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +The specialization preserves CSR page ownership, GQA head sharing, optional +attention sinks, and the output plus log-sum-exp contract. + +Page-table gather and score normalization follow +tests/fixtures/placed/mha_decode_paged.py. +""" + +from __future__ import annotations + +from dataclasses import replace + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, Q_HEADS, KV_HEADS, HEAD_DIM = 4, 32, 8, 128 +PAGES, PAGE_SIZE, PAGES_PER_REQUEST = 128, 16, 32 +CONTEXT = PAGES_PER_REQUEST * PAGE_SIZE +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchAttentionModule: + """FlashInfer BatchAttention entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k_cache: Tensor[(PAGES, PAGE_SIZE, KV_HEADS, HEAD_DIM), "bf16"], + v_cache: Tensor[(PAGES, PAGE_SIZE, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, PAGES_PER_REQUEST), "i32"], + sink: Tensor[(Q_HEADS,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("request", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.request, + 1, + Q_HEADS @ (mesh.kv_head, mesh.warp), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + k_reg = tf.reshard( + k_cache, + (PAGES, PAGE_SIZE, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + v_reg = tf.reshard( + v_cache, + (PAGES, PAGE_SIZE, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k_reg, + (PAGES, PAGE_SIZE, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v_reg, + (PAGES, PAGE_SIZE, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.request, PAGES_PER_REQUEST), "rmem") + sink_reg = tf.reshard(sink, (Q_HEADS @ (mesh.kv_head, mesh.warp),), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * PAGES_PER_REQUEST,) + ) + k_pages = tf.reshard( + k_smem, + (PAGES, PAGE_SIZE, KV_HEADS, HEAD_DIM), + "gmem", + ) + v_pages = tf.reshard( + v_smem, + (PAGES, PAGE_SIZE, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(k_pages, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + ( + BATCH @ mesh.request, + CONTEXT, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(v_pages, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + ( + BATCH @ mesh.request, + CONTEXT, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=( + BATCH, + 1, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + sink_score = tf.reshape( + sink_reg, + new_shape=(1, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, 1), + ) + peak = tf.max(peak, sink_score) + weight = tf.exp(score - peak) + sink_weight = tf.exp(sink_score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + tf.reduce( + sink_weight, + axes=(-2,), + keepdim=False, + kind="sum", + ) + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + output = tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + lse = tf.reshape( + tf.log(total) + + tf.reshape(peak, new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1)), + new_shape=(BATCH, 1, Q_HEADS), + ) + return output, lse + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchAttentionWithSinkModule: + """FlashInfer BatchAttentionWithAttentionSinkWrapper entry.""" + + kernel = replace(BatchAttentionModule, name="kernel", target=None) + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k_cache: Tensor[(PAGES, PAGE_SIZE, KV_HEADS, HEAD_DIM), "bf16"], + v_cache: Tensor[(PAGES, PAGE_SIZE, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, PAGES_PER_REQUEST), "i32"], + sink: Tensor[(Q_HEADS,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("request", "kv_head", "warp", "lane"), + ) as mesh: + q_local = tf.reshard( + q, + ( + BATCH @ mesh.request, + 1, + Q_HEADS @ (mesh.kv_head, mesh.warp), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) + return kernel(q_local, k_cache, v_cache, page_indices, sink) + + +__all__ = ["BatchAttentionModule", "BatchAttentionWithSinkModule"] diff --git a/tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha.py b/tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha.py new file mode 100644 index 00000000..c0744871 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha.py @@ -0,0 +1,195 @@ +"""Placed ragged prefill attention written from existing HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/attention/cute_dsl/fmha.py +license: Apache-2.0 (no upstream source is vendored) + +The score, mask, normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py rather than a black-box op. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CuteDslFmhaRaggedPrefillModule1: + """FlashInfer cute_dsl_fmha_ragged_prefill entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha_blockscaled.blocked.py b/tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha_blockscaled.blocked.py new file mode 100644 index 00000000..5d2a0939 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_attention_cute_dsl_fmha_blockscaled.blocked.py @@ -0,0 +1,60 @@ +"""Packed block-scaled attention payload and scale production. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/attention/cute_dsl/fmha_blockscaled.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS * (QUERY // 16), 8 * 16 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CuteDslFmhaBlockscaledPrefillModule1: + """FlashInfer cute_dsl_fmha_blockscaled_prefill entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales diff --git a/tests/fixtures/flashinfer/attention_py_cudnn_decode.blocked.py b/tests/fixtures/flashinfer/attention_py_cudnn_decode.blocked.py new file mode 100644 index 00000000..12a7ed64 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cudnn_decode.blocked.py @@ -0,0 +1,116 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cudnn/decode.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/mha_decode_paged.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CudnnBatchDecodeWithKvCacheModule1: + """FlashInfer cudnn_batch_decode_with_kv_cache entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) diff --git a/tests/fixtures/flashinfer/attention_py_cudnn_prefill.py b/tests/fixtures/flashinfer/attention_py_cudnn_prefill.py new file mode 100644 index 00000000..23e65cbd --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cudnn_prefill.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cudnn/prefill.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CudnnBatchPrefillWithKvCacheModule1: + """FlashInfer cudnn_batch_prefill_with_kv_cache entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_dsa_hca_fp8.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_dsa_hca_fp8.blocked.py new file mode 100644 index 00000000..3b601996 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_dsa_hca_fp8.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/dsa/hca_fp8.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellHeavilyCompressedAttentionForwardFP8Module1: + """FlashInfer BlackwellHeavilyCompressedAttentionForwardFP8 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha.py new file mode 100644 index 00000000..0f18478b --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/fmha/fmha.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellFusedMultiHeadAttentionForwardModule1: + """FlashInfer BlackwellFusedMultiHeadAttentionForward entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha_blockscaled.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha_blockscaled.blocked.py new file mode 100644 index 00000000..dee2a35c --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_fmha_blockscaled.blocked.py @@ -0,0 +1,60 @@ +"""Packed block-scaled attention payload and scale production. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/fmha/fmha_blockscaled.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS * (QUERY // 16), 8 * 16 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellFusedMultiHeadBlockScaledAttentionForwardModule1: + """FlashInfer BlackwellFusedMultiHeadBlockScaledAttentionForward entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_quantize.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_quantize.blocked.py new file mode 100644 index 00000000..c6aa6c52 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fmha_quantize.blocked.py @@ -0,0 +1,60 @@ +"""Packed block-scaled attention payload and scale production. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/fmha/quantize.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS * (QUERY // 16), 8 * 16 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class QuantizeBlockscaledQkModule1: + """FlashInfer quantize_blockscaled_qk entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_mask.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_mask.py new file mode 100644 index 00000000..fba55081 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_mask.py @@ -0,0 +1,76 @@ +"""Placed band-mask construction used by FlashInfer CuTe attention. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/fusion/mask.py +license: Apache-2.0 (no upstream source is vendored) + +CTA query rows and thread-partitioned KV columns retain causal, left-window, +right-window, and sequence-tail bounds before materializing the mask. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 2, 128, 2048 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", BATCH * QUERY), Topology("thread", 4 * 32)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AttentionMaskModule: + """FlashInfer AttentionMask entry. + + predicted-ns: 1952 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run(scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + score_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + q_pos = tf.reshape( + tf.arange(Tensor[(QUERY,), "i64"]), + new_shape=(1, QUERY, 1), + ) + k_pos = tf.reshape( + tf.arange(Tensor[(CONTEXT,), "i64"]), + new_shape=(1, 1, CONTEXT), + ) + q_reg = tf.reshard( + q_pos, + (1, QUERY @ mesh.query, 1), + "rmem", + ) + k_reg = tf.reshard( + k_pos, + (1, 1, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + offset = CONTEXT - QUERY + visible = (k_reg <= q_reg + offset) and (k_reg >= q_reg + offset - 512) + masked = tf.where(visible, score_reg, -1000000.0) + mask_smem = tf.reshard( + masked, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + return tf.reshard( + mask_smem, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +__all__ = ["AttentionMaskModule"] diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.blocked.py new file mode 100644 index 00000000..b62bda0b --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.blocked.py @@ -0,0 +1,109 @@ +"""Placed positional-bias transforms refused by current local contracts. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/fusion/variant.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: Binary: operands have conflicting storage (rmem, smem); a multi-input op requires its concrete operands to share one residency +ledger: REG-11, EXT-03, FE-07 + +ALiBi reaches the recorded mixed-residency error first. Selecting RPEAttention +independently reaches the preserved table gather, then exceeds smem capacity: +value v4 needs 262144 B in smem, above the target's 232448 B limit. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 2, 128, 2048 +HEADS = 32 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", BATCH * QUERY), Topology("thread", 4 * 32)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ALiBiAttentionModule: + """FlashInfer ALiBiAttention entry with per-head linear position bias.""" + + @func + def run( + scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"], + slopes: ConstTensor[(HEADS,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + positions = tf.reshape( + tf.arange(Tensor[(CONTEXT,), "i64"]), + new_shape=(1, 1, CONTEXT), + ) + positions_reg = tf.reshard( + positions, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + slope_smem = tf.reshard(slopes, (HEADS @ mesh.lane,), "smem") + slope = tf.reduce(slope_smem, axes=(0,), keepdim=True, kind="max") + biased = scores_smem + tf.cast(positions_reg, "f32") * slope + probabilities = tf.softmax(biased, axis=-1) + return tf.reshard( + probabilities, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class RPEAttentionModule: + """FlashInfer RPEAttention entry with a learned relative-position table.""" + + @func + def run( + scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"], + table: ConstTensor[(HEADS, 129), "f32"], + relative_indices: Tensor[(CONTEXT,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + table_smem = tf.reshard(table, (HEADS @ mesh.lane, 129), "smem") + indices_reg = tf.reshard(relative_indices, (CONTEXT @ mesh.warp,), "rmem") + bias = tf.index_select(table_smem, indices_reg, dim=1) + biased = scores_smem + tf.reduce(bias, axes=(0,), keepdim=True, kind="max") + probabilities = tf.softmax(biased, axis=-1) + return tf.reshard( + probabilities, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +__all__ = ["ALiBiAttentionModule", "RPEAttentionModule"] diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.py new file mode 100644 index 00000000..ad66aabf --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_fusion_variant.py @@ -0,0 +1,240 @@ +"""Placed score and statistics transforms for runnable CuTe attention variants. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/fusion/variant.py +license: Apache-2.0 (no upstream source is vendored) + +CTA query rows and thread-partitioned context columns preserve each entry's +default, sink, sigmoid, or soft-capping score mechanism. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 2, 128, 2048 +HEADS = 32 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", BATCH * QUERY), Topology("thread", 4 * 32)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AttentionVariantModule: + """FlashInfer AttentionVariant entry with its default softmax behavior. + + predicted-ns: 914 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run(scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + probabilities = tf.softmax(scores_smem, axis=-1) + return tf.reshard( + probabilities, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class StandardAttentionModule: + """FlashInfer StandardAttention entry with unmodified softmax logits. + + predicted-ns: 914 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run(scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + probabilities = tf.softmax(scores_smem, axis=-1) + return tf.reshard( + probabilities, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AttentionWithSinkModule: + """FlashInfer AttentionWithSink entry with one learned sink per head. + + predicted-ns: 934 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"], + sinks: ConstTensor[(HEADS,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + sink_smem = tf.reshard(sinks, (HEADS @ mesh.lane,), "smem") + sink_stat = tf.reduce(sink_smem, axes=(0,), keepdim=True, kind="max") + probabilities = tf.softmax(scores_smem - sink_stat, axis=-1) + return tf.reshard( + probabilities, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SigmoidAttentionModule: + """FlashInfer SigmoidAttention entry using exp2 and reciprocal logits. + + predicted-ns: 1064 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run(scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + weights = 1.0 / (1.0 + tf.exp2(-scores_smem)) + return tf.reshard( + weights, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SigmoidTanhAttentionModule: + """FlashInfer SigmoidTanhAttention entry using the tanh sigmoid identity. + + predicted-ns: 944 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run(scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + weights = 0.5 + 0.5 * tf.tanh(scores_smem * 0.5) + return tf.reshard( + weights, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SoftCappingAttentionModule: + """FlashInfer SoftCappingAttention entry with bounded tanh logits. + + predicted-ns: 944 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run(scores: Tensor[(BATCH, QUERY, CONTEXT), "f32"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, QUERY, 4, 32), + names=("batch", "query", "warp", "lane"), + ) as mesh: + scores_reg = tf.reshard( + scores, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "rmem", + ) + scores_smem = tf.reshard( + scores_reg, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "smem", + ) + capped = 50.0 * tf.tanh(scores_smem / 50.0) + probabilities = tf.softmax(capped, axis=-1) + return tf.reshard( + probabilities, + (BATCH @ mesh.batch, QUERY @ mesh.query, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +__all__ = [ + "AttentionVariantModule", + "AttentionWithSinkModule", + "SigmoidAttentionModule", + "SigmoidTanhAttentionModule", + "SoftCappingAttentionModule", + "StandardAttentionModule", +] diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode.blocked.py new file mode 100644 index 00000000..f52d7e3e --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode.blocked.py @@ -0,0 +1,116 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/gqa_decode.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/mha_decode_paged.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class GroupedQueryAttentionDecodeModule1: + """FlashInfer GroupedQueryAttentionDecode entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode_paged.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode_paged.blocked.py new file mode 100644 index 00000000..bf9808a2 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_gqa_decode_paged.blocked.py @@ -0,0 +1,116 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/gqa_decode_paged.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/mha_decode_paged.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class GroupedQueryAttentionDecodePagedModule1: + """FlashInfer GroupedQueryAttentionDecodePaged entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode.blocked.py new file mode 100644 index 00000000..e08ce9f1 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/mla_decode.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellMultiLatentAttentionForwardModule1: + """FlashInfer BlackwellMultiLatentAttentionForward entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode_fp8.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode_fp8.blocked.py new file mode 100644 index 00000000..ed3fc808 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_decode_fp8.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/mla_decode_fp8.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellMultiLatentAttentionForwardFP8Module1: + """FlashInfer BlackwellMultiLatentAttentionForwardFP8 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_dispatch.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_dispatch.blocked.py new file mode 100644 index 00000000..eb6d0498 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_mla_dispatch.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/mla_dispatch.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CuteDslMlaDecodeModule1: + """FlashInfer cute_dsl_mla_decode entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode.blocked.py new file mode 100644 index 00000000..7a88ceb2 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/monolithic/mla_decode.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CuteDslMlaDecodeModule1: + """FlashInfer cute_dsl_mla_decode entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp16.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp16.blocked.py new file mode 100644 index 00000000..f1bcf870 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp16.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellMultiHeadLatentAttentionForwardFP16Module1: + """FlashInfer BlackwellMultiHeadLatentAttentionForwardFP16 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp8.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp8.blocked.py new file mode 100644 index 00000000..58836dcf --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_monolithic_mla_decode_fp8.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellMultiHeadLatentAttentionForwardFP8Module1: + """FlashInfer BlackwellMultiHeadLatentAttentionForwardFP8 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_prefill.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_prefill.py new file mode 100644 index 00000000..00a2c589 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_prefill.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/prefill.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BlackwellFusedMultiHeadAttentionForwardModule1: + """FlashInfer BlackwellFusedMultiHeadAttentionForward entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_decode.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_decode.blocked.py new file mode 100644 index 00000000..062aa8b8 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_decode.blocked.py @@ -0,0 +1,207 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/wrappers/batch_decode.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/mha_decode_paged.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeCuteDSLWrapperModule1: + """FlashInfer BatchDecodeCuteDSLWrapper entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodePagedCuteDSLWrapperModule2: + """FlashInfer BatchDecodePagedCuteDSLWrapper entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_hca.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_hca.blocked.py new file mode 100644 index 00000000..72b454dd --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_hca.blocked.py @@ -0,0 +1,101 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/wrappers/batch_hca.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CuteDslHcaDecodeModule1: + """FlashInfer cute_dsl_hca_decode entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_mla.blocked.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_mla.blocked.py new file mode 100644 index 00000000..07948f5f --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_mla.blocked.py @@ -0,0 +1,177 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/wrappers/batch_mla.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchMLADecodeCuteDSLWrapperModule1: + """FlashInfer BatchMLADecodeCuteDSLWrapper entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CuteDslMlaDecodeModule2: + """FlashInfer cute_dsl_mla_decode entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_prefill.py b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_prefill.py new file mode 100644 index 00000000..e47306a2 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_cute_dsl_attention_wrappers_batch_prefill.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/cute_dsl/attention/wrappers/batch_prefill.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchPrefillCuteDSLWrapperModule1: + """FlashInfer BatchPrefillCuteDSLWrapper entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_decode.blocked.py b/tests/fixtures/flashinfer/attention_py_decode.blocked.py new file mode 100644 index 00000000..36465b88 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_decode.blocked.py @@ -0,0 +1,757 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/decode.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py, tests/fixtures/placed/mha_decode_paged.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * KV_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SingleDecodeWithKvCacheWithJitModuleModule1: + """FlashInfer single_decode_with_kv_cache_with_jit_module entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SingleDecodeWithKvCacheModule2: + """FlashInfer single_decode_with_kv_cache entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeWithPagedKVCacheWrapperModule3: + """FlashInfer BatchDecodeWithPagedKVCacheWrapper entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CUDAGraphBatchDecodeWithPagedKVCacheWrapperModule4: + """FlashInfer CUDAGraphBatchDecodeWithPagedKVCacheWrapper entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchDecodeMlaWithPagedKVCacheWrapperModule5: + """FlashInfer BatchDecodeMlaWithPagedKVCacheWrapper entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmGenDecodeModuleModule6: + """FlashInfer TrtllmGenDecodeModule entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmBatchDecodeWithKvCacheModule7: + """FlashInfer trtllm_batch_decode_with_kv_cache entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XqaBatchDecodeWithKvCacheModule8: + """FlashInfer xqa_batch_decode_with_kv_cache entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + page_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (BATCH @ mesh.batch, 1, Q_HEADS @ (mesh.kv_head, mesh.warp), HEAD_DIM @ mesh.lane), + "rmem", + ) + k_smem = tf.reshard( + k, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v, + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "smem", + ) + pages_reg = tf.reshard(page_indices, (BATCH @ mesh.batch, CONTEXT // 16), "rmem") + flat_pages = tf.reshape( + pages_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + paged_k = tf.reshard( + tf.reshape( + k_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + paged_v = tf.reshard( + tf.reshape( + v_smem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ), + (BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + "gmem", + ) + selected_k = tf.reshard( + tf.reshape( + tf.index_select(paged_k, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + tf.reshape( + tf.index_select(paged_v, flat_pages, dim=0), + new_shape=(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), + ), + (BATCH @ mesh.batch, CONTEXT, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), + new_shape=(BATCH, 1, KV_HEADS, Q_HEADS // KV_HEADS, 1, HEAD_DIM), + ) + key = tf.reshape( + tf.transpose(tf.cast(selected_k, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + value = tf.reshape( + tf.transpose(tf.cast(selected_v, "f32"), perm=(0, 2, 1, 3)), + new_shape=(BATCH, 1, KV_HEADS, 1, CONTEXT, HEAD_DIM), + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast( + tf.reshape(blended / total, new_shape=(BATCH, 1, Q_HEADS, HEAD_DIM)), + "bf16", + ) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class FastDecodePlanModule9: + """FlashInfer fast_decode_plan entry.""" + + @func + def run( + qo_indptr: Tensor[(BATCH + 1,), "i32"], + kv_indptr: Tensor[(BATCH + 1,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 4, 32), + names=("request", "head", "warp", "lane"), + ) as mesh: + qo_reg = tf.reshard(qo_indptr, (BATCH + 1,), "rmem") + kv_smem = tf.reshard(kv_indptr, (BATCH + 1,), "smem") + return tf.dynamic_cta_dispatch(qo_reg, kv_smem, mesh.request, mesh.head) diff --git a/tests/fixtures/flashinfer/attention_py_mla_core.blocked.py b/tests/fixtures/flashinfer/attention_py_mla_core.blocked.py new file mode 100644 index 00000000..de7a31d3 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_mla_core.blocked.py @@ -0,0 +1,557 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/mla/_core.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ConvertCompressedPageAlignedSparseIndicesToHcaMetadataModule1: + """FlashInfer convert_compressed_page_aligned_sparse_indices_to_hca_metadata entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmBatchDecodeSparseMlaDsv4Module2: + """FlashInfer trtllm_batch_decode_sparse_mla_dsv4 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchMLAPagedAttentionWrapperModule3: + """FlashInfer BatchMLAPagedAttentionWrapper entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmGenMlaDecodeRunnerModule4: + """FlashInfer TrtllmGenMlaDecodeRunner entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class CuteDslMlaDecodeRunnerModule5: + """FlashInfer CuteDslMlaDecodeRunner entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmBatchDecodeWithKvCacheMlaModule6: + """FlashInfer trtllm_batch_decode_with_kv_cache_mla entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XqaBatchDecodeWithKvCacheMlaModule7: + """FlashInfer xqa_batch_decode_with_kv_cache_mla entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_mla_sparse_mla_sm120.blocked.py b/tests/fixtures/flashinfer/attention_py_mla_sparse_mla_sm120.blocked.py new file mode 100644 index 00000000..f482b37b --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_mla_sparse_mla_sm120.blocked.py @@ -0,0 +1,177 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/mla/_sparse_mla_sm120.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/derived_prefill.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * 4, 12 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaSm120DecodeDsv32Module1: + """FlashInfer sparse_mla_sm120_decode_dsv3_2 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseMlaSm120DecodeDsv4Module2: + """FlashInfer sparse_mla_sm120_decode_dsv4 entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") diff --git a/tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_decode_sm12x.py b/tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_decode_sm12x.py new file mode 100644 index 00000000..59946f48 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_decode_sm12x.py @@ -0,0 +1,124 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/msa_ops/cute_dsl/sparse_decode_sm12x.py +license: Apache-2.0 (no upstream source is vendored) + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/gqa_decode.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparseDecodeForwardSm12xModule1: + """FlashInfer SparseDecodeForwardSm12x entry. + + predicted-ns: 196255608 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + block_indices: Tensor[(BATCH, QUERY, 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 8, 32), + names=("batch", "head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, QUERY, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), "rmem" + ) + k_gmem = tf.reshard( + k, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + v_gmem = tf.reshard( + v, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + index_reg = tf.reshard(block_indices, (BATCH @ mesh.batch, QUERY, 16), "rmem") + blocked_k = tf.reshape( + k_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + blocked_v = tf.reshape( + v_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + output = tf.zeros(Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]) + + for query_index in range(QUERY): + query_blocks = tf.slice( + index_reg, + (0, query_index, 0), + sizes=(BATCH, 1, 16), + strides=(1, 1, 1), + ) + flat_blocks = tf.reshape(query_blocks, new_shape=(BATCH * 16,)) + selected_k = tf.reshape( + tf.index_select(blocked_k, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_v = tf.reshape( + tf.index_select(blocked_v, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_k = tf.repeat_interleave( + selected_k, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_v = tf.repeat_interleave( + selected_v, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_k = tf.reshard( + selected_k, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + selected_v, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.slice( + q_reg, + (0, query_index, 0, 0), + sizes=(BATCH, 1, Q_HEADS, HEAD_DIM), + strides=(1, 1, 1, 1), + ) + query = tf.reshape( + tf.cast(query, "f32"), + new_shape=(BATCH, 1, 1, Q_HEADS, HEAD_DIM), + ) + key = tf.cast(selected_k, "f32") + value = tf.cast(selected_v, "f32") + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-3,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-3,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-3,), keepdim=False, kind="sum") + attended = tf.cast(blended / total, "bf16") + attended = tf.reshard( + attended, + (BATCH @ mesh.batch, 1, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "gmem", + ) + output = tf.insert_slice(output, attended, (0, query_index, 0, 0)) + + return output diff --git a/tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_prefill_sm12x.py b/tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_prefill_sm12x.py new file mode 100644 index 00000000..3d314276 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_msa_ops_cute_dsl_sparse_prefill_sm12x.py @@ -0,0 +1,124 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/msa_ops/cute_dsl/sparse_prefill_sm12x.py +license: Apache-2.0 (no upstream source is vendored) + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/gqa_decode.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SparsePrefillSm12xModule1: + """FlashInfer SparsePrefillSm12x entry. + + predicted-ns: 196255608 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + block_indices: Tensor[(BATCH, QUERY, 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 8, 32), + names=("batch", "head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, QUERY, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), "rmem" + ) + k_gmem = tf.reshard( + k, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + v_gmem = tf.reshard( + v, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + index_reg = tf.reshard(block_indices, (BATCH @ mesh.batch, QUERY, 16), "rmem") + blocked_k = tf.reshape( + k_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + blocked_v = tf.reshape( + v_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + output = tf.zeros(Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]) + + for query_index in range(QUERY): + query_blocks = tf.slice( + index_reg, + (0, query_index, 0), + sizes=(BATCH, 1, 16), + strides=(1, 1, 1), + ) + flat_blocks = tf.reshape(query_blocks, new_shape=(BATCH * 16,)) + selected_k = tf.reshape( + tf.index_select(blocked_k, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_v = tf.reshape( + tf.index_select(blocked_v, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_k = tf.repeat_interleave( + selected_k, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_v = tf.repeat_interleave( + selected_v, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_k = tf.reshard( + selected_k, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + selected_v, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.slice( + q_reg, + (0, query_index, 0, 0), + sizes=(BATCH, 1, Q_HEADS, HEAD_DIM), + strides=(1, 1, 1, 1), + ) + query = tf.reshape( + tf.cast(query, "f32"), + new_shape=(BATCH, 1, 1, Q_HEADS, HEAD_DIM), + ) + key = tf.cast(selected_k, "f32") + value = tf.cast(selected_v, "f32") + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-3,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-3,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-3,), keepdim=False, kind="sum") + attended = tf.cast(blended / total, "bf16") + attended = tf.reshard( + attended, + (BATCH @ mesh.batch, 1, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "gmem", + ) + output = tf.insert_slice(output, attended, (0, query_index, 0, 0)) + + return output diff --git a/tests/fixtures/flashinfer/attention_py_msa_ops_sparse_decode.py b/tests/fixtures/flashinfer/attention_py_msa_ops_sparse_decode.py new file mode 100644 index 00000000..feb99fe3 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_msa_ops_sparse_decode.py @@ -0,0 +1,124 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/msa_ops/sparse_decode.py +license: Apache-2.0 (no upstream source is vendored) + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/gqa_decode.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class MsaSparseDecodeAttentionModule1: + """FlashInfer msa_sparse_decode_attention entry. + + predicted-ns: 196255608 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + block_indices: Tensor[(BATCH, QUERY, 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 8, 32), + names=("batch", "head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, QUERY, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), "rmem" + ) + k_gmem = tf.reshard( + k, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + v_gmem = tf.reshard( + v, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + index_reg = tf.reshard(block_indices, (BATCH @ mesh.batch, QUERY, 16), "rmem") + blocked_k = tf.reshape( + k_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + blocked_v = tf.reshape( + v_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + output = tf.zeros(Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]) + + for query_index in range(QUERY): + query_blocks = tf.slice( + index_reg, + (0, query_index, 0), + sizes=(BATCH, 1, 16), + strides=(1, 1, 1), + ) + flat_blocks = tf.reshape(query_blocks, new_shape=(BATCH * 16,)) + selected_k = tf.reshape( + tf.index_select(blocked_k, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_v = tf.reshape( + tf.index_select(blocked_v, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_k = tf.repeat_interleave( + selected_k, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_v = tf.repeat_interleave( + selected_v, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_k = tf.reshard( + selected_k, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + selected_v, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.slice( + q_reg, + (0, query_index, 0, 0), + sizes=(BATCH, 1, Q_HEADS, HEAD_DIM), + strides=(1, 1, 1, 1), + ) + query = tf.reshape( + tf.cast(query, "f32"), + new_shape=(BATCH, 1, 1, Q_HEADS, HEAD_DIM), + ) + key = tf.cast(selected_k, "f32") + value = tf.cast(selected_v, "f32") + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-3,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-3,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-3,), keepdim=False, kind="sum") + attended = tf.cast(blended / total, "bf16") + attended = tf.reshard( + attended, + (BATCH @ mesh.batch, 1, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "gmem", + ) + output = tf.insert_slice(output, attended, (0, query_index, 0, 0)) + + return output diff --git a/tests/fixtures/flashinfer/attention_py_msa_ops_sparse_prefill.py b/tests/fixtures/flashinfer/attention_py_msa_ops_sparse_prefill.py new file mode 100644 index 00000000..71740475 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_msa_ops_sparse_prefill.py @@ -0,0 +1,124 @@ +"""Placed indexed attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/msa_ops/sparse_prefill.py +license: Apache-2.0 (no upstream source is vendored) + +Page/block gather, score normalization, and value reduction follow the +primitive structure in tests/fixtures/placed/gqa_decode.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class MsaSparseAttentionModule1: + """FlashInfer msa_sparse_attention entry. + + predicted-ns: 196255608 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + block_indices: Tensor[(BATCH, QUERY, 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 8, 32), + names=("batch", "head", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, QUERY, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), "rmem" + ) + k_gmem = tf.reshard( + k, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + v_gmem = tf.reshard( + v, (BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "gmem" + ) + index_reg = tf.reshard(block_indices, (BATCH @ mesh.batch, QUERY, 16), "rmem") + blocked_k = tf.reshape( + k_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + blocked_v = tf.reshape( + v_gmem, + new_shape=(BATCH * (CONTEXT // 16), 16, KV_HEADS, HEAD_DIM), + ) + output = tf.zeros(Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]) + + for query_index in range(QUERY): + query_blocks = tf.slice( + index_reg, + (0, query_index, 0), + sizes=(BATCH, 1, 16), + strides=(1, 1, 1), + ) + flat_blocks = tf.reshape(query_blocks, new_shape=(BATCH * 16,)) + selected_k = tf.reshape( + tf.index_select(blocked_k, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_v = tf.reshape( + tf.index_select(blocked_v, flat_blocks, dim=0), + new_shape=(BATCH, 1, 16 * 16, KV_HEADS, HEAD_DIM), + ) + selected_k = tf.repeat_interleave( + selected_k, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_v = tf.repeat_interleave( + selected_v, repeats=Q_HEADS // KV_HEADS, axis=3 + ) + selected_k = tf.reshard( + selected_k, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + selected_v = tf.reshard( + selected_v, + (BATCH @ mesh.batch, 1, 16 * 16, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "rmem", + ) + query = tf.slice( + q_reg, + (0, query_index, 0, 0), + sizes=(BATCH, 1, Q_HEADS, HEAD_DIM), + strides=(1, 1, 1, 1), + ) + query = tf.reshape( + tf.cast(query, "f32"), + new_shape=(BATCH, 1, 1, Q_HEADS, HEAD_DIM), + ) + key = tf.cast(selected_k, "f32") + value = tf.cast(selected_v, "f32") + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.08838834764831845) + peak = tf.reduce(score, axes=(-3,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-3,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-3,), keepdim=False, kind="sum") + attended = tf.cast(blended / total, "bf16") + attended = tf.reshard( + attended, + (BATCH @ mesh.batch, 1, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "gmem", + ) + output = tf.insert_slice(output, attended, (0, query_index, 0, 0)) + + return output diff --git a/tests/fixtures/flashinfer/attention_py_nvfp4_attention_sm120.blocked.py b/tests/fixtures/flashinfer/attention_py_nvfp4_attention_sm120.blocked.py new file mode 100644 index 00000000..409bcbff --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_nvfp4_attention_sm120.blocked.py @@ -0,0 +1,95 @@ +"""Packed block-scaled attention payload and scale production. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/nvfp4_attention_sm120.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS * (QUERY // 16), 8 * 16 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class Nvfp4AttentionSm120QuantizeQkvModule1: + """FlashInfer nvfp4_attention_sm120_quantize_qkv entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class Nvfp4AttentionSm120FwdModule2: + """FlashInfer nvfp4_attention_sm120_fwd entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales diff --git a/tests/fixtures/flashinfer/attention_py_page.blocked.py b/tests/fixtures/flashinfer/attention_py_page.blocked.py new file mode 100644 index 00000000..5b02340e --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_page.blocked.py @@ -0,0 +1,197 @@ +"""Packed block-scaled attention payload and scale production. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/page.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS * (QUERY // 16), 8 * 16 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class Nvfp4QuantizeAppendPagedKvCacheModule1: + """FlashInfer nvfp4_quantize_append_paged_kv_cache entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class Nvfp4QuantizeAppendPagedKvCacheWithSlotMappingModule2: + """FlashInfer nvfp4_quantize_append_paged_kv_cache_with_slot_mapping entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, QUERY // 16, 8, 16), + names=("batch", "head", "token_block", "token", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + QUERY @ (mesh.token_block, mesh.token), + Q_HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + packed, scales = tf.quant(tf.cast(x_smem, "f32"), group=16, target_dtype="nvfp4") + return packed, scales + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class GetBatchIndicesPositionsModule3: + """FlashInfer get_batch_indices_positions entry.""" + + @func + def run( + cache: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + values: Tensor[(BATCH, 1, KV_HEADS, HEAD_DIM), "bf16"], + slots: Tensor[(BATCH,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + cache_smem = tf.reshard( + cache, + ( + BATCH @ mesh.batch, + CONTEXT @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + values_reg = tf.reshard( + values, + (BATCH @ mesh.batch, 1, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + slots_reg = tf.reshard(slots, (BATCH @ mesh.batch,), "rmem") + return tf.scatter_update(cache_smem, slots_reg, values_reg, axis=1) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AppendPagedMlaKvCacheModule4: + """FlashInfer append_paged_mla_kv_cache entry.""" + + @func + def run( + cache: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + values: Tensor[(BATCH, 1, KV_HEADS, HEAD_DIM), "bf16"], + slots: Tensor[(BATCH,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + cache_smem = tf.reshard( + cache, + ( + BATCH @ mesh.batch, + CONTEXT @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + values_reg = tf.reshard( + values, + (BATCH @ mesh.batch, 1, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + slots_reg = tf.reshard(slots, (BATCH @ mesh.batch,), "rmem") + return tf.scatter_update(cache_smem, slots_reg, values_reg, axis=1) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AppendPagedKvCacheModule5: + """FlashInfer append_paged_kv_cache entry.""" + + @func + def run( + cache: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + values: Tensor[(BATCH, 1, KV_HEADS, HEAD_DIM), "bf16"], + slots: Tensor[(BATCH,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, 4, 32), + names=("batch", "kv_head", "warp", "lane"), + ) as mesh: + cache_smem = tf.reshard( + cache, + ( + BATCH @ mesh.batch, + CONTEXT @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + values_reg = tf.reshard( + values, + (BATCH @ mesh.batch, 1, KV_HEADS @ mesh.kv_head, HEAD_DIM @ mesh.lane), + "rmem", + ) + slots_reg = tf.reshard(slots, (BATCH @ mesh.batch,), "rmem") + return tf.scatter_update(cache_smem, slots_reg, values_reg, axis=1) diff --git a/tests/fixtures/flashinfer/attention_py_parallel_attention_attention_ops.blocked.py b/tests/fixtures/flashinfer/attention_py_parallel_attention_attention_ops.blocked.py new file mode 100644 index 00000000..9e9cddc6 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_parallel_attention_attention_ops.blocked.py @@ -0,0 +1,44 @@ +"""Data-dependent attention planning and CTA work dispatch. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/parallel_attention/attention_ops.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.dynamic_cta_dispatch' (4 positional, no keywords) +ledger: OP-12 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AttentionOpManagerModule1: + """FlashInfer AttentionOpManager entry.""" + + @func + def run( + qo_indptr: Tensor[(BATCH + 1,), "i32"], + kv_indptr: Tensor[(BATCH + 1,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 4, 32), + names=("request", "head", "warp", "lane"), + ) as mesh: + qo_reg = tf.reshard(qo_indptr, (BATCH + 1,), "rmem") + kv_smem = tf.reshard(kv_indptr, (BATCH + 1,), "smem") + return tf.dynamic_cta_dispatch(qo_reg, kv_smem, mesh.request, mesh.head) diff --git a/tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_attention.py b/tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_attention.py new file mode 100644 index 00000000..195ba258 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_attention.py @@ -0,0 +1,195 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/parallel_attention/parallel_attention.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ParallelAttentionModule1: + """FlashInfer ParallelAttention entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_wrapper.blocked.py b/tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_wrapper.blocked.py new file mode 100644 index 00000000..c8b5c91a --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_parallel_attention_parallel_wrapper.blocked.py @@ -0,0 +1,45 @@ +"""Placed parallel-attention exchange across participant shards. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/parallel_attention/parallel_wrapper.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.all_to_all' (1 positional, keywords ['split_axis', 'concat_axis']) +ledger: OP-10 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class AllToAllModule1: + """FlashInfer all_to_all entry.""" + + @func + def run(x: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"]): + with Mesh( + ("cta", "thread"), + layout=(BATCH, Q_HEADS, 4, 32), + names=("rank", "head", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, (BATCH @ mesh.rank, QUERY, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), "rmem" + ) + x_smem = tf.reshard( + x_reg, (BATCH @ mesh.rank, QUERY, Q_HEADS @ mesh.head, HEAD_DIM @ mesh.lane), "smem" + ) + return tf.all_to_all(x_smem, split_axis=2, concat_axis=1) diff --git a/tests/fixtures/flashinfer/attention_py_pod.py b/tests/fixtures/flashinfer/attention_py_pod.py new file mode 100644 index 00000000..eef608f4 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_pod.py @@ -0,0 +1,366 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/pod.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class PODWithPagedKVCacheWrapperModule1: + """FlashInfer PODWithPagedKVCacheWrapper entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchPODWithPagedKVCacheWrapperModule2: + """FlashInfer BatchPODWithPagedKVCacheWrapper entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_prefill.py b/tests/fixtures/flashinfer/attention_py_prefill.py new file mode 100644 index 00000000..2d73e0e9 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_prefill.py @@ -0,0 +1,1563 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/prefill.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SinglePrefillWithKvCacheWithJitModuleModule1: + """FlashInfer single_prefill_with_kv_cache_with_jit_module entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class SinglePrefillWithKvCacheModule2: + """FlashInfer single_prefill_with_kv_cache entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchPrefillWithPagedKVCacheWrapperModule3: + """FlashInfer BatchPrefillWithPagedKVCacheWrapper entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class BatchPrefillWithRaggedKVCacheWrapperModule4: + """FlashInfer BatchPrefillWithRaggedKVCacheWrapper entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmSageAttentionQuantizeModule5: + """FlashInfer trtllm_sage_attention_quantize entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmRaggedAttentionDeepseekModule6: + """FlashInfer trtllm_ragged_attention_deepseek entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmBatchContextWithKvCacheModule7: + """FlashInfer trtllm_batch_context_with_kv_cache entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class FmhaV2PrefillDeepseekModule8: + """FlashInfer fmha_v2_prefill_deepseek entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class TrtllmFmhaV2PrefillModule9: + """FlashInfer trtllm_fmha_v2_prefill entry. + + predicted-ns: 68434 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) diff --git a/tests/fixtures/flashinfer/attention_py_rope.blocked.py b/tests/fixtures/flashinfer/attention_py_rope.blocked.py new file mode 100644 index 00000000..72d34881 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_rope.blocked.py @@ -0,0 +1,454 @@ +"""Placed rotary-position attention boundary with rank-three inputs. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/rope.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.rope_with_positions' (5 positional, no keywords) +ledger: EXT-02 + +Each SURVEY entry remains an independent Module with the source group's +representative supported specialization and upstream mechanism intact. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 2048 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * QUERY, 4 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class RopeQuantizeFp8Module1: + """FlashInfer rope_quantize_fp8 entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class MlaRopeQuantizeFp8Module2: + """FlashInfer mla_rope_quantize_fp8 entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class RopeQuantizeFp8AppendPagedKvCacheModule3: + """FlashInfer rope_quantize_fp8_append_paged_kv_cache entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyRopeInplaceModule4: + """FlashInfer apply_rope_inplace entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyRopePosIdsInplaceModule5: + """FlashInfer apply_rope_pos_ids_inplace entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyLlama31RopeInplaceModule6: + """FlashInfer apply_llama31_rope_inplace entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyLlama31RopePosIdsInplaceModule7: + """FlashInfer apply_llama31_rope_pos_ids_inplace entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyRopeModule8: + """FlashInfer apply_rope entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyRopePosIdsModule9: + """FlashInfer apply_rope_pos_ids entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyLlama31RopeModule10: + """FlashInfer apply_llama31_rope entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyLlama31RopePosIdsModule11: + """FlashInfer apply_llama31_rope_pos_ids entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyRopeWithCosSinCacheModule12: + """FlashInfer apply_rope_with_cos_sin_cache entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class ApplyRopeWithCosSinCacheInplaceModule13: + """FlashInfer apply_rope_with_cos_sin_cache_inplace entry.""" + + @func + def run( + q: Tensor[(BATCH * QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH * QUERY, KV_HEADS, HEAD_DIM), "bf16"], + cos: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + sin: ConstTensor[(BATCH * QUERY, HEAD_DIM // 2), "f32"], + positions: Tensor[(BATCH * QUERY,), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH * QUERY, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, ((BATCH * QUERY) @ mesh.token, Q_HEADS @ mesh.warp, HEAD_DIM @ mesh.lane), "rmem" + ) + k_smem = tf.reshard( + k, ((BATCH * QUERY) @ mesh.token, KV_HEADS, HEAD_DIM @ mesh.lane), "smem" + ) + cos_smem = tf.reshard( + cos, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + sin_smem = tf.reshard( + sin, ((BATCH * QUERY) @ mesh.token, (HEAD_DIM // 2) @ mesh.lane), "smem" + ) + pos_reg = tf.reshard(positions, ((BATCH * QUERY) @ mesh.token,), "rmem") + return tf.rope_with_positions(q_reg, k_smem, cos_smem, sin_smem, pos_reg) diff --git a/tests/fixtures/flashinfer/attention_py_xqa.blocked.py b/tests/fixtures/flashinfer/attention_py_xqa.blocked.py new file mode 100644 index 00000000..e62d8bba --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_xqa.blocked.py @@ -0,0 +1,104 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/xqa.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: pending primitive rewrite probe +ledger: None + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) +MLA_TOPOLOGIES = (Topology("cta", BATCH * 4), Topology("thread", 12 * 32)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XqaMlaModule2: + """FlashInfer xqa_mla entry.""" + + @func + def run( + q: Tensor[(BATCH, 1, 128, 192), "bf16"], + compressed_kv: Tensor[(BATCH, CONTEXT, 576), "bf16"], + sparse_indices: Tensor[(BATCH, CONTEXT // 16), "i32"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, 4, 12, 32), + names=("batch", "head_group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, (BATCH @ mesh.batch, 1, 128 @ mesh.head_group, 192 @ mesh.lane), "rmem" + ) + kv_smem = tf.reshard( + compressed_kv, (BATCH @ mesh.batch, CONTEXT @ mesh.warp, 576 @ mesh.lane), "smem" + ) + index_reg = tf.reshard( + sparse_indices, (BATCH @ mesh.batch, (CONTEXT // 16) @ mesh.warp), "rmem" + ) + flat_blocks = tf.reshape( + index_reg, new_shape=(BATCH * (CONTEXT // 16),) + ) + blocked_kv = tf.reshard( + tf.reshape( + kv_smem, new_shape=(BATCH * (CONTEXT // 16), 16, 576) + ), + (BATCH * (CONTEXT // 16), 16, 576), + "gmem", + ) + selected_kv = tf.reshape( + tf.index_select(blocked_kv, flat_blocks, dim=0), + new_shape=(BATCH, CONTEXT, 576), + ) + key = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 0), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + value = tf.reshard( + tf.slice( + selected_kv, + (0, 0, 192), + sizes=(BATCH, CONTEXT, 192), + strides=(1, 1, 1), + ), + (BATCH @ mesh.batch, CONTEXT, 192 @ mesh.lane), + "rmem", + ) + query = tf.reshape( + tf.cast(q_reg, "f32"), new_shape=(BATCH, 1, 128, 1, 192) + ) + key = tf.reshape( + tf.cast(key, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + value = tf.reshape( + tf.cast(value, "f32"), new_shape=(BATCH, 1, 1, CONTEXT, 192) + ) + raw = tf.reduce(query * key, axes=(-1,), keepdim=True, kind="sum") + score = raw * tf.full_like(raw, value=0.07216878364870322) + peak = tf.reduce(score, axes=(-2,), keepdim=True, kind="max") + weight = tf.exp(score - peak) + total = tf.reduce(weight, axes=(-2,), keepdim=False, kind="sum") + blended = tf.reduce(weight * value, axes=(-2,), keepdim=False, kind="sum") + return tf.cast(blended / total, "bf16") + diff --git a/tests/fixtures/flashinfer/attention_py_xqa.py b/tests/fixtures/flashinfer/attention_py_xqa.py new file mode 100644 index 00000000..37f583ce --- /dev/null +++ b/tests/fixtures/flashinfer/attention_py_xqa.py @@ -0,0 +1,197 @@ +"""Placed blockwise attention composed from HIR primitives. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 flashinfer/xqa.py +license: Apache-2.0 (no upstream source is vendored) + +The blockwise score, mask, online normalization, and value reduction follow +tests/fixtures/placed/prefill_decode_attention.py. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Mesh, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +BATCH, QUERY, CONTEXT = 4, 128, 1920 +Q_HEADS, KV_HEADS, HEAD_DIM = 32, 8, 128 +BLOCK = 64 +TARGET = CudaTarget("nvidia.h200_sxm") +CTA_COUNT, THREAD_COUNT = BATCH * Q_HEADS, 8 * 32 +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", THREAD_COUNT)) +MLA_TOPOLOGIES = (Topology("cta", BATCH * 4), Topology("thread", 12 * 32)) + + +@module(entry="run", target=TARGET, topologies=TOPOLOGIES) +class XqaModule1: + """FlashInfer xqa entry. + + predicted-ns: 64276 + waves: 1 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def run( + q: Tensor[(BATCH, QUERY, Q_HEADS, HEAD_DIM), "bf16"], + k: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + v: Tensor[(BATCH, CONTEXT, KV_HEADS, HEAD_DIM), "bf16"], + mask: Tensor[(BATCH, QUERY, CONTEXT), "bool"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, KV_HEADS, Q_HEADS // KV_HEADS, 8, 32), + names=("batch", "kv_head", "group", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + query = tf.reshape( + tf.cast(q_smem, "f32"), + new_shape=( + BATCH, + QUERY, + KV_HEADS, + Q_HEADS // KV_HEADS, + 1, + HEAD_DIM, + ), + ) + state_partial = tf.reduce(query, axes=(-1,), keepdim=True, kind="sum") + state = tf.reshard( + state_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + 1, + ), + "smem", + ) + running_max = tf.full_like(state, value=-1e30) + running_sum = tf.full_like(state, value=0.0) + running_out = tf.full_like(query, value=0.0) + + for start in range(0, CONTEXT, BLOCK): + k_smem = tf.reshard( + k[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + v_smem = tf.reshard( + v[:, start : start + BLOCK, :, :], + ( + BATCH @ mesh.batch, + BLOCK @ mesh.warp, + KV_HEADS @ mesh.kv_head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + mask_smem = tf.reshard( + mask[:, :, start : start + BLOCK], + (BATCH @ mesh.batch, QUERY, BLOCK @ mesh.warp), + "smem", + ) + key = tf.transpose( + tf.reshape( + tf.cast(k_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 5, 2), + ) + value = tf.transpose( + tf.reshape( + tf.cast(v_smem, "f32"), + new_shape=(BATCH, 1, BLOCK, KV_HEADS, 1, HEAD_DIM), + ), + perm=(0, 1, 3, 4, 2, 5), + ) + score_partial = tf.matmul(query, key) + score = tf.reshard( + score_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + BLOCK @ mesh.warp, + ), + "smem", + ) + score = score * 0.08838834764831845 + live = tf.where( + tf.reshape( + mask_smem, + new_shape=(BATCH, QUERY, 1, 1, 1, BLOCK), + ), + score, + tf.full_like(score, value=-1e30), + ) + block_max = tf.reduce(live, axes=(-1,), keepdim=True, kind="max") + next_max = tf.max(running_max, block_max) + correction = tf.exp(running_max - next_max) + weight = tf.exp(live - next_max) + next_sum = running_sum * correction + tf.reduce( + weight, axes=(-1,), keepdim=True, kind="sum" + ) + block_out_partial = tf.matmul(weight, value) + block_out = tf.reshard( + block_out_partial, + ( + BATCH @ mesh.batch, + QUERY, + KV_HEADS @ mesh.kv_head, + (Q_HEADS // KV_HEADS) @ mesh.group, + 1, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + next_out = running_out * correction + block_out + running_max = next_max + running_sum = next_sum + running_out = next_out + + output = tf.reshape( + tf.cast(running_out / running_sum, "bf16"), + new_shape=(BATCH, QUERY, Q_HEADS, HEAD_DIM), + ) + return tf.reshard( + output, + ( + BATCH @ mesh.batch, + QUERY, + Q_HEADS @ (mesh.kv_head, mesh.group), + HEAD_DIM @ mesh.lane, + ), + "gmem", + ) + diff --git a/tests/fixtures/flashinfer/attention_state.py b/tests/fixtures/flashinfer/attention_state.py new file mode 100644 index 00000000..063bec72 --- /dev/null +++ b/tests/fixtures/flashinfer/attention_state.py @@ -0,0 +1,91 @@ +"""Placed A07 online attention-state merge boundary. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +include/flashinfer/attention/state.cuh:state_t.merge +license: Apache-2.0 (no upstream source is vendored) + +CTAs own independent state rows. Four warps and 32 lanes split the output +vector, matching the upstream vectorized merge while scalar maxima and sums +remain broadcast within each row. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, WIDTH = 256, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 4 * 32)) + + +@module(entry="merge", target=TARGET, topologies=TOPOLOGIES) +class AttentionStateMerge: + """FlashInfer attention state merge kernel. + + predicted-ns: 126 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def merge( + o0: Tensor[(ROWS, WIDTH), "f32"], + m0: Tensor[(ROWS, 1), "f32"], + d0: Tensor[(ROWS, 1), "f32"], + o1: Tensor[(ROWS, WIDTH), "f32"], + m1: Tensor[(ROWS, 1), "f32"], + d1: Tensor[(ROWS, 1), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 4, 32), + names=("state", "warp", "lane"), + ) as mesh: + o0_reg = tf.reshard( + o0, + (ROWS @ mesh.state, 4 @ mesh.warp, 32 @ mesh.lane), + "rmem", + ) + o1_reg = tf.reshard( + o1, + (ROWS @ mesh.state, 4 @ mesh.warp, 32 @ mesh.lane), + "rmem", + ) + m0_reg = tf.reshard(m0, (ROWS @ mesh.state, 1), "rmem") + m1_reg = tf.reshard(m1, (ROWS @ mesh.state, 1), "rmem") + d0_reg = tf.reshard(d0, (ROWS @ mesh.state, 1), "rmem") + d1_reg = tf.reshard(d1, (ROWS @ mesh.state, 1), "rmem") + a = tf.reshard( + o0_reg, + (ROWS @ mesh.state, 4 @ mesh.warp, 32 @ mesh.lane), + "smem", + ) + b = tf.reshard( + o1_reg, + (ROWS @ mesh.state, 4 @ mesh.warp, 32 @ mesh.lane), + "smem", + ) + ml0 = tf.reshard(m0_reg, (ROWS @ mesh.state, 1), "smem") + ml1 = tf.reshard(m1_reg, (ROWS @ mesh.state, 1), "smem") + dl0 = tf.reshard(d0_reg, (ROWS @ mesh.state, 1), "smem") + dl1 = tf.reshard(d1_reg, (ROWS @ mesh.state, 1), "smem") + maximum = tf.maximum(ml0, ml1) + w0 = tf.exp2(ml0 - maximum) + w1 = tf.exp2(ml1 - maximum) + denominator = w0 * dl0 + w1 * dl1 + output = (w0 * a + w1 * b) / denominator + return ( + tf.reshard( + output, + (ROWS @ mesh.state, 4 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ), + tf.reshard(maximum, (ROWS @ mesh.state, 1), "gmem"), + tf.reshard(denominator, (ROWS @ mesh.state, 1), "gmem"), + ) + + +__all__ = ["AttentionStateMerge"] diff --git a/tests/fixtures/flashinfer/attention_variants.py b/tests/fixtures/flashinfer/attention_variants.py new file mode 100644 index 00000000..d0d3bbfb --- /dev/null +++ b/tests/fixtures/flashinfer/attention_variants.py @@ -0,0 +1,88 @@ +"""Placed A03 logits soft-cap attention boundary. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +include/flashinfer/attention/variants.cuh:DefaultAttention +license: Apache-2.0 (no upstream source is vendored) + +CTAs own query rows while four warps and 32 lanes vectorize each head. This +matches the upstream attention variant's row-parallel, vectorized score path. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, WIDTH = 256, 128 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 4 * 32)) + + +@module(entry="softcap_attention", target=TARGET, topologies=TOPOLOGIES) +class SoftcapAttention: + """FlashInfer DefaultAttention logits soft-cap kernel. + + predicted-ns: 14814 waves: 2 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def softcap_attention( + q: Tensor[(ROWS, WIDTH), "f32"], + k: Tensor[(ROWS, WIDTH), "f32"], + v: Tensor[(ROWS, WIDTH), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 4, 32), + names=("query", "warp", "lane"), + ) as mesh: + q_reg = tf.reshard( + q, + (ROWS @ mesh.query, 4 @ mesh.warp, 32 @ mesh.lane), + "rmem", + ) + k_reg = tf.reshard( + k, + (ROWS, 4 @ mesh.warp, 32 @ mesh.lane), + "rmem", + ) + v_reg = tf.reshard( + v, + (ROWS, 4 @ mesh.warp, 32 @ mesh.lane), + "rmem", + ) + q_smem = tf.reshard( + q_reg, + (ROWS @ mesh.query, 4 @ mesh.warp, 32 @ mesh.lane), + "smem", + ) + k_smem = tf.reshard( + k_reg, + (ROWS, 4 @ mesh.warp, 32 @ mesh.lane), + "smem", + ) + v_smem = tf.reshard( + v_reg, + (ROWS, 4 @ mesh.warp, 32 @ mesh.lane), + "smem", + ) + scores = q_smem @ tf.transpose(k_smem, perm=(1, 0)) + reduced_scores = tf.reshard( + scores, + (ROWS @ mesh.query, ROWS), + "smem", + ) + probabilities = tf.softmax(8.0 * tf.tanh(reduced_scores / 8.0), axis=-1) + output = probabilities @ v_smem + return tf.reshard( + output, + (ROWS @ mesh.query, 4 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ) + + +__all__ = ["SoftcapAttention"] diff --git a/tests/fixtures/flashinfer/cute_dsl_add_rmsnorm_fp4quant.blocked.py b/tests/fixtures/flashinfer/cute_dsl_add_rmsnorm_fp4quant.blocked.py new file mode 100644 index 00000000..12109379 --- /dev/null +++ b/tests/fixtures/flashinfer/cute_dsl_add_rmsnorm_fp4quant.blocked.py @@ -0,0 +1,115 @@ +"""Placed CuTe Add-RMSNorm followed by packed NVFP4 quantization. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/cute_dsl/add_rmsnorm_fp4quant.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +The return preserves residual writeback, packed payload, and E4M3 scales. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 128, 4096 +CTA_COUNT, ROWS_PER_CTA, WARPS_PER_ROW = 64, 2, 2 +TARGET = CudaTarget("nvidia.b200_sxm") +TOPOLOGIES = ( + Topology("cta", CTA_COUNT), + Topology("thread", ROWS_PER_CTA * WARPS_PER_ROW * 32), +) + + +# noqa +@module(entry="add_rmsnorm_fp4quant", target=TARGET, topologies=TOPOLOGIES) +class AddRMSNormFP4QuantKernel: + @func + def add_rmsnorm_fp4quant( + x: Tensor[(ROWS, HIDDEN), "bf16"], + residual: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + global_scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(CTA_COUNT, ROWS_PER_CTA, WARPS_PER_ROW, 32), + names=("block", "row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "rmem", + ) + residual_reg = tf.reshard( + residual, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "smem", + ) + residual_smem = tf.reshard( + residual_reg, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "smem", + ) + weight_smem = tf.reshard(weight, (HIDDEN,), "smem") + scale_smem = tf.reshard(global_scale, (1,), "smem") + summed = x_smem + residual_smem + normalized = tf.rms_norm(summed, weight_smem, eps=1e-6) + scaled = tf.cast(normalized, "f32") / scale_smem + packed, block_scales = tf.quant( + scaled, + group=16, + target_dtype="nvfp4", + ) + return ( + tf.reshard( + packed, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "gmem", + ), + tf.reshard( + block_scales, + ( + ROWS @ (mesh.block, mesh.row), + (HIDDEN // 16) @ (mesh.warp, mesh.lane), + ), + "gmem", + ), + tf.reshard( + summed, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "gmem", + ), + ) + + +__all__ = ["AddRMSNormFP4QuantKernel"] diff --git a/tests/fixtures/flashinfer/cute_dsl_rmsnorm_fp4quant.blocked.py b/tests/fixtures/flashinfer/cute_dsl_rmsnorm_fp4quant.blocked.py new file mode 100644 index 00000000..ab53289b --- /dev/null +++ b/tests/fixtures/flashinfer/cute_dsl_rmsnorm_fp4quant.blocked.py @@ -0,0 +1,86 @@ +"""Placed CuTe RMSNorm followed by packed MXFP4 quantization. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/cute_dsl/rmsnorm_fp4quant.py +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +The program preserves 32-element MXFP4 blocks and their UE8M0 scale grid. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 128, 4096 +CTA_COUNT, ROWS_PER_CTA, WARPS_PER_ROW = 64, 2, 2 +TARGET = CudaTarget("nvidia.b200_sxm") +TOPOLOGIES = ( + Topology("cta", CTA_COUNT), + Topology("thread", ROWS_PER_CTA * WARPS_PER_ROW * 32), +) + + +# noqa +@module(entry="rmsnorm_fp4quant", target=TARGET, topologies=TOPOLOGIES) +class RMSNormFP4QuantKernel: + @func + def rmsnorm_fp4quant( + x: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(CTA_COUNT, ROWS_PER_CTA, WARPS_PER_ROW, 32), + names=("block", "row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "smem", + ) + weight_smem = tf.reshard(weight, (HIDDEN,), "smem") + normalized = tf.rms_norm(x_smem, weight_smem, eps=1e-6) + packed, block_scales = tf.quant( + tf.cast(normalized, "f32"), + group=32, + target_dtype="nvfp4", + ) + return ( + tf.reshard( + packed, + ( + ROWS @ (mesh.block, mesh.row), + HIDDEN @ (mesh.warp, mesh.lane), + ), + "gmem", + ), + tf.reshard( + block_scales, + ( + ROWS @ (mesh.block, mesh.row), + (HIDDEN // 32) @ (mesh.warp, mesh.lane), + ), + "gmem", + ), + ) + + +__all__ = ["RMSNormFP4QuantKernel"] diff --git a/tests/fixtures/flashinfer/fp4_quantization.blocked.py b/tests/fixtures/flashinfer/fp4_quantization.blocked.py new file mode 100644 index 00000000..c560bf38 --- /dev/null +++ b/tests/fixtures/flashinfer/fp4_quantization.blocked.py @@ -0,0 +1,83 @@ +"""Placed Q01 SwiGLU followed by packed NVFP4 block quantization. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/quantization/fp4_quantization.py:silu_and_mul_nvfp4_quantize +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: REG-05, OP-01 + +Token CTAs and vectorized threads preserve the packed NVFP4 payload and \ +per-16-element scales; substituting FP8 Quant would change the kernel. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 256, 1024 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 4 * 32)) + + +# noqa +@module(entry="silu_and_mul_nvfp4_quantize", target=TARGET, topologies=TOPOLOGIES) +class SwiGLUNVFP4Quantize: + @func + def silu_and_mul_nvfp4_quantize( + gate: Tensor[(ROWS, HIDDEN), "bf16"], + up: Tensor[(ROWS, HIDDEN), "bf16"], + global_scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + gate_reg = tf.reshard( + gate, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + up_reg = tf.reshard( + up, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + gate_smem = tf.reshard( + gate_reg, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + up_smem = tf.reshard( + up_reg, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + scale_smem = tf.reshard(global_scale, (1,), "smem") + activated = tf.silu(gate_smem) * up_smem + scaled = tf.cast(activated, "f32") / scale_smem + quantized, scales = tf.quant( + scaled, + group=16, + target_dtype="nvfp4", + ) + return ( + tf.reshard( + quantized, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ), + tf.reshard( + scales, + (ROWS @ mesh.token, 8 @ mesh.warp), + "gmem", + ), + ) + + +__all__ = ["SwiGLUNVFP4Quantize"] diff --git a/tests/fixtures/flashinfer/fused_dit_layernorm.py b/tests/fixtures/flashinfer/fused_dit_layernorm.py new file mode 100644 index 00000000..3338767c --- /dev/null +++ b/tests/fixtures/flashinfer/fused_dit_layernorm.py @@ -0,0 +1,120 @@ +"""Placed gated DiT residual LayerNorm scale-shift kernel. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +include/flashinfer/norm/fused_dit_layernorm.cuh +license: Apache-2.0 (no upstream source is vendored) +The specialization keeps WAN's gate, scale, shift, and residual writeback while +spelling out sum and sumsq reduction state. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 256, 3072 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 12 * 32)) + + +@module(entry="meta_fused_layernorm", target=TARGET, topologies=TOPOLOGIES) +class MetaFusedLayerNormKernel: + """FlashInfer fused DiT LayerNorm scale-shift kernel. + + predicted-ns: 4676 + waves: 2 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def meta_fused_layernorm( + x: Tensor[(ROWS, HIDDEN), "bf16"], + residual: Tensor[(ROWS, HIDDEN), "bf16"], + gate: Tensor[(ROWS, HIDDEN), "bf16"], + gate_bias: ConstTensor[(HIDDEN,), "f32"], + scale: Tensor[(ROWS, HIDDEN), "bf16"], + scale_bias: ConstTensor[(HIDDEN,), "f32"], + shift: Tensor[(ROWS, HIDDEN), "bf16"], + shift_bias: ConstTensor[(HIDDEN,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 12, 32), + names=("row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), "rmem" + ) + residual_reg = tf.reshard( + residual, + (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + gate_reg = tf.reshard( + gate, (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), "rmem" + ) + scale_reg = tf.reshard( + scale, (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), "rmem" + ) + shift_reg = tf.reshard( + shift, (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), "rmem" + ) + x_smem = tf.reshard( + x_reg, (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), "smem" + ) + residual_smem = tf.reshard( + residual_reg, + (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + gate_smem = tf.reshard( + gate_reg, + (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + scale_smem = tf.reshard( + scale_reg, + (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + shift_smem = tf.reshard( + shift_reg, + (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + gate_bias_smem = tf.reshard(gate_bias, (HIDDEN,), "smem") + scale_bias_smem = tf.reshard(scale_bias, (HIDDEN,), "smem") + shift_bias_smem = tf.reshard(shift_bias, (HIDDEN,), "smem") + gated = tf.cast(x_smem, "f32") * ( + tf.cast(gate_smem, "f32") + gate_bias_smem + ) + summed = gated + tf.cast(residual_smem, "f32") + count = tf.reduce( + tf.full_like(summed, value=1.0), axes=(-1,), keepdim=True, kind="sum" + ) + reduced = tf.reduce(summed, axes=(-1,), keepdim=True, kind="sum") + sumsq = tf.reduce(summed * summed, axes=(-1,), keepdim=True, kind="sum") + mean = reduced / count + variance = sumsq / count - mean * mean + normalized = (summed - mean) * tf.rsqrt(variance + 1e-6) + shifted = normalized * ( + tf.cast(scale_smem, "f32") + scale_bias_smem + 1.0 + ) + (tf.cast(shift_smem, "f32") + shift_bias_smem) + return ( + tf.reshard( + tf.cast(shifted, "bf16"), + (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ), + tf.reshard( + tf.cast(summed, "bf16"), + (ROWS @ mesh.row, 24 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ), + ) + + +__all__ = ["MetaFusedLayerNormKernel"] diff --git a/tests/fixtures/flashinfer/fused_qk_rmsnorm_rope.blocked.py b/tests/fixtures/flashinfer/fused_qk_rmsnorm_rope.blocked.py new file mode 100644 index 00000000..81af8c01 --- /dev/null +++ b/tests/fixtures/flashinfer/fused_qk_rmsnorm_rope.blocked.py @@ -0,0 +1,116 @@ +"""Placed fused QK RMSNorm and three-dimensional RoPE kernel. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +include/flashinfer/norm/fused_qk_rmsnorm_rope.cuh +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: Split: Split: axis 1 is divided into 3 parts and is already Split across participants +ledger: REG-08 + +The program keeps the packed QKV input, across-head RMSNorm, derived 3D +positions, and the upstream three-output contract. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +TOKENS, HEADS, HEAD_DIM = 120, 24, 128 +PPF, PPH, PPW = 1, 5, 24 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", TOKENS * 3), Topology("thread", HEADS * 32)) + + +@module(entry="fused_qk_norm_rope", target=TARGET, topologies=TOPOLOGIES) +class FusedQKNormRopeKernel: + @func + def fused_qk_norm_rope( + qkv: Tensor[(TOKENS, 3 * HEADS * HEAD_DIM), "bf16"], + q_weight: ConstTensor[(HEADS * HEAD_DIM,), "bf16"], + k_weight: ConstTensor[(HEADS * HEAD_DIM,), "bf16"], + freq: ConstTensor[(HEAD_DIM // 2,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(TOKENS, 3, HEADS, 32), + names=("token", "kind", "head", "lane"), + ) as mesh: + qkv_reg = tf.reshard( + qkv, + ( + TOKENS @ mesh.token, + 3, + HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "rmem", + ) + qkv_smem = tf.reshard( + qkv_reg, + ( + TOKENS @ mesh.token, + 3, + HEADS @ mesh.head, + HEAD_DIM @ mesh.lane, + ), + "smem", + ) + q_smem, k_smem, v_smem = tf.split(qkv_smem, axis=1, num_splits=3) + q_weight_smem = tf.reshard(q_weight, (HEADS * HEAD_DIM,), "smem") + k_weight_smem = tf.reshard(k_weight, (HEADS * HEAD_DIM,), "smem") + freq_smem = tf.reshard(freq, (HEAD_DIM // 2,), "smem") + positions = tf.arange(Tensor[(TOKENS,), "i64"]) + frame_reg = tf.reshard( + positions // (PPH * PPW), (TOKENS @ mesh.token,), "rmem" + ) + height_reg = tf.reshard( + (positions % (PPH * PPW)) // PPW, + (TOKENS @ mesh.token,), + "rmem", + ) + width_reg = tf.reshard( + positions % PPW, (TOKENS @ mesh.token,), "rmem" + ) + q_normalized_flat = tf.rms_norm(q_smem, q_weight_smem, eps=1e-6) + k_normalized_flat = tf.rms_norm(k_smem, k_weight_smem, eps=1e-6) + q_normalized = tf.reshape( + q_normalized_flat, + (TOKENS, HEADS, HEAD_DIM), + ) + k_normalized = tf.reshape( + k_normalized_flat, + (TOKENS, HEADS, HEAD_DIM), + ) + v = tf.reshape(v_smem, (TOKENS, HEADS, HEAD_DIM)) + q_rope, k_rope = tf.rope_3d( + q_normalized, + k_normalized, + freq_smem, + frame_reg, + height_reg, + width_reg, + ) + return ( + tf.reshard( + q_rope, + (TOKENS @ mesh.token, HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "gmem", + ), + tf.reshard( + k_rope, + (TOKENS @ mesh.token, HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "gmem", + ), + tf.reshard( + v, + (TOKENS @ mesh.token, HEADS @ mesh.head, HEAD_DIM @ mesh.lane), + "gmem", + ), + ) + + +__all__ = ["FusedQKNormRopeKernel"] diff --git a/tests/fixtures/flashinfer/gemm_svdquant.blocked.py b/tests/fixtures/flashinfer/gemm_svdquant.blocked.py new file mode 100644 index 00000000..1a08fcd0 --- /dev/null +++ b/tests/fixtures/flashinfer/gemm_svdquant.blocked.py @@ -0,0 +1,86 @@ +"""Placed G01 NVFP4 residual GEMM with LoRA-up and bias epilogue. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/gemm/gemm_svdquant.py:mm_nvfp4_svdquant +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: EXT-04, OP-01, OP-02 + +Row CTAs and vectorized threads preserve packed operands, block scales, \ +BF16 LoRA correction, alpha, and bias instead of substituting a BF16 twin. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +M, K, N, R = 256, 128, 128, 32 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", M), Topology("thread", 4 * 32)) + + +# noqa +@module(entry="mm_nvfp4_svdquant", target=TARGET, topologies=TOPOLOGIES) +class NVFP4SVDQuantGemm: + @func + def mm_nvfp4_svdquant( + a: Tensor[(M, K), "nvfp4"], + b: ConstTensor[(N, K), "nvfp4"], + a_scale: Tensor[(M, K // 16), "fp8e4m3"], + b_scale: ConstTensor[(N, K // 16), "fp8e4m3"], + alpha: ConstTensor[(1,), "f32"], + down: Tensor[(M, R), "bf16"], + lora_up: ConstTensor[(N, R), "bf16"], + bias: ConstTensor[(N,), "bf16"], + ) -> Tensor[(M, N), "bf16"]: + with Mesh( + ("cta", "thread"), + layout=(M, 4, 32), + names=("row", "warp", "lane"), + ) as mesh: + a_reg = tf.reshard( + a, + (M @ mesh.row, 1 @ mesh.warp, K @ mesh.lane), + "rmem", + ) + b_reg = tf.reshard( + b, + (N, 1 @ mesh.warp, K @ mesh.lane), + "rmem", + ) + a_smem = tf.reshard( + a_reg, + (M @ mesh.row, 1 @ mesh.warp, K @ mesh.lane), + "smem", + ) + b_smem = tf.reshard( + b_reg, + (N, 1 @ mesh.warp, K @ mesh.lane), + "smem", + ) + a_scale_smem = tf.reshard(a_scale, (M @ mesh.row, K // 16), "smem") + b_scale_smem = tf.reshard(b_scale, (N, K // 16), "smem") + down_smem = tf.reshard(down, (M @ mesh.row, R), "smem") + lora_up_smem = tf.reshard(lora_up, (N, R), "smem") + bias_smem = tf.reshard(bias, (N,), "smem") + residual = tf.block_scaled_matmul( + a_smem, + tf.transpose(b_smem, perm=(1, 0)), + a_scale_smem, + b_scale_smem, + ) + correction = down_smem @ tf.transpose(lora_up_smem, perm=(1, 0)) + output = alpha * residual + correction + bias_smem + return tf.reshard( + output, + (M @ mesh.row, 4 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ) + + +__all__ = ["NVFP4SVDQuantGemm"] diff --git a/tests/fixtures/flashinfer/gemm_transforms.blocked.py b/tests/fixtures/flashinfer/gemm_transforms.blocked.py new file mode 100644 index 00000000..03baaa07 --- /dev/null +++ b/tests/fixtures/flashinfer/gemm_transforms.blocked.py @@ -0,0 +1,151 @@ +"""Expected authored HIR for G01, T01-T02, and L01. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c58fdd2392914ea05e2a8714946ac0eef6 +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: DType: unknown value 'nvfp4' +ledger: EXT-04, OP-04, OP-14 +This is negative conformance input using proposed authored-HIR operations. +""" + +from tilefoundry import func, module +from tilefoundry.dsl import Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOS = (Topology("cta", 8), Topology("thread", 128)) +M, K, N, R = 128, 64, 128, 16 + + +# noqa +@module(entry="fused_cta", target=TARGET, topologies=TOPOS) +class G01ResidualGemmLoraBias: + @func + def unfused( + a: Tensor[(M, K), "bf16"], + b: Tensor[(N, K), "nvfp4"], + d: Tensor[(M, R), "bf16"], + l1: Tensor[(N, R), "bf16"], + bias: Tensor[(N,), "bf16"], + ) -> Tensor[(M, N), "bf16"]: + base = tf.nvfp4_matmul(a, tf.transpose(b, perm=(1, 0))) + delta = tf.matmul(d, tf.transpose(l1, perm=(1, 0))) + return base + delta + bias + + @func + def fused_program( + a: Tensor[(M, K), "bf16"], + b: Tensor[(N, K), "nvfp4"], + d: Tensor[(M, R), "bf16"], + l1: Tensor[(N, R), "bf16"], + bias: Tensor[(N,), "bf16"], + ) -> Tensor[(M, N), "bf16"]: + + acc = tf.nvfp4_mma_accumulate(a, b, accumulator="f32") + acc = tf.lora_up_accumulate(acc, d, l1) + return tf.epilogue(acc, bias=bias, alpha=1.0, output_dtype="bf16") + + @func + def fused_cta( + a: Tensor[(M, K), "bf16"], + b: Tensor[(N, K), "nvfp4"], + d: Tensor[(M, R), "bf16"], + l1: Tensor[(N, R), "bf16"], + bias: Tensor[(N,), "bf16"], + ) -> Tensor[(M, N), "bf16"]: + with Mesh(("cta",), layout=(8,)) as cta: + a_tile = tf.reshard(a, (M @ cta, K), "smem") + d_tile = tf.reshard(d, (M @ cta, R), "smem") + b_tile = tf.reshard(b, (N, K), "smem") + l1_tile = tf.reshard(l1, (N, R), "smem") + bias_tile = tf.reshard(bias, (N,), "smem") + acc = tf.nvfp4_mma_accumulate(a_tile, b_tile, accumulator="rmem") + acc = tf.lora_up_accumulate(acc, d_tile, l1_tile) + out = tf.epilogue(acc, bias=bias_tile, alpha=1.0, output_dtype="bf16") + return tf.reshard(out, (M, N), "gmem") + + @func + def fused_thread( + a: Tensor[(M, K), "bf16"], + b: Tensor[(N, K), "nvfp4"], + d: Tensor[(M, R), "bf16"], + l1: Tensor[(N, R), "bf16"], + bias: Tensor[(N,), "bf16"], + ) -> Tensor[(M, N), "bf16"]: + with Mesh(("thread",), layout=(4, 32), names=("warp", "lane")) as thread: + acc = tf.mma_fragment(a, b, owner=thread, storage="rmem", input_format="nvfp4") + acc = tf.lora_up_accumulate(acc, d, l1, owner=thread) + out = tf.epilogue(acc, bias=bias, output_dtype="bf16") + return tf.reshard(out, (M, N), "gmem") + + +# noqa +@module(entry="topk_page_cta", target=TARGET, topologies=TOPOS) +class TopKIndexTransforms: + @func + def topk_page_unfused( + scores: Tensor[(M, N), "bf16"], page_table: Tensor[(M, N), "i32"] + ) -> Tensor[(M, 8), "i32"]: + _, indices = tf.topk(scores, k=8) + return tf.page_table_transform(indices, page_table) + + @func + def topk_page_cta( + scores: Tensor[(M, N), "bf16"], page_table: Tensor[(M, N), "i32"] + ) -> Tensor[(M, 8), "i32"]: + with Mesh(("cta",), layout=(8,)) as cta: + local = tf.reshard(scores, (M @ cta, N), "smem") + _, indices = tf.topk(local, k=8) + transformed = tf.page_table_transform(indices, page_table, storage="smem") + return tf.reshard(transformed, (M, 8), "gmem") + + @func + def topk_page_thread( + scores: Tensor[(M, N), "bf16"], page_table: Tensor[(M, N), "i32"] + ) -> Tensor[(M, 8), "i32"]: + with Mesh(("thread",), layout=(128,)) as thread: + local = tf.reshard(scores, (M @ thread, N), "rmem") + _, indices = tf.topk(local, k=8) + transformed = tf.page_table_transform(indices, page_table, storage="rmem") + return tf.reshard(transformed, (M, 8), "gmem") + + @func + def topk_ragged_cta( + scores: Tensor[(M, N), "bf16"], indptr: Tensor[(M + 1,), "i32"] + ) -> Tensor[(M, 8), "i32"]: + with Mesh(("cta",), layout=(8,)) as cta: + local = tf.reshard(scores, (M @ cta, N), "smem") + _, indices = tf.topk(local, k=8) + return tf.reshard(tf.ragged_index_transform(indices, indptr), (M, 8), "gmem") + + +# noqa +@module(entry="fused_thread", target=TARGET, topologies=TOPOS) +class LogitsProcessorChain: + @func + def unfused( + logits: Tensor[(M, N), "f32"], temperature: Tensor[(M, 1), "f32"] + ) -> Tensor[(M, N), "f32"]: + scaled = logits / temperature + penalized = tf.repetition_penalty(scaled) + return tf.top_p_mask(tf.top_k_mask(penalized)) + + @func + def fused_thread( + logits: Tensor[(M, N), "f32"], temperature: Tensor[(M, 1), "f32"] + ) -> Tensor[(M, N), "f32"]: + with Mesh(("thread",), layout=(128,)) as thread: + local = tf.reshard(logits, (M @ thread, N), "rmem") + scaled = local / temperature + penalized = tf.repetition_penalty(scaled, storage="rmem") + selected = tf.fused_top_k_top_p_mask(penalized, storage="rmem") + return tf.reshard(selected, (M, N), "gmem") + + +# noqa +# noqa +# noqa +# noqa diff --git a/tests/fixtures/flashinfer/ln_fwd_silu_kernel.py b/tests/fixtures/flashinfer/ln_fwd_silu_kernel.py new file mode 100644 index 00000000..02251811 --- /dev/null +++ b/tests/fixtures/flashinfer/ln_fwd_silu_kernel.py @@ -0,0 +1,63 @@ +"""Placed SM100 RMSNorm and SiLU specialization. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +include/flashinfer/norm/ln_fwd_silu_kernel.cuh +license: Apache-2.0 (no upstream source is vendored) + +This is the upstream LUT specialization for 1,560 tokens, hidden size 1,024, +and bf16 output: 390 CTAs, four rows per CTA, and one warp per row. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 1560, 1024 +CTA_COUNT, ROWS_PER_CTA = 390, 4 +TARGET = CudaTarget("nvidia.b200_sxm") +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", ROWS_PER_CTA * 32)) + + +@module(entry="ln_fwd_kernel", target=TARGET, topologies=TOPOLOGIES) +class LnFwdSiluKernel: + """FlashInfer SM100 ln_fwd_kernel RMSNorm-SiLU bf16 specialization. + + predicted-ns: 1272 waves: 3 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def ln_fwd_kernel( + x: Tensor[(ROWS, HIDDEN), "bf16"], + gamma: ConstTensor[(HIDDEN,), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(CTA_COUNT, ROWS_PER_CTA, 32), + names=("block", "row", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "smem", + ) + gamma_smem = tf.reshard(gamma, (HIDDEN,), "smem") + normalized = tf.rms_norm(x_smem, gamma_smem, eps=1e-6) + activated = tf.silu(normalized) + return tf.reshard( + activated, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "gmem", + ) + + +__all__ = ["LnFwdSiluKernel"] diff --git a/tests/fixtures/flashinfer/moe_recurrent.blocked.py b/tests/fixtures/flashinfer/moe_recurrent.blocked.py new file mode 100644 index 00000000..ea1a3f9c --- /dev/null +++ b/tests/fixtures/flashinfer/moe_recurrent.blocked.py @@ -0,0 +1,41 @@ +"""MoE dispatch into a grouped matmul that has no current HIR surface. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c5 csrc/moe.cu:grouped_gemm +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.grouped_matmul' (2 positional, no keywords) +ledger: OP-05 +placement_refusal: IndexSelect rejects the sharded smem operand on dim 0. +placement_error: IndexSelect: dim 0 index_select over a shard layout with multiple +Split axes including the selected dim; cannot derive an output layout +placement_workaround: materialize the staged tensor back to gmem before IndexSelect. +""" + +from tilefoundry import module +from tilefoundry.dsl import Mesh, Tensor, Topology, func, tf +from tilefoundry.target import CudaTarget + + +@module( + entry="grouped", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 2), Topology("thread", 4)), +) +class MoeRecurrent: + @func + def grouped( + tokens: Tensor[(8, 16), "bf16"], + weights: Tensor[(8, 16), "bf16"], + route: Tensor[(4,), "i32"], + ): + with Mesh(("cta",), layout=(2,), names=("tile",)) as cta: + with Mesh(("thread",), layout=(4,), names=("lane",)) as thread: + tokens_smem = tf.reshard(tokens, (8 @ cta.tile, 16 @ thread.lane), "smem") + tokens_for_gather = tf.reshard(tokens_smem, (8, 16), "gmem") + weights_smem = tf.reshard(weights, (8, 16), "smem") + route_rmem = tf.reshard(route, (4,), "rmem") + routed = tf.index_select(tokens_for_gather, route_rmem, dim=0) + routed_smem = tf.reshard(routed, (4 @ cta.tile, 16 @ thread.lane), "smem") + return tf.grouped_matmul(routed_smem, weights_smem) diff --git a/tests/fixtures/flashinfer/moe_recurrent_spec.blocked.py b/tests/fixtures/flashinfer/moe_recurrent_spec.blocked.py new file mode 100644 index 00000000..6e6f18c8 --- /dev/null +++ b/tests/fixtures/flashinfer/moe_recurrent_spec.blocked.py @@ -0,0 +1,139 @@ +"""Complete equation inventory for MoE, recurrent, state-space, and mHC. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c58fdd2392914ea05e2a8714946ac0eef6 +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: selection/analysis +error: source defines no TileFoundry Module +classification: expected-spec notation; no authored Module is declared. +Equations retain unfused/fused boundaries and placement variants as test data. +""" + +# noqa +M01_U = """def m01_unfused(x:gmem, bias:gmem): + p = sigmoid(x); s = group_score(p, bias); g = topk(s, groups=4) + e = topk(gather(p, g), k=8); return renorm(e) -> gmem +""" +M01 = ( + M01_U + + """def m01_fused(x:gmem,bias:gmem): + e = renorm(topk(group_topk(sigmoid(x)+bias, groups=4), k=8)) -> smem + return e # [program=token, cta=group, thread=expert, gpu=rank] +""" +) + +M02_U = """def m02_unfused(x:gmem, w1:gmem): + r = gather(x, route(x)) -> gmem; z = grouped_gemm(r,w1) -> gmem + return silu_mul(z) -> gmem +""" +M02 = ( + M02_U + + """def m02_fused(x:gmem,w1:gmem): + z = silu_mul(grouped_gemm(gather(x,route(x)) -> smem,w1) -> smem) -> gmem + return z # [program=expert, cta=GEMM-MN, thread=fragment] +""" +) + +M03_U = """def m03_unfused(z:gmem,w2:gmem,weights:gmem,route_ids:gmem): + y = grouped_gemm(z,w2) -> gmem; return weighted_scatter_add(y,weights,route_ids) +""" +M03 = ( + M03_U + + """def m03_fused(z:gmem,w2:gmem,weights:gmem,route_ids:gmem): + y = weighted_scatter_add(grouped_gemm(z,w2) -> smem,weights,route_ids) -> gmem + return y # [program=expert, cta=output-tile, thread=lane, gpu=rank] +""" +) + +M04_U = """def m04_unfused(x:gmem,w1:gmem,w2:gmem): + r=route(x); a=grouped_gemm(gather(x,r),w1); h=silu_mul(a); y=grouped_gemm(h,w2) + return scatter_add(y,r) +""" +M04 = ( + M04_U + + """def m04_fused(x:gmem,w1:gmem,w2:gmem): + return scatter_add(grouped_gemm(silu_mul(grouped_gemm(gather(x,route(x))->smem,w1)->smem),w2)->smem,route(x)) + # [program=token, cta=expert-tile, thread=mma-lane, gpu=rank] +""" +) + +M05_U = """def m05_unfused(x:gmem,w1:gmem,w2:gmem): + p=renorm(topk(score(x))); return reduce_experts([grouped_gemm(silu_mul(grouped_gemm(x,w1)),w2)],p) +""" +M05 = ( + M05_U + + """def m05_fused(x:gmem,w1:gmem,w2:gmem): + return reduce_experts_fused(x,w1,w2,route=renorm(topk(score(x)))) -> gmem + # [program=token, cta=expert, thread=lane] +""" +) + +M06_U = """def m06_unfused(x:gmem,rw:gmem,sw:gmem): + y=routed_moe(x,rw); sh=shared_expert(x,sw); return y+sh +""" +M06 = ( + M06_U + + """def m06_fused(x:gmem,rw:gmem,sw:gmem): + return finalize(routed_moe(x,rw)->smem + shared_expert(x,sw)->smem) -> gmem + # [program=token, cta=expert-tile, thread=accumulator, gpu=rank] +""" +) + +M07 = """def m07_fused(x:gmem,w:gmem,scale:gmem): + q=block_quant(x,scale)->smem; return finalize(grouped_gemm(q,w)->smem)->gmem +""" +M08 = """def m08_fused(x:gmem,w:gmem): + q=pack_quant(x, dtype=fp4)->smem; return finalize(grouped_gemm(q,w)->smem)->gmem +""" +M09 = """def m09_fused(x:gmem,w:gmem): + q=dequant(x,dtype=mxint4)->smem; return finalize(grouped_gemm(q,w)->smem)->gmem +""" +M10 = """def m10_fused(x:gmem,w:gmem,A:gmem,B:gmem): + return grouped_gemm(x,w)->smem + bgmv(x,A,B)->smem -> gmem +""" +M11 = """def m11_fused(x:gmem,w:gmem,A:gmem,B:gmem,route:gmem): + return finalize(grouped_gemm(x,w)->smem + bgmv(x,A,B,route)->smem)->gmem +""" +M12 = """def m12_fused(x:gmem,w:gmem,pack:gmem): + r=route_pack(x,pack)->smem; return w4a16_expert_pipeline(r,w)->gmem +""" + +# noqa +K01 = """def k01_fused(x:gmem,state:gmem,w:gmem): + c=depthwise_conv4(x,w)->smem; u=silu(c); state2=kda_update(state,u)->gmem + return rms_norm(state2)*silu(x) # [program=token,cta=head,thread=channel,gpu=rank] +""" +K02 = """def k02_fused(q:gmem,k:gmem,v:gmem,state:gmem): + qn=l2_norm(q)->smem; kn=l2_norm(k)->smem; b=sigmoid(beta(q,k)); s=kda_update(state,qn,kn,v,b)->gmem + return s # [program=sequence, cta=head, thread=channel] +""" +K03 = """def k03_fused(tokens:gmem,state:gmem,accepted:gmem): + s=state + for t in tokens: s=kda_step(s,t)->rmem + return select_checkpoint(s,accepted)->gmem # [program=sequence,cta=head,thread=channel] +""" + +# noqa +S01 = """def s01_fused(x:gmem,A:gmem,B:gmem,C:gmem,D:gmem): + z=chunk_cumsum(x,A)->smem; s=ssd_scan(z,B,C)->gmem + return s + D*x # [program=chunk,cta=sequence-tile,thread=channel,gpu=rank] +""" +S02 = """def s02_fused(old:gmem,new:gmem,checkpoint:gmem,pred:gmem): + replay=ssu_replay(old,checkpoint)->smem; out,st=ssu_step(replay,new)->smem + cache_write_if(pred,st)->gmem; return out # [program=sequence,cta=chunk,thread=channel] +""" + +# noqa +H01 = """def h01_fused(x:gmem,H:gmem): + P=sinkhorn(H)->smem; y=residual_mix(x,P)->gmem + return y # [program=batch,cta=hidden-tile,thread=element,gpu=rank] +""" +H02 = """def h02_fused(x:gmem,H:gmem,gamma:gmem,beta:gmem): + n=layer_norm(x,gamma,beta)->smem; P=sinkhorn(H)->smem + return residual_mix(n,P)->gmem # [program=batch,cta=hidden-tile,thread=element,gpu=rank] +""" + +ALL = { + k: v for k, v in globals().items() if k.startswith(("M", "K", "S", "H")) and k[1:2].isdigit() +} diff --git a/tests/fixtures/flashinfer/moe_state_hir.blocked.py b/tests/fixtures/flashinfer/moe_state_hir.blocked.py new file mode 100644 index 00000000..cc1d309a --- /dev/null +++ b/tests/fixtures/flashinfer/moe_state_hir.blocked.py @@ -0,0 +1,162 @@ +"""Executable-shaped expected HIR for representative MoE/state mechanisms. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c58fdd2392914ea05e2a8714946ac0eef6 +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.route_topk' (1 positional, keywords ['experts', 'k']) +ledger: OP-05, OP-06, OP-11, OP-13 +Toy extents and proposed operations retain the complete mechanism inventory. +""" + +from tilefoundry import func, module +from tilefoundry.dsl import Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +TARGET = CudaTarget("nvidia.h200_sxm") +CUDA_TOPOS = (Topology("cta", 8), Topology("thread", 128)) +TOKENS, HIDDEN, EXPERTS, TOPK = 128, 64, 16, 2 + + +# noqa +@module(entry="fused_cta", target=TARGET, topologies=CUDA_TOPOS) +class M04FullMoE: + @func + def unfused( + x: Tensor[(TOKENS, HIDDEN), "bf16"], + w1: Tensor[(EXPERTS, HIDDEN, 2 * HIDDEN), "bf16"], + w2: Tensor[(EXPERTS, HIDDEN, HIDDEN), "bf16"], + ) -> Tensor[(TOKENS, HIDDEN), "bf16"]: + weights, experts = tf.route_topk(x, experts=EXPERTS, k=TOPK) + permuted = tf.expert_gather(x, experts) + gate_up = tf.grouped_matmul(permuted, w1) + hidden = tf.silu_and_mul(gate_up) + expert_out = tf.grouped_matmul(hidden, w2) + return tf.weighted_scatter_reduce(expert_out, weights, experts) + + @func + def fused_program( + x: Tensor[(TOKENS, HIDDEN), "bf16"], + w1: Tensor[(EXPERTS, HIDDEN, 2 * HIDDEN), "bf16"], + w2: Tensor[(EXPERTS, HIDDEN, HIDDEN), "bf16"], + ) -> Tensor[(TOKENS, HIDDEN), "bf16"]: + return tf.fused_moe(x, w1, w2, experts=EXPERTS, k=TOPK) + + @func + def fused_cta( + x: Tensor[(TOKENS, HIDDEN), "bf16"], + w1: Tensor[(EXPERTS, HIDDEN, 2 * HIDDEN), "bf16"], + w2: Tensor[(EXPERTS, HIDDEN, HIDDEN), "bf16"], + ) -> Tensor[(TOKENS, HIDDEN), "bf16"]: + with Mesh(("cta",), layout=(8,), names=("expert_tile",)) as cta: + weights, experts = tf.route_topk(x, experts=EXPERTS, k=TOPK) + routed = tf.expert_gather(x, experts, owner=cta, storage="smem") + gate_up = tf.grouped_mma(routed, w1, accumulator="rmem") + hidden = tf.silu_and_mul(gate_up, storage="rmem") + down = tf.grouped_mma(hidden, w2, accumulator="rmem") + out = tf.weighted_scatter_reduce(down, weights, experts, storage="smem") + return tf.reshard(out, (TOKENS, HIDDEN), "gmem") + + @func + def fused_thread( + x: Tensor[(TOKENS, HIDDEN), "bf16"], + w1: Tensor[(EXPERTS, HIDDEN, 2 * HIDDEN), "bf16"], + w2: Tensor[(EXPERTS, HIDDEN, HIDDEN), "bf16"], + ) -> Tensor[(TOKENS, HIDDEN), "bf16"]: + with Mesh(("thread",), layout=(4, 32), names=("warp", "lane")) as thread: + return tf.fused_moe_fragment(x, w1, w2, owner=thread, storage="rmem") + + @func + def fused_gpu( + x: Tensor[(TOKENS, HIDDEN), "bf16"], + w1: Tensor[(EXPERTS, HIDDEN, 2 * HIDDEN), "bf16"], + w2: Tensor[(EXPERTS, HIDDEN, HIDDEN), "bf16"], + ) -> Tensor[(TOKENS, HIDDEN), "bf16"]: + with Mesh(("gpu",), layout=(8,), names=("rank",)) as gpu: + routed = tf.all_to_all_expert_dispatch(x, owner=gpu, storage="gmem") + local = tf.fused_moe(routed, w1, w2, owner=gpu) + return tf.all_to_all_expert_combine(local, owner=gpu) + + +# noqa +@module(entry="fused_cta", target=TARGET, topologies=CUDA_TOPOS) +class K01FusedDecode: + @func + def unfused(x, conv_weight, state, norm_weight, gate): + conv = tf.depthwise_conv4(x, conv_weight) + activated = tf.silu(conv) + next_state, read = tf.kda_update(state, activated) + return tf.rms_norm(read, norm_weight) * tf.silu(gate), next_state + + @func + def fused_cta(x, conv_weight, state, norm_weight, gate): + with Mesh(("cta",), layout=(8,), names=("head",)) as cta: + x_local = tf.place(x, owner=cta, storage="smem") + conv = tf.depthwise_conv4(x_local, conv_weight, storage="smem") + activated = tf.silu(conv) + next_state, read = tf.kda_update(state, activated, owner=cta, storage="smem") + out = tf.rms_norm(read, norm_weight) * tf.silu(gate) + return tf.reshard(out, tf.logical_layout(out), "gmem"), next_state + + @func + def fused_thread(x, conv_weight, state, norm_weight, gate): + with Mesh(("thread",), layout=(128,), names=("channel",)) as thread: + return tf.fused_kda_decode( + x, conv_weight, state, norm_weight, gate, owner=thread, storage="rmem" + ) + + +# noqa +@module(entry="fused_cta", target=TARGET, topologies=CUDA_TOPOS) +class S01SSDCombined: + @func + def unfused(x, a, b, c, d, z): + cumulative = tf.chunk_cumsum(x, a) + states = tf.ssd_state_passing(cumulative, b) + scanned = tf.ssd_scan(states, c) + return tf.silu(z) * (scanned + d * x) + + @func + def fused_cta(x, a, b, c, d, z): + with Mesh(("cta",), layout=(8,), names=("chunk",)) as cta: + cumulative = tf.chunk_cumsum(x, a, owner=cta, storage="smem") + states = tf.ssd_state_passing(cumulative, b, storage="smem") + scanned = tf.ssd_scan(states, c, storage="rmem") + out = tf.silu(z) * (scanned + d * x) + return tf.reshard(out, tf.logical_layout(out), "gmem") + + @func + def fused_thread(x, a, b, c, d, z): + with Mesh(("thread",), layout=(128,), names=("channel",)) as thread: + return tf.ssd_combined(x, a, b, c, d, z, owner=thread, storage="rmem") + + +# noqa +@module(entry="fused_cta", target=TARGET, topologies=CUDA_TOPOS) +class H02MHCWithPrenorm: + @func + def unfused(x, h, gamma, beta): + normalized = tf.layer_norm(x, gamma, beta, axis=-1, eps=1e-5) + transform = tf.sinkhorn_transform(h) + return tf.residual_mix(normalized, transform) + + @func + def fused_cta(x, h, gamma, beta): + with Mesh(("cta",), layout=(8,), names=("row",)) as cta: + local = tf.place(x, owner=cta, storage="smem") + normalized = tf.layer_norm(local, gamma, beta, axis=-1, eps=1e-5) + transform = tf.sinkhorn_transform(h, owner=cta, storage="smem") + mixed = tf.residual_mix(normalized, transform, storage="smem") + return tf.reshard(mixed, tf.logical_layout(mixed), "gmem") + + @func + def fused_thread(x, h, gamma, beta): + with Mesh(("thread",), layout=(128,), names=("element",)) as thread: + return tf.mhc_pre_with_prenorm(x, h, gamma, beta, owner=thread, storage="rmem") + + +# noqa +# noqa +# noqa diff --git a/tests/fixtures/flashinfer/norm.py b/tests/fixtures/flashinfer/norm.py new file mode 100644 index 00000000..a2e79498 --- /dev/null +++ b/tests/fixtures/flashinfer/norm.py @@ -0,0 +1,157 @@ +"""Placed FlashInfer LayerNorm fusion boundaries blocked by reduction sharding. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/norm/__init__.py +license: Apache-2.0 (no upstream source is vendored) +The explicit sum and sumsq reductions preserve token-CTA and hidden-thread +placement without hiding cross-participant combination inside LayerNorm. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN, DIT_HIDDEN = 256, 1024, 3072 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 4 * 32)) + + +@module(entry="layernorm_quant", target=TARGET, topologies=TOPOLOGIES) +class LayerNormFixedScaleFP8: + """FlashInfer rmsnorm_fp8_quant kernel. + + predicted-ns: 700 + waves: 2 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def layernorm_quant( + x: Tensor[(ROWS, HIDDEN), "bf16"], + gamma: ConstTensor[(HIDDEN,), "f32"], + beta: ConstTensor[(HIDDEN,), "f32"], + scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + gamma_smem = tf.reshard(gamma, (HIDDEN,), "smem") + beta_smem = tf.reshard(beta, (HIDDEN,), "smem") + scale_smem = tf.reshard(scale, (1,), "smem") + x32 = tf.cast(x_smem, "f32") + count = tf.reduce( + tf.full_like(x32, value=1.0), axes=(-1,), keepdim=True, kind="sum" + ) + summed = tf.reduce(x32, axes=(-1,), keepdim=True, kind="sum") + sumsq = tf.reduce(x32 * x32, axes=(-1,), keepdim=True, kind="sum") + mean = summed / count + variance = sumsq / count - mean * mean + normalized = (x32 - mean) * tf.rsqrt(variance + 1e-6) + affine = normalized * gamma_smem + beta_smem + quantized = tf.cast(affine / scale_smem, "fp8e4m3") + return tf.reshard( + quantized, + (ROWS @ mesh.token, 8 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ) + + +@module(entry="fused_dit_gate_layernorm", target=TARGET, topologies=TOPOLOGIES) +class DiTGateLayerNorm: + """FlashInfer fused DiT gate and LayerNorm kernel. + + predicted-ns: 4266 + waves: 2 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def fused_dit_gate_layernorm( + x: Tensor[(ROWS, DIT_HIDDEN), "bf16"], + residual: Tensor[(ROWS, DIT_HIDDEN), "bf16"], + gate: Tensor[(ROWS, DIT_HIDDEN), "bf16"], + gate_bias: ConstTensor[(DIT_HIDDEN,), "f32"], + gamma: ConstTensor[(DIT_HIDDEN,), "f32"], + beta: ConstTensor[(DIT_HIDDEN,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 4, 32), + names=("token", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + residual_reg = tf.reshard( + residual, + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + gate_reg = tf.reshard( + gate, + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + residual_smem = tf.reshard( + residual_reg, + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + gate_smem = tf.reshard( + gate_reg, + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + gate_bias_smem = tf.reshard(gate_bias, (DIT_HIDDEN,), "smem") + gamma_smem = tf.reshard(gamma, (DIT_HIDDEN,), "smem") + beta_smem = tf.reshard(beta, (DIT_HIDDEN,), "smem") + gate_f32 = tf.cast(gate_smem, "f32") + gate_bias_smem + summed = tf.cast(residual_smem, "f32") + tf.cast(x_smem, "f32") * gate_f32 + count = tf.reduce( + tf.full_like(summed, value=1.0), axes=(-1,), keepdim=True, kind="sum" + ) + reduced = tf.reduce(summed, axes=(-1,), keepdim=True, kind="sum") + sumsq = tf.reduce(summed * summed, axes=(-1,), keepdim=True, kind="sum") + mean = reduced / count + variance = sumsq / count - mean * mean + normalized = (summed - mean) * tf.rsqrt(variance + 1e-6) + normalized = normalized * gamma_smem + beta_smem + return ( + tf.reshard( + tf.cast(summed, "bf16"), + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ), + tf.reshard( + normalized, + (ROWS @ mesh.token, 24 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ), + ) + + +__all__ = ["LayerNormFixedScaleFP8", "DiTGateLayerNorm"] diff --git a/tests/fixtures/flashinfer/norm_cuda.py b/tests/fixtures/flashinfer/norm_cuda.py new file mode 100644 index 00000000..c82e86e3 --- /dev/null +++ b/tests/fixtures/flashinfer/norm_cuda.py @@ -0,0 +1,65 @@ +"""Placed CUDA RMSNorm followed by fixed-scale FP8 quantization. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +csrc/norm.cu:rmsnorm_quant +license: Apache-2.0 (no upstream source is vendored) + +Each CTA owns one token row. Four warps and 32 lanes vectorize the hidden +dimension, matching the row-wise CUDA normalization kernel. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 256, 1536 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 4 * 32)) + + +@module(entry="norm_quant", target=TARGET, topologies=TOPOLOGIES) +class RMSNormQuant: + """FlashInfer CUDA fixed-scale FP8 RMSNorm kernel. + + predicted-ns: 496 waves: 2 + measured-ns: 11916 (flashinfer 0.6.18, NVIDIA H200, 2026-08-25) + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def norm_quant( + a: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 4, 32), + names=("row", "warp", "lane"), + ) as mesh: + a_reg = tf.reshard( + a, + (ROWS @ mesh.row, 12 @ mesh.warp, 128 @ mesh.lane), + "rmem", + ) + weight_reg = tf.reshard(weight, (HIDDEN,), "rmem") + a_smem = tf.reshard( + a_reg, + (ROWS @ mesh.row, 12 @ mesh.warp, 128 @ mesh.lane), + "smem", + ) + weight_smem = tf.reshard(weight_reg, (HIDDEN,), "smem") + scale_smem = tf.reshard(scale, (1,), "smem") + normalized = tf.rms_norm(a_smem, weight_smem, eps=1e-6) + quantized = tf.cast(tf.cast(normalized, "f32") / scale_smem, "fp8e4m3") + return tf.reshard( + quantized, + (ROWS @ mesh.row, 12 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ) + + +__all__ = ["RMSNormQuant"] diff --git a/tests/fixtures/flashinfer/norm_cuh.py b/tests/fixtures/flashinfer/norm_cuh.py new file mode 100644 index 00000000..3d546605 --- /dev/null +++ b/tests/fixtures/flashinfer/norm_cuh.py @@ -0,0 +1,71 @@ +"""Placed CUDA general LayerNorm kernel with a split reduction axis. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +include/flashinfer/norm.cuh +license: Apache-2.0 (no upstream source is vendored) +The specialization keeps gamma, beta, FP8 output, and hidden-axis vectorization +while spelling out sum and sumsq reduction state. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 256, 1024 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 32 * 32)) + + +@module(entry="general_layer_norm", target=TARGET, topologies=TOPOLOGIES) +class GeneralLayerNormKernel: + """FlashInfer GeneralLayerNorm kernel. + + predicted-ns: 700 + waves: 2 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def general_layer_norm( + x: Tensor[(ROWS, HIDDEN), "bf16"], + gamma: ConstTensor[(HIDDEN,), "f32"], + beta: ConstTensor[(HIDDEN,), "f32"], + scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 32, 32), + names=("row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, (ROWS @ mesh.row, 32 @ mesh.warp, 32 @ mesh.lane), "rmem" + ) + x_smem = tf.reshard( + x_reg, (ROWS @ mesh.row, 32 @ mesh.warp, 32 @ mesh.lane), "smem" + ) + gamma_smem = tf.reshard(gamma, (HIDDEN,), "smem") + beta_smem = tf.reshard(beta, (HIDDEN,), "smem") + scale_smem = tf.reshard(scale, (1,), "smem") + x32 = tf.cast(x_smem, "f32") + count = tf.reduce( + tf.full_like(x32, value=1.0), axes=(-1,), keepdim=True, kind="sum" + ) + summed = tf.reduce(x32, axes=(-1,), keepdim=True, kind="sum") + sumsq = tf.reduce(x32 * x32, axes=(-1,), keepdim=True, kind="sum") + mean = summed / count + variance = sumsq / count - mean * mean + normalized = (x32 - mean) * tf.rsqrt(variance + 1e-6) + affine = normalized * gamma_smem + beta_smem + quantized = tf.cast(affine / scale_smem, "fp8e4m3") + return tf.reshard( + quantized, + (ROWS @ mesh.row, 32 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ) + + +__all__ = ["GeneralLayerNormKernel"] diff --git a/tests/fixtures/flashinfer/norm_kernels_fused_add_rmsnorm.py b/tests/fixtures/flashinfer/norm_kernels_fused_add_rmsnorm.py new file mode 100644 index 00000000..ae9c5917 --- /dev/null +++ b/tests/fixtures/flashinfer/norm_kernels_fused_add_rmsnorm.py @@ -0,0 +1,142 @@ +"""Placed CuTe fused residual-add and RMSNorm kernels. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/norm/kernels/fused_add_rmsnorm.py +license: Apache-2.0 (no upstream source is vendored) + +Each 128-thread CTA owns four rows, one row per warp, matching the CuTe +``rows_per_block`` dispatch for hidden size 1024. Functional tuple returns +preserve the upstream in-place writes. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 256, 1024 +CTA_COUNT, ROWS_PER_CTA = 64, 4 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", CTA_COUNT), Topology("thread", ROWS_PER_CTA * 32)) + + +@module(entry="fused_add_rmsnorm_cute", target=TARGET, topologies=TOPOLOGIES) +class FusedAddRMSNormKernel: + """FlashInfer CuTe fused residual-add and RMSNorm kernel. + + predicted-ns: 1027 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def fused_add_rmsnorm_cute( + x: Tensor[(ROWS, HIDDEN), "bf16"], + residual: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(CTA_COUNT, ROWS_PER_CTA, 32), + names=("block", "row", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "rmem", + ) + residual_reg = tf.reshard( + residual, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "smem", + ) + residual_smem = tf.reshard( + residual_reg, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "smem", + ) + weight_smem = tf.reshard(weight, (HIDDEN,), "smem") + summed = x_smem + residual_smem + normalized = tf.rms_norm(summed, weight_smem, eps=1e-6) + return ( + tf.reshard( + normalized, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "gmem", + ), + tf.reshard( + summed, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "gmem", + ), + ) + + +@module(entry="fused_add_rmsnorm_quant_cute", target=TARGET, topologies=TOPOLOGIES) +class FusedAddRMSNormQuantKernel: + """FlashInfer CuTe fused residual-add, RMSNorm, and FP8 kernel. + + predicted-ns: 934 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def fused_add_rmsnorm_quant_cute( + x: Tensor[(ROWS, HIDDEN), "bf16"], + residual: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(CTA_COUNT, ROWS_PER_CTA, 32), + names=("block", "row", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "rmem", + ) + residual_reg = tf.reshard( + residual, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "smem", + ) + residual_smem = tf.reshard( + residual_reg, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "smem", + ) + weight_smem = tf.reshard(weight, (HIDDEN,), "smem") + scale_smem = tf.reshard(scale, (1,), "smem") + summed = x_smem + residual_smem + normalized = tf.rms_norm(summed, weight_smem, eps=1e-6) + quantized = tf.cast(tf.cast(normalized, "f32") / scale_smem, "fp8e4m3") + return ( + tf.reshard( + quantized, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "gmem", + ), + tf.reshard( + summed, + (ROWS @ (mesh.block, mesh.row), HIDDEN @ mesh.lane), + "gmem", + ), + ) + + +__all__ = ["FusedAddRMSNormKernel", "FusedAddRMSNormQuantKernel"] diff --git a/tests/fixtures/flashinfer/norm_kernels_layernorm.py b/tests/fixtures/flashinfer/norm_kernels_layernorm.py new file mode 100644 index 00000000..89a465a7 --- /dev/null +++ b/tests/fixtures/flashinfer/norm_kernels_layernorm.py @@ -0,0 +1,71 @@ +"""Placed CuTe LayerNorm kernel with explicit distributed statistics. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/norm/kernels/layernorm.py +license: Apache-2.0 (no upstream source is vendored) +The hidden dimension stays vectorized. Separate sum and sumsq reductions carry +the two statistics needed to combine sharded LayerNorm state. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 256, 1024 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 4 * 32)) + + +@module(entry="layernorm_cute", target=TARGET, topologies=TOPOLOGIES) +class LayerNormKernel: + """FlashInfer CuTe DSL LayerNorm kernel. + + predicted-ns: 748 + waves: 2 + measured-ns: not taken + note: recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def layernorm_cute( + x: Tensor[(ROWS, HIDDEN), "bf16"], + gamma: ConstTensor[(HIDDEN,), "f32"], + beta: ConstTensor[(HIDDEN,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 4, 32), + names=("row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, (ROWS @ mesh.row, 8 @ mesh.warp, 128 @ mesh.lane), "rmem" + ) + x_smem = tf.reshard( + x_reg, (ROWS @ mesh.row, 8 @ mesh.warp, 128 @ mesh.lane), "smem" + ) + gamma_smem = tf.reshard(gamma, (HIDDEN,), "smem") + beta_smem = tf.reshard(beta, (HIDDEN,), "smem") + x32 = tf.cast(x_smem, "f32") + count = tf.reduce( + tf.full_like(x32, value=1.0), + axes=(-1,), + keepdim=True, + kind="sum", + ) + summed = tf.reduce(x32, axes=(-1,), keepdim=True, kind="sum") + sumsq = tf.reduce(x32 * x32, axes=(-1,), keepdim=True, kind="sum") + mean = summed / count + variance = sumsq / count - mean * mean + normalized = (x32 - mean) * tf.rsqrt(variance + 1e-6) + affine = normalized * gamma_smem + beta_smem + return tf.reshard( + tf.cast(affine, "bf16"), + (ROWS @ mesh.row, 8 @ mesh.warp, 128 @ mesh.lane), + "gmem", + ) + + +__all__ = ["LayerNormKernel"] diff --git a/tests/fixtures/flashinfer/norm_kernels_rmsnorm.py b/tests/fixtures/flashinfer/norm_kernels_rmsnorm.py new file mode 100644 index 00000000..627d3d6a --- /dev/null +++ b/tests/fixtures/flashinfer/norm_kernels_rmsnorm.py @@ -0,0 +1,169 @@ +"""Placed CuTe RMSNorm, QK RMSNorm, and fixed-scale FP8 kernels. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/norm/kernels/rmsnorm.py +license: Apache-2.0 (no upstream source is vendored) + +The selected specializations preserve the CuTe multi-row dispatch: independent +warps own rows or heads, while lanes own the reduction dimension. The QK +variant retains its three-dimensional batch, head, and head-dimension layout. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +RMS_ROWS, RMS_HIDDEN = 128, 1024 +QUANT_ROWS, QUANT_HIDDEN = 64, 4096 +BATCH, HEADS, HEAD_DIM = 4, 32, 128 +RMS_CTA_COUNT, RMS_ROWS_PER_CTA = 32, 4 +QUANT_CTA_COUNT, QUANT_ROWS_PER_CTA = 32, 2 +QK_CTA_COUNT, QK_ROWS_PER_CTA = 16, 8 +TARGET = CudaTarget("nvidia.h200_sxm") +RMS_TOPOLOGIES = (Topology("cta", RMS_CTA_COUNT), Topology("thread", 4 * 32)) +QUANT_TOPOLOGIES = (Topology("cta", QUANT_CTA_COUNT), Topology("thread", 4 * 32)) +QK_TOPOLOGIES = (Topology("cta", QK_CTA_COUNT), Topology("thread", 4 * 32)) + + +@module(entry="rmsnorm_cute", target=TARGET, topologies=RMS_TOPOLOGIES) +class RMSNormKernel: + """FlashInfer CuTe row-wise RMSNorm kernel. + + predicted-ns: 574 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def rmsnorm_cute( + x: Tensor[(RMS_ROWS, RMS_HIDDEN), "bf16"], + weight: ConstTensor[(RMS_HIDDEN,), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(RMS_CTA_COUNT, RMS_ROWS_PER_CTA, 32), + names=("block", "row", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + (RMS_ROWS @ (mesh.block, mesh.row), RMS_HIDDEN @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + (RMS_ROWS @ (mesh.block, mesh.row), RMS_HIDDEN @ mesh.lane), + "smem", + ) + weight_smem = tf.reshard(weight, (RMS_HIDDEN,), "smem") + normalized = tf.rms_norm(x_smem, weight_smem, eps=1e-6) + return tf.reshard( + normalized, + (RMS_ROWS @ (mesh.block, mesh.row), RMS_HIDDEN @ mesh.lane), + "gmem", + ) + + +@module(entry="qk_rmsnorm_cute", target=TARGET, topologies=QK_TOPOLOGIES) +class QKRMSNormKernel: + """FlashInfer CuTe three-dimensional QK RMSNorm kernel. + + predicted-ns: 139 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def qk_rmsnorm_cute( + x: Tensor[(BATCH, HEADS, HEAD_DIM), "bf16"], + weight: ConstTensor[(HEAD_DIM,), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(BATCH, HEADS // QK_ROWS_PER_CTA, QK_ROWS_PER_CTA, 16), + names=("batch", "head_group", "row", "lane_in_row"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + BATCH @ mesh.batch, + HEADS @ (mesh.head_group, mesh.row), + HEAD_DIM @ mesh.lane_in_row, + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + BATCH @ mesh.batch, + HEADS @ (mesh.head_group, mesh.row), + HEAD_DIM @ mesh.lane_in_row, + ), + "smem", + ) + weight_smem = tf.reshard(weight, (HEAD_DIM,), "smem") + normalized = tf.rms_norm(x_smem, weight_smem, eps=1e-6) + return tf.reshard( + normalized, + ( + BATCH @ mesh.batch, + HEADS @ (mesh.head_group, mesh.row), + HEAD_DIM @ mesh.lane_in_row, + ), + "gmem", + ) + + +@module(entry="rmsnorm_quant_cute", target=TARGET, topologies=QUANT_TOPOLOGIES) +class RMSNormQuantKernel: + """FlashInfer CuTe RMSNorm and fixed-scale FP8 kernel. + + predicted-ns: 1069 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def rmsnorm_quant_cute( + x: Tensor[(QUANT_ROWS, QUANT_HIDDEN), "bf16"], + weight: ConstTensor[(QUANT_HIDDEN,), "bf16"], + scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(QUANT_CTA_COUNT, QUANT_ROWS_PER_CTA, 2, 32), + names=("block", "row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, + ( + QUANT_ROWS @ (mesh.block, mesh.row), + QUANT_HIDDEN @ (mesh.warp, mesh.lane), + ), + "rmem", + ) + x_smem = tf.reshard( + x_reg, + ( + QUANT_ROWS @ (mesh.block, mesh.row), + QUANT_HIDDEN @ (mesh.warp, mesh.lane), + ), + "smem", + ) + weight_smem = tf.reshard(weight, (QUANT_HIDDEN,), "smem") + scale_smem = tf.reshard(scale, (1,), "smem") + normalized = tf.rms_norm(x_smem, weight_smem, eps=1e-6) + quantized = tf.cast(tf.cast(normalized, "f32") / scale_smem, "fp8e4m3") + return tf.reshard( + quantized, + ( + QUANT_ROWS @ (mesh.block, mesh.row), + QUANT_HIDDEN @ (mesh.warp, mesh.lane), + ), + "gmem", + ) + + +__all__ = ["RMSNormKernel", "QKRMSNormKernel", "RMSNormQuantKernel"] diff --git a/tests/fixtures/flashinfer/norm_quant_rope.blocked.py b/tests/fixtures/flashinfer/norm_quant_rope.blocked.py new file mode 100644 index 00000000..2b7c801f --- /dev/null +++ b/tests/fixtures/flashinfer/norm_quant_rope.blocked.py @@ -0,0 +1,49 @@ +"""Complete equation inventory for FlashInfer norm/quant/RoPE boundaries. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c58fdd2392914ea05e2a8714946ac0eef6 +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: selection/analysis +error: source defines no TileFoundry Module +classification: expected-spec notation; no authored Module is declared. +Notation retains storage handoffs, placement variants, and capability gaps. +""" + +# noqa +N01 = """def rmsnorm_quant(x:gmem, gamma:gmem, scale:gmem):\n y = rms_norm(x, gamma, eps=1e-6)\n return cast(y / scale, fp8_e4m3fn) # [program,cta-smem,thread-rmem]\n""" +N02 = """def add_rmsnorm(x:gmem, residual:gmem, gamma:gmem):\n s = add(x, residual) -> smem\n y = rms_norm(s, gamma, eps=1e-6)\n return (s, y) # [program,cta-smem,thread-rmem]\n""" +N03 = """def add_rmsnorm_quant(x:gmem, residual:gmem, gamma:gmem, scale:gmem):\n s = add(x, residual) -> smem\n y = rms_norm(s, gamma, eps=1e-6)\n q = cast(y / scale, fp8_e4m3fn)\n return (s, q) # [program,cta-smem,thread-rmem]\n""" +N04 = """def gemma_add_rmsnorm(x:gmem, residual:gmem, gamma:gmem):\n s = add(x, residual) -> smem\n return s, rms_norm(s, add(gamma, 1), eps=1e-6)\n""" +N05 = """def layernorm_quant(x:gmem, gamma:gmem, beta:gmem, scale:gmem):\n y = layer_norm(x, gamma, beta, axis=-1, eps=1e-5)\n return cast(y / scale, fp8_e4m3fn)\n""" +N06 = """def rmsnorm_silu(x:gmem, gamma:gmem):\n y = rms_norm(x, gamma, eps=1e-6) -> smem\n return silu(y)\n""" +N07 = """def rmsnorm_silu_fp8(x:gmem, gamma:gmem):\n y = rms_norm(x, gamma, eps=1e-6) -> smem\n return cast(silu(y), fp8_e4m3fn)\n""" +N08 = """def rmsnorm_silu_nvfp4(x:gmem, gamma:gmem):\n y = rms_norm(x, gamma, eps=1e-6) -> smem\n return block_quant(silu(y), dtype=nvfp4, block=16)\n""" +N09 = """def qk_rmsnorm_rope(q:gmem, k:gmem, v:gmem, gq:gmem, gk:gmem, cos:gmem, sin:gmem):\n qr = rope(rms_norm(q,gq), rms_norm(k,gk), cos, sin) -> smem\n return qr.q, qr.k, v\n""" +N10 = N09.replace( + "return qr.q, qr.k, v", + "return quant(qr.q,fp8_e4m3fn), quant(qr.k,fp8_e4m3fn), quant(v,fp8_e4m3fn)", +) +N11 = """def dit_residual_ln(x:gmem, residual:gmem, scale:gmem, shift:gmem):\n s = add(x,residual) -> smem\n return layer_norm(s) * (1 + scale) + shift\n""" +N12 = """def dit_gate_residual_ln(x:gmem, residual:gmem, gate:gmem, scale:gmem, shift:gmem):\n s = add(mul(x,gate),residual) -> smem\n return layer_norm(s) * (1 + scale) + shift\n""" +N13 = """def dit_gate_residual_ln_gamma_beta(x:gmem, residual:gmem, gate:gmem, bias:gmem, gamma:gmem, beta:gmem):\n s = add(mul(x,add(gate,bias)),residual) -> smem\n return layer_norm(s,gamma,beta,axis=-1,eps=1e-5)\n""" + +# noqa +Q01 = """def silu_mul_nvfp4(gate:gmem, up:gmem):\n z = mul(silu(gate), up) -> smem\n return block_quant(z, dtype=nvfp4, block=16)\n""" +Q02 = """def scaled_silu_mul_nvfp4(gate:gmem, up:gmem, expert_scale:gmem):\n z = mul(silu(gate), up) * gather(expert_scale) -> smem\n return block_quant(z, dtype=nvfp4, block=16)\n""" +Q03 = """def smooth_nvfp4(x:gmem, pre_scale:gmem):\n return block_quant(x * pre_scale, dtype=nvfp4, block=16)\n""" +Q04 = """def mxfp4(x:gmem):\n scale = block_absmax(x, block=32) -> smem\n return pack_fp4(x / scale), scale\n""" +Q05 = """def mxfp8(x:gmem):\n scale = block_absmax(x, block=32) -> smem\n return cast(x / scale, fp8_e4m3fn), scale\n""" + +# noqa +R01 = """def rope_quantize(q:gmem, k:gmem, cos:gmem, sin:gmem):\n qr = rope(q,k,cos,sin) -> smem\n return quant(qr.q,fp8_e4m3fn), quant(qr.k,fp8_e4m3fn)\n""" +R02 = """def mla_rope_quantize(q:gmem, k_rank2:gmem, cos:gmem, sin:gmem):\n qr = rope(q,k_rank2,cos,sin) -> smem\n return quant(qr.q,fp8_e4m3fn), quant(qr.k,fp8_e4m3fn)\n""" +R03 = """def rope_quant_append(q:gmem,k:gmem,cos:gmem,sin:gmem,cache:gmem,slots:gmem):\n qr = rope(q,k,cos,sin) -> smem\n q8,k8 = quant(qr,fp8_e4m3fn) -> rmem\n cache[slots] = k8\n return q8, cache\n""" +R04 = """def nvfp4_append(k:gmem,v:gmem,cache:gmem,slots:gmem):\n kq = block_quant(k,dtype=nvfp4,block=16) -> rmem\n vq = block_quant(v,dtype=nvfp4,block=16) -> rmem\n cache[slots] = (kq,vq)\n return cache\n""" +R05 = """def nvfp4_append_slot(k:gmem,v:gmem,cache:gmem,slot_map:gmem):\n kq,vq = block_quant((k,v),dtype=nvfp4,block=16) -> rmem\n cache[slot_map] = (kq,vq)\n return cache\n""" + +ALL = { + **{f"N{i:02d}": globals()[f"N{i:02d}"] for i in range(1, 14)}, + **{f"Q{i:02d}": globals()[f"Q{i:02d}"] for i in range(1, 6)}, + **{f"R{i:02d}": globals()[f"R{i:02d}"] for i in range(1, 6)}, +} diff --git a/tests/fixtures/flashinfer/paged_gather.blocked.py b/tests/fixtures/flashinfer/paged_gather.blocked.py new file mode 100644 index 00000000..8ec067c1 --- /dev/null +++ b/tests/fixtures/flashinfer/paged_gather.blocked.py @@ -0,0 +1,36 @@ +"""Paged KV gather whose base index selection reports the wrong traffic. + +Notes: +upstream: flashinfer-ai/flashinfer@2ab910c5 page.cu:gather_paged_kv +license: Apache-2.0 +blocked: mis-analyzed +phase: selection/analysis +got: traffic traffic=gmem:r136/w32@r136/w32 +expected: only the selected page bytes +why: the access relation charges the whole cache instead of selected pages +ledger: OP-04 +placement_refusal: IndexSelect rejects the sharded smem operand on dim 0. +placement_error: IndexSelect: dim 0 index_select over a shard layout with multiple +Split axes including the selected dim; cannot derive an output layout +placement_workaround: materialize the staged tensor back to gmem before IndexSelect. +""" + +from tilefoundry import module +from tilefoundry.dsl import Mesh, Tensor, Topology, func, tf +from tilefoundry.target import CudaTarget + + +@module( + entry="paged_gather", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 2), Topology("thread", 4)), +) +class PagedGather: + @func + def paged_gather(cache: Tensor[(8, 4), "f32"], indices: Tensor[(2,), "i32"]): + with Mesh(("cta",), layout=(2,), names=("tile",)) as cta: + with Mesh(("thread",), layout=(4,), names=("lane",)) as thread: + cache_rmem = tf.reshard(cache, (2 @ cta.tile, 1, 4 @ thread.lane), "rmem") + cache_smem = tf.reshard(cache_rmem, (2 @ cta.tile, 1, 4 @ thread.lane), "smem") + cache_for_gather = tf.reshard(cache_smem, (8, 4), "gmem") + return tf.index_select(cache_for_gather, indices, dim=0) diff --git a/tests/fixtures/flashinfer/paged_gather_placement.blocked.py b/tests/fixtures/flashinfer/paged_gather_placement.blocked.py new file mode 100644 index 00000000..bfd6a9fa --- /dev/null +++ b/tests/fixtures/flashinfer/paged_gather_placement.blocked.py @@ -0,0 +1,30 @@ +"""Direct paged gather placement that IndexSelect currently rejects. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c58fdd2392914ea05e2a8714946ac0eef6 +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: IndexSelect: dim 0 index_select over a shard layout +ledger: EXT-01a +placement_refusal: IndexSelect rejects the sharded smem operand on dim 0. +placement_workaround: materialize the staged tensor back to gmem before IndexSelect. +""" + +from tilefoundry import module +from tilefoundry.dsl import Mesh, Tensor, Topology, func, tf +from tilefoundry.target import CudaTarget + + +@module( + entry="paged_gather_placement", + target=CudaTarget("nvidia.h200_sxm"), + topologies=(Topology("cta", 2), Topology("thread", 4)), +) +class PagedGatherPlacement: + @func + def paged_gather_placement(cache: Tensor[(8, 4), "f32"], indices: Tensor[(2,), "i32"]): + with Mesh(("cta",), layout=(2,), names=("tile",)) as cta: + with Mesh(("thread",), layout=(4,), names=("lane",)) as thread: + cache_smem = tf.reshard(cache, (2 @ cta.tile, 1, 4 @ thread.lane), "smem") + return tf.index_select(cache_smem, indices, dim=0) diff --git a/tests/fixtures/flashinfer/routing_gather.blocked.py b/tests/fixtures/flashinfer/routing_gather.blocked.py new file mode 100644 index 00000000..30e833ee --- /dev/null +++ b/tests/fixtures/flashinfer/routing_gather.blocked.py @@ -0,0 +1,81 @@ +"""FlashInfer M01 routing needs a batch-aware gather surface. + +Notes: +upstream: flashinfer-ai/flashinfer @ 2ab910c58fdd2392914ea05e2a8714946ac0eef6 +license: Apache-2.0 (no upstream source is vendored) +blocked: refused +phase: load +error: runtime_expression: unsupported call 'tf.gather' (2 positional, keywords ['axis', 'batch_dims']) +ledger: EXT-01 +The three golden level variants preserve the routing equations; ``tf.gather`` +with ``batch_dims=1`` has no current authored-HIR surface. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN, K = 128, 16, 8 +DT = "bf16" +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", 8), Topology("thread", 128)) + + +@module(entry="fused", target=TARGET, topologies=TOPOLOGIES) +class RoutingProgram: + @func + def score(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + return tf.sigmoid(x) + bias + + @func + def unfused(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + s = score(x, bias) + _, indices = tf.topk(s, k=K, axis=-1) + return tf.gather(s, indices, axis=-1, batch_dims=1) + + @func + def fused(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + s = tf.sigmoid(x) + bias + _, indices = tf.topk(s, k=K, axis=-1, largest=True, sorted=True) + return tf.gather(s, indices, axis=-1, batch_dims=1) + + +@module(entry="fused", target=TARGET, topologies=TOPOLOGIES) +class RoutingCTA: + @func + def score(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + return tf.sigmoid(x) + bias + + @func + def unfused(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + s = score(x, bias) + _, indices = tf.topk(s, k=K, axis=-1) + return tf.gather(s, indices, axis=-1, batch_dims=1) + + @func + def fused(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + s = tf.sigmoid(x) + bias + _, indices = tf.topk(s, k=K, axis=-1, largest=True, sorted=True) + return tf.gather(s, indices, axis=-1, batch_dims=1) + + +@module(entry="fused", target=TARGET, topologies=TOPOLOGIES) +class RoutingThread: + @func + def score(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + return tf.sigmoid(x) + bias + + @func + def unfused(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + s = score(x, bias) + _, indices = tf.topk(s, k=K, axis=-1) + return tf.gather(s, indices, axis=-1, batch_dims=1) + + @func + def fused(x: Tensor[(ROWS, HIDDEN), DT], bias: ConstTensor[(HIDDEN,), DT]): + s = tf.sigmoid(x) + bias + _, indices = tf.topk(s, k=K, axis=-1, largest=True, sorted=True) + return tf.gather(s, indices, axis=-1, batch_dims=1) diff --git a/tests/fixtures/flashinfer/triton_kernels_norm.py b/tests/fixtures/flashinfer/triton_kernels_norm.py new file mode 100644 index 00000000..a1dc64aa --- /dev/null +++ b/tests/fixtures/flashinfer/triton_kernels_norm.py @@ -0,0 +1,86 @@ +"""Placed Triton RMSNorm kernel with residual and scaling specializations. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/triton/kernels/norm.py +license: Apache-2.0 (no upstream source is vendored) + +The concrete specialization keeps input dequantization, residual writeback, +FP8 output scaling, and one row per CTA from the parameterized Triton kernel. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 128, 2048 +TARGET = CudaTarget("nvidia.h200_sxm") +TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 32 * 32)) + + +@module(entry="rms_norm_kernel", target=TARGET, topologies=TOPOLOGIES) +class TritonRMSNormKernel: + """FlashInfer Triton scaled residual RMSNorm kernel. + + predicted-ns: 521 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def rms_norm_kernel( + x: Tensor[(ROWS, HIDDEN), "fp8e4m3"], + residual: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + input_scale: ConstTensor[(1,), "f32"], + output_scale: ConstTensor[(1,), "f32"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 32, 32), + names=("row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), "rmem" + ) + residual_reg = tf.reshard( + residual, + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), "smem" + ) + residual_smem = tf.reshard( + residual_reg, + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "smem", + ) + weight_smem = tf.reshard(weight, (HIDDEN,), "smem") + input_scale_smem = tf.reshard(input_scale, (1,), "smem") + output_scale_smem = tf.reshard(output_scale, (1,), "smem") + dequantized = tf.cast(x_smem, "f32") * input_scale_smem + summed = dequantized + tf.cast(residual_smem, "f32") + normalized = tf.rms_norm( + tf.cast(summed, "bf16"), + weight_smem, + eps=1e-6, + ) + quantized = tf.cast(tf.cast(normalized, "f32") * output_scale_smem, "fp8e4m3") + return ( + tf.reshard( + quantized, + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ), + tf.reshard( + tf.cast(summed, "bf16"), + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ), + ) + + +__all__ = ["TritonRMSNormKernel"] diff --git a/tests/fixtures/flashinfer/triton_norm.py b/tests/fixtures/flashinfer/triton_norm.py new file mode 100644 index 00000000..e43d1695 --- /dev/null +++ b/tests/fixtures/flashinfer/triton_norm.py @@ -0,0 +1,111 @@ +"""Placed public Triton RMSNorm wrappers. + +upstream: flashinfer-ai/flashinfer@2ab910c58fdd2392914ea05e2a8714946ac0eef6 \ +flashinfer/triton/norm.py +license: Apache-2.0 (no upstream source is vendored) + +The two independent modules preserve the plain output API and the residual +wrapper's paired output and in-place residual effects. +""" + +from __future__ import annotations + +from tilefoundry import func, module +from tilefoundry.dsl import ConstTensor, Tensor, tf +from tilefoundry.ir.types.shard import Mesh, Topology +from tilefoundry.target import CudaTarget + +ROWS, HIDDEN = 128, 2048 +TARGET = CudaTarget("nvidia.h200_sxm") +PLAIN_TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 8 * 32)) +RESIDUAL_TOPOLOGIES = (Topology("cta", ROWS), Topology("thread", 32 * 32)) + + +@module(entry="rms_norm", target=TARGET, topologies=PLAIN_TOPOLOGIES) +class TritonRMSNorm: + """FlashInfer public Triton RMSNorm wrapper kernel. + + predicted-ns: 372 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def rms_norm( + x: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 8, 32), + names=("row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, (ROWS @ mesh.row, 8 @ mesh.warp, 256 @ mesh.lane), "rmem" + ) + x_smem = tf.reshard( + x_reg, (ROWS @ mesh.row, 8 @ mesh.warp, 256 @ mesh.lane), "smem" + ) + weight_smem = tf.reshard(weight, (HIDDEN,), "smem") + normalized = tf.rms_norm(x_smem, weight_smem, eps=1e-6) + return tf.reshard( + normalized, + (ROWS @ mesh.row, 8 @ mesh.warp, 256 @ mesh.lane), + "gmem", + ) + + +@module(entry="rms_norm_add_residual", target=TARGET, topologies=RESIDUAL_TOPOLOGIES) +class TritonRMSNormAddResidual: + """FlashInfer public Triton residual-add RMSNorm wrapper kernel. + + predicted-ns: 599 waves: 1 + measured-ns: not taken + note: Recorded evidence, not an assertion; open a ledger row if predicted exceeds measured. + """ + + @func + def rms_norm_add_residual( + x: Tensor[(ROWS, HIDDEN), "bf16"], + residual: Tensor[(ROWS, HIDDEN), "bf16"], + weight: ConstTensor[(HIDDEN,), "bf16"], + ): + with Mesh( + ("cta", "thread"), + layout=(ROWS, 32, 32), + names=("row", "warp", "lane"), + ) as mesh: + x_reg = tf.reshard( + x, (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), "rmem" + ) + residual_reg = tf.reshard( + residual, + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "rmem", + ) + x_smem = tf.reshard( + x_reg, (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), "smem" + ) + residual_smem = tf.reshard( + residual_reg, + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "smem", + ) + weight_smem = tf.reshard(weight, (HIDDEN,), "smem") + summed = x_smem + residual_smem + normalized = tf.rms_norm(summed, weight_smem, eps=1e-6) + return ( + tf.reshard( + normalized, + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ), + tf.reshard( + summed, + (ROWS @ mesh.row, 64 @ mesh.warp, 32 @ mesh.lane), + "gmem", + ), + ) + + +__all__ = ["TritonRMSNorm", "TritonRMSNormAddResidual"] diff --git a/tests/fixtures/test_corpus.py b/tests/fixtures/test_corpus.py new file mode 100644 index 00000000..0996623e --- /dev/null +++ b/tests/fixtures/test_corpus.py @@ -0,0 +1,126 @@ +"""Run the fixture corpus through the authored HIR command boundary.""" + +from __future__ import annotations + +import ast +import subprocess +import sys +from pathlib import Path + +import pytest + +from tilefoundry.cli.source import load_namespace +from tilefoundry.ir.core.module import Module + +CORPUS = Path(__file__).parent / "flashinfer" +CLI = Path(sys.executable).with_name("tilefoundry") + + +def _fixtures(pattern: str) -> tuple[Path, ...]: + return tuple(sorted(CORPUS.glob(pattern))) + + +def _ok_fixtures() -> tuple[Path, ...]: + return tuple( + path + for path in _fixtures("*.py") + if path.name != "__init__.py" and not path.name.endswith(".blocked.py") + ) + + +def _doc_fields(path: Path) -> dict[str, str]: + document = ast.get_docstring(ast.parse(path.read_text(encoding="utf-8"))) + assert document is not None + fields: dict[str, str] = {} + for line in document.splitlines(): + key, separator, value = line.partition(":") + if separator and key in {"blocked", "phase", "error", "got", "expected", "why"}: + fields[key] = value.strip() + return fields + + +def _ok_sources() -> tuple[str, ...]: + sources = [] + for path in _ok_fixtures(): + namespace, _ = load_namespace(str(path)) + modules = sorted(value.name for value in namespace.values() if isinstance(value, Module)) + sources.extend(f"{path}:{name}" for name in modules) + return tuple(sources) + + +def _analyze(source: str, report: Path) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + str(CLI), + "analyze", + source, + str(report), + "--compute-cost", + "--memory", + "--roofline", + "--performance", + "--topology", + "cta", + ], + capture_output=True, + text=True, + check=False, + ) + + +OK_SOURCES = _ok_sources() + + +@pytest.mark.parametrize( + "source", + OK_SOURCES, + ids=lambda source: source.removeprefix(f"{CORPUS}/"), +) +def test_ok_fixtures_load_and_analyze(source: str, tmp_path: Path) -> None: + """Every Module in every unblocked fixture must load and produce a report.""" + selector = source.rsplit(":", maxsplit=1)[-1] + report = tmp_path / f"{selector}.report" + completed = _analyze(source, report) + assert completed.returncode == 0, completed.stderr + assert completed.stdout == "" + assert completed.stderr == "" + assert report.is_file() + assert "# analysis" in report.read_text(encoding="utf-8") + + +@pytest.mark.parametrize("path", _fixtures("*.blocked.py"), ids=lambda path: path.name) +def test_blocked_fixtures_preserve_their_observed_phase_and_result( + path: Path, tmp_path: Path +) -> None: + """Blocked fixtures must fail or report exactly where their docstring says.""" + fields = _doc_fields(path) + assert fields["phase"] in {"load", "selection/analysis"} + report = tmp_path / f"{path.stem}.report" + prefix = path.stem.partition(".")[0] + before = {name for name in sys.modules if name.partition(".")[0] == prefix} + try: + load_namespace(str(path)) + loaded = True + except Exception: + loaded = False + assert {name for name in sys.modules if name.partition(".")[0] == prefix} == before + + completed = _analyze(str(path), report) + if fields["blocked"] == "refused": + if fields["phase"] == "load": + assert not loaded + else: + assert fields["phase"] == "selection/analysis" + assert loaded + assert completed.returncode != 0 + assert not report.exists() + assert fields["error"] in completed.stderr + return + + assert fields["blocked"] == "mis-analyzed" + assert fields["phase"] == "selection/analysis" + assert loaded + assert completed.returncode == 0, completed.stderr + assert completed.stdout == "" + assert report.is_file() + assert fields["got"] in report.read_text(encoding="utf-8")