diff --git a/.gitattributes b/.gitattributes index 4381cbe52236..b7766e7fff52 100644 --- a/.gitattributes +++ b/.gitattributes @@ -20,3 +20,5 @@ cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cubin/xqa_kernel_cubin. docs/source/blogs/media/tech_blog26_deepseek_v4_hybrid_attention.png filter=lfs diff=lfs merge=lfs -text docs/source/blogs/media/tech_blog26_deepseek_v4_mhc_moe.png filter=lfs diff=lfs merge=lfs -text docs/source/blogs/media/tech_blog28_video_*.gif filter=lfs diff=lfs merge=lfs -text +3rdparty/vendor_patches/*.patch -whitespace +3rdparty/patches/*.patch -whitespace diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index c9acc9b61ece..6471c71c878c 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -472,6 +472,9 @@ /.github/CODEOWNERS @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance /.github/tava_architecture_diagram.md @NVIDIA/trt-llm-TAVA-design-change /3rdparty/** @NVIDIA/trt-llm-oss-compliance +/3rdparty/vendor_patches/** @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance +/3rdparty/vendor_patches/flashinfer-prims-ts.patch @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-oss-compliance +/3rdparty/vendor_sources.lock.yaml @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance /ATTRIBUTIONS-*.md @NVIDIA/trt-llm-oss-compliance /LICENSE @NVIDIA/trt-llm-oss-compliance /constraints.txt @NVIDIA/trt-llm-oss-compliance @@ -488,7 +491,9 @@ /requirements-grpc-smg.txt @NVIDIA/trt-llm-oss-compliance /requirements-openengine.txt @NVIDIA/trt-llm-oss-compliance /requirements.txt @NVIDIA/trt-llm-oss-compliance +/scripts/vendor_sources.py @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance /setup.py @NVIDIA/trt-llm-oss-compliance +/tensorrt_llm/_torch/attention/backends/prims_ts/** @NVIDIA/trt-llm-torch-attention-devs @NVIDIA/trt-llm-oss-compliance /tests/unittest/api_stability/ @NVIDIA/trt-llm-noncommitted-api-review-committee /tests/unittest/api_stability/references_committed/ @NVIDIA/trt-llm-committed-api-review-committee /triton_kernels/** @NVIDIA/trt-llm-oss-compliance diff --git a/.gitignore b/.gitignore index fd757612a86f..2579eb1ed3e2 100644 --- a/.gitignore +++ b/.gitignore @@ -86,6 +86,7 @@ ad-test-workspace/ */tllm_debug/** *.patch !cpp/tensorrt_llm/deep_ep/*.patch +!3rdparty/vendor_patches/*.patch examples/disaggregated/slurm/benchmark/logs/ scripts/attribution/data/checksum_to_paths.yml diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0f7bbe4b3a2e..ab8a7ab62c84 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1527,7 +1527,7 @@ legacy-files: &legacy_files | # list; the hook's own `files:` pattern only gates *when* the hook triggers. # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' +exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)' default_install_hook_types: [pre-commit, commit-msg] repos: @@ -1638,6 +1638,13 @@ repos: additional_dependencies: [jinja2] always_run: true pass_filenames: false + - id: vendor-sources-check + name: verify vendored sources are in sync + entry: python scripts/vendor_sources.py check + language: python + additional_dependencies: [PyYAML] + always_run: true + pass_filenames: false - id: test lists format name: Check for tabs and multiple spaces in test_lists txt files entry: ./scripts/format_test_list.py diff --git a/3rdparty/README.md b/3rdparty/README.md index 4b28139757b0..57113bd1fbee 100644 --- a/3rdparty/README.md +++ b/3rdparty/README.md @@ -1,8 +1,9 @@ # `3rdparty/` -This directory holds TensorRT-LLM's third-party C++ dependencies (driven by -cmake `FetchContent` from `fetch_content.json`) plus tooling that -accelerates repeat clones of those dependencies. +This directory holds TensorRT-LLM's third-party dependency metadata and +tooling. C++ dependencies are driven by CMake `FetchContent` from +`fetch_content.json`; source vendors use a generated lock and patches. It also +contains tooling that accelerates repeat clones of C++ dependencies. ## Adding new third-party dependencies @@ -14,6 +15,8 @@ dependency you want to add: and re-distributed with the wheel, see [cpp-thirdparty.md](cpp-thirdparty.md) * For python dependencies declared via wheel metadata and installed in the container via pip, see [py-thirdparty.md](py-thirdparty.md) +* For source trees copied into this repository and pinned to an upstream Git + commit, see [vendor-sources.md](vendor-sources.md) ## FetchContent cache (`--use-3rdparty-cache`) diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index 140e9c90f94d..136be361a4ae 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -718,3 +718,20 @@ index 306b416..a564d47 100644 out.copy_(result) return out, None return result, None +diff --git a/python/fmha_sm100/cute/src/common/cute_dsl_utils.py b/python/fmha_sm100/cute/src/common/cute_dsl_utils.py +index e3473fb..dc4cfb0 100644 +--- a/python/fmha_sm100/cute/src/common/cute_dsl_utils.py ++++ b/python/fmha_sm100/cute/src/common/cute_dsl_utils.py +@@ -128,10 +128,8 @@ if os.getenv("MINIMAX_LOG_COMPILE", "0") == "1": + logger.setLevel(logging.DEBUG) + + +-# Monkey-patch cute.compile so every JIT compile across the repo gets timed +-# without touching individual call sites. Idempotent: only patches once. +-if cute.compile is not cute_compile_patched: +- cute.compile = cute_compile_patched ++# TensorRT-LLM shares CUTLASS across attention backends, so importing MSA must ++# not replace the process-global cute.compile callable. + + + def assume_strides_aligned(t): diff --git a/3rdparty/vendor-sources.md b/3rdparty/vendor-sources.md new file mode 100644 index 000000000000..ad6939511da8 --- /dev/null +++ b/3rdparty/vendor-sources.md @@ -0,0 +1,285 @@ +# Vendored Sources + +TensorRT-LLM keeps some upstream source trees in this repository so they can be +built, packaged, and reviewed with the code that uses them. The generic +vendoring tool records where each tree came from, materializes it reproducibly, +and rejects destination edits that are not represented by its lock entry. + +The generated lock is `3rdparty/vendor_sources.lock.yaml`. It records an +upstream Git URL, an immutable commit, the source and destination directories, +the selected files, any persistent compatibility patch and its content digest, +and a digest of the materialized destination. A short branch or tag may explain +where the commit came from, but the full commit is authoritative. + +Use `scripts/vendor_sources.py` for every lock or vendor-state change. Do not +edit the YAML, generated patches, or digests by hand. All examples below use the +default lock. For an isolated test or another consumer repository, place +`--lock PATH` before the subcommand. + +## Lock contract + +A locked vendor has one of two durable states: + +- **Exact**: the selected destination files are byte-for-byte copies of the + selected files at the locked upstream commit. +- **Patched**: applying a deterministic, persistent compatibility patch to + those upstream files reproduces the destination exactly. Use this patch only + for TensorRT-LLM-specific adaptations that do not belong upstream. + +A destination edit is not a third state. While such an edit is pending, +`status`, the default offline `check`, and the pre-commit check intentionally +fail. Resolve it by discarding it with `sync`, recording a TensorRT-LLM-only +adaptation with `patch`, or exporting an upstream-worthy change and pinning the +resulting commit. `export` accepts this pending destination delta by default and +does not change the lock or persistent patch. + +## Choose a command + +```mermaid +flowchart TD + A{Lock entry exists?} + A -- No --> B{Destination exists?} + B -- No --> C[create] + B -- Yes --> D[create --adopt exact or patched] + A -- Yes --> E{What do you need?} + E -- Inspect --> F[list, status, or check] + E -- Restore locked bytes --> G[sync current immutable pin] + E -- Use a newer upstream commit --> H[Prepare matching destination, then pin] + E -- Destination changed --> I{Should the change go upstream?} + I -- No, TensorRT-LLM only --> J[patch create or refresh] + I -- Yes --> K[Temporary branch, export, commit and push, then pin] + E -- Stop vendoring --> L[remove] +``` + +`sync` only restores the commit and compatibility patch already recorded in the +lock. It never discovers, imports, or pins a newer upstream commit. To move to a +new upstream revision, first make the destination equal that revision plus the +existing compatibility patch, then use `pin`. + +## Inspect vendors + +List entries, or run the offline integrity status for all or one vendor: + +```bash +python scripts/vendor_sources.py list +python scripts/vendor_sources.py status +python scripts/vendor_sources.py status VENDOR +python scripts/vendor_sources.py check VENDOR +``` + +`status` and the default `check` exit unsuccessfully if the destination has a +pending delta. That failure is expected during an export workflow and remains +until `pin` succeeds. + +## Add or adopt a vendor + +When neither the lock entry nor destination exists, create both from an +immutable commit and a local upstream checkout: + +```bash +python scripts/vendor_sources.py create VENDOR \ + --url https://example.com/organization/repository.git \ + --branch main \ + --commit FULL_COMMIT \ + --source path/in/upstream \ + --destination path/in/tensorrt-llm \ + --include '**/*.py' \ + --repo /path/to/upstream +``` + +Use `--tag TAG` instead of `--branch BRANCH` for a tagged source. Without +`--repo`, the tool obtains the commit from the recorded URL. + +If the destination already exists but has no lock entry, adopt it. Use `exact` +to require an exact upstream match: + +```bash +python scripts/vendor_sources.py create VENDOR \ + --url https://example.com/organization/repository.git \ + --commit FULL_COMMIT \ + --source path/in/upstream \ + --destination path/in/tensorrt-llm \ + --include '**/*.py' \ + --adopt exact \ + --repo /path/to/upstream +``` + +Use `--adopt patched` instead to capture intentional TensorRT-LLM compatibility +adaptations. Adoption never silently accepts an unrepresented difference. + +## Restore the current lock + +Discard destination edits and reproduce the currently locked upstream commit +plus its persistent patch: + +```bash +python scripts/vendor_sources.py sync VENDOR --repo /path/to/upstream +``` + +This overwrites the selected destination files. It does not update the lock, +look at a branch tip, or choose a newer commit. + +## Maintain a TensorRT-LLM compatibility patch + +After editing an exact destination for a change that must remain downstream, +create its persistent patch: + +```bash +python scripts/vendor_sources.py patch VENDOR create --repo /path/to/upstream +``` + +After intentionally changing an already patched destination, regenerate the +patch: + +```bash +python scripts/vendor_sources.py patch VENDOR refresh --repo /path/to/upstream +``` + +Drop a no-longer-needed patch only after the destination exactly matches the +currently locked upstream selection: + +```bash +python scripts/vendor_sources.py patch VENDOR drop --repo /path/to/upstream +``` + +Generated patches live under `3rdparty/vendor_patches/`. Review them, but update +them only through the tool. Do not use a persistent patch for a change that +should be contributed upstream; use the export workflow instead. + +## Export a destination change upstream + +Start with the desired change in the TensorRT-LLM destination. The offline +check now fails by design. In a clean upstream checkout, create a temporary +branch at the currently locked commit **before** exporting: + +```bash +git -C /path/to/upstream switch -c trtllm-vendor-fix LOCKED_FULL_COMMIT +python scripts/vendor_sources.py export VENDOR --repo /path/to/upstream +``` + +The upstream checkout's selected source must be clean before export and its +`HEAD` must equal the locked commit. `export` computes the pending destination +delta relative to the locked materialization, applies only that delta to the +raw upstream source, and leaves the vendor lock, destination, and persistent +compatibility patch unchanged. + +Run the upstream tests, review the result, then commit and push the temporary +branch: + +```bash +git -C /path/to/upstream add path/in/upstream +git -C /path/to/upstream commit -s -m 'Apply exported fix' +git -C /path/to/upstream push -u origin trtllm-vendor-fix +``` + +Finally, pin the committed revision from that checkout: + +```bash +python scripts/vendor_sources.py pin VENDOR \ + --url https://example.com/my-fork/repository.git \ + --branch trtllm-vendor-fix \ + --commit NEW_FULL_COMMIT \ + --repo /path/to/upstream +``` + +`pin` first tries the selected files at `NEW_FULL_COMMIT` plus the existing +persistent compatibility patch. They must exactly equal the checked-in +destination. One exception is safe: if the raw new commit itself exactly equals +the destination, upstream has absorbed the compatibility patch, so `pin` drops +that patch and its metadata. Otherwise `pin` does not absorb a mismatch, +regenerate the patch, or copy candidate files into the destination. On success +it durably updates the immutable lock before removing an absorbed patch and +restores passing offline checks. If the patch cannot be removed after that +commit, `pin` succeeds with a warning and leaves a safe, unreferenced orphan; +delete the reported file manually. A failure before the durable lock commit +does not remove the existing patch. If directory synchronization fails after +the atomic replacement, the lock may already show the new pin, but the retained +patch keeps either recovered lock version reproducible. + +The same rule applies when adopting a newer commit that was developed upstream +first: prepare the destination to exactly match the proposed commit plus the +existing patch, then run `pin`. Do not use `sync` to look for that commit. + +## Remove a vendor + +Remove a lock entry and its generated compatibility patch while preserving the +destination: + +```bash +python scripts/vendor_sources.py remove VENDOR +``` + +The preserved destination is no longer protected by the lock. Delete or move +it separately as part of the reviewed migration that removes the vendor. + +## Source access and checks + +The default check is deliberately offline: + +```bash +python scripts/vendor_sources.py check +python scripts/vendor_sources.py check --offline +``` + +It validates the lock schema and path safety, patch metadata, and the checked-in +destination digest. It never invokes Git, performs DNS resolution, or contacts +a recorded URL. This is the always-run pre-commit check. A pending destination +delta therefore blocks a commit until it is synchronized, patched, or pinned. + +When network access is available, attempt verification against every recorded +upstream: + +```bash +python scripts/vendor_sources.py check --upstream +``` + +An inaccessible repository is reported as unavailable rather than failing. If +a commit can be obtained, a source, patch, or destination mismatch is an error. +Trusted maintainer CI can require access to every source: + +```bash +python scripts/vendor_sources.py check --upstream --require-access +``` + +To verify one vendor against an existing checkout without contacting the +recorded URL, provide it explicitly: + +```bash +python scripts/vendor_sources.py check VENDOR --repo /path/to/upstream +``` + +The checkout's configured remote may differ from the lock URL; it only needs to +contain the locked commit. Source-consuming commands accept the same `--repo` +form. + +An offline digest proves that the committed destination matches the lock. It +cannot independently prove that a URL, commit, and source directory produced +that destination. Creating and pinning vendors therefore require a fetched or +local repository, and URL or commit changes require vendor CODEOWNER review. +Never put credentials in a lock URL. Run checks that use internal credentials +only in a trusted environment, not with pull-request-controlled scripts. + +## License and attribution + +The vendor lock is a reproducibility record, not a license manifest. Before +adding a vendor, verify that the selected upstream files carry the required +notices and follow [the Python third-party process](py-thirdparty.md) or +[the C++ third-party process](cpp-thirdparty.md), as applicable. Exact upstream +files retain their upstream copyright headers. Add an NVIDIA header only to +files that TensorRT-LLM modifies. + +## PrimTS + +The `flashinfer-prims-ts` entry selects the complete Python tree under +`flashinfer/attention/prims_ts` and materializes it at +`tensorrt_llm/_torch/attention/backends/prims_ts`. The `**/*.py` selection +deliberately omits upstream README files. Its persistent patch contains only +TensorRT-LLM integration and compatibility adaptations; all other selected +files remain exact upstream copies. + +Use the normal commands with `flashinfer-prims-ts`, for example: + +```bash +python scripts/vendor_sources.py status flashinfer-prims-ts +python scripts/vendor_sources.py check flashinfer-prims-ts +``` diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch new file mode 100644 index 000000000000..44f3e7a4581d --- /dev/null +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -0,0 +1,164 @@ +diff --git a/block_sparse.py b/block_sparse.py +index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481e949f6b8 100644 +--- a/block_sparse.py ++++ b/block_sparse.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -35,12 +36,6 @@ from typing import Literal + import torch + + from flashinfer.api_logging import flashinfer_api +-from flashinfer.trace.templates.attention import ( +- prims_ts_block_sparse_trace, +- prims_ts_block_sparse_wrapper_trace_dispatch, +- prims_ts_paged_block_sparse_trace_dispatch, +- prims_ts_paged_block_sparse_wrapper_trace_dispatch, +-) + + from ._block_sparse.config import _validate_block_sparse_static_profile + from ._block_sparse.inspection import ( +@@ -220,7 +215,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): + # previously published revision intact and runnable. + self._plan_state = candidate + +- @flashinfer_api(trace=prims_ts_block_sparse_wrapper_trace_dispatch) ++ @flashinfer_api + def run( + self, + q: torch.Tensor, +@@ -302,7 +297,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): + return self._launch_validated_run(state, run_args, run_stream) + + +-@flashinfer_api(trace=prims_ts_block_sparse_trace) ++@flashinfer_api + def block_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, +@@ -512,7 +507,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): + ) + self._plan_state = candidate + +- @flashinfer_api(trace=prims_ts_paged_block_sparse_wrapper_trace_dispatch) ++ @flashinfer_api + def run( + self, + q: torch.Tensor, +@@ -618,7 +613,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): + return self._launch_validated_run(state, run_args, run_stream) + + +-@flashinfer_api(trace=prims_ts_paged_block_sparse_trace_dispatch) ++@flashinfer_api + def block_sparse_attention_with_paged_kv_cache( + q: torch.Tensor, + paged_kv_cache: PagedKVCache, +diff --git a/context.py b/context.py +index 7245bee1a0f725086171c9c5002115757e425d84..47996c867a962a3684c78e07d3a14eebf34b8452 100644 +--- a/context.py ++++ b/context.py +@@ -29,8 +29,7 @@ position is ``q + (S_kv - S_q)`` and ``window_left`` is measured from that + position. + + PrimTS context entry points are intentionally excluded from ``fi_trace`` for +-now; unlike the decode APIs, their ``@flashinfer_api`` decorators do not +-register trace templates. ++now; their ``@flashinfer_api`` decorators do not register trace templates. + """ + + from dataclasses import dataclass +@@ -424,7 +423,7 @@ def _validate_device(device: torch.device) -> int: + # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 + # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. + if capability == (10, 7): +- from ...cute_dsl.utils import require_cute_dsl_arch ++ from flashinfer.cute_dsl.utils import require_cute_dsl_arch + + require_cute_dsl_arch(device_index) + return device_index +diff --git a/decode.py b/decode.py +index 86797a2ad65b3cd433d829134e84fe25aedeaf88..3117a992956d62a13280063c268b937f1844c5f3 100644 +--- a/decode.py ++++ b/decode.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -25,11 +26,6 @@ from typing import TYPE_CHECKING, Literal, Optional, Union + import torch + + from flashinfer.api_logging import flashinfer_api +-from flashinfer.trace.templates.attention import ( +- attention_ts_decode_trace_dispatch, +- prims_ts_decode_trace_dispatch, +- prims_ts_decode_wrapper_trace_dispatch, +-) + + from ._tensor_aliasing import ( + _validate_out_does_not_overlap_inputs, +@@ -701,7 +697,7 @@ def _validate_runtime_device(device: torch.device) -> int: + # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 + # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. + if capability == (10, 7): +- from ...cute_dsl.utils import require_cute_dsl_arch ++ from flashinfer.cute_dsl.utils import require_cute_dsl_arch + + require_cute_dsl_arch(device_index) + return device_index +@@ -2154,7 +2150,7 @@ def _validate_decode_run_metadata_values( + ) + + +-@flashinfer_api(trace=prims_ts_decode_trace_dispatch) ++@flashinfer_api + def prims_ts_batch_decode_with_kv_cache( + query: torch.Tensor, + kv_cache: PagedKVCache, +@@ -2703,7 +2699,7 @@ class BatchDecodePagedTSWrapper: + # previous complete plan revision usable. + self._plan_state = candidate + +- @flashinfer_api(trace=prims_ts_decode_wrapper_trace_dispatch) ++ @flashinfer_api + def run( + self, + q: torch.Tensor, +@@ -2864,7 +2860,7 @@ class BatchDecodePagedTSWrapper: + ) + + +-@flashinfer_api(trace=attention_ts_decode_trace_dispatch) ++@flashinfer_api + def batch_decode_with_paged_kv_cache( + q: torch.Tensor, + paged_kv_cache: PagedKVCache, +diff --git a/mla_decode.py b/mla_decode.py +index c656f0023f2cb009e95f761299b0a97f098d18c8..9098a7c58a80fe1933e1239e920ae7894d0d9e6c 100644 +--- a/mla_decode.py ++++ b/mla_decode.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -25,7 +26,6 @@ from flashinfer.api_logging import flashinfer_api + from flashinfer.trace.templates.attention import ( + prims_ts_decode_mla_one_shot_trace_dispatch, + prims_ts_decode_mla_trace_dispatch, +- prims_ts_decode_mla_wrapper_trace_dispatch, + ) + + from ._tensor_aliasing import ( +@@ -1793,7 +1793,7 @@ class BatchMLADecodePagedTSWrapper: + split_kv=int(dict(policy)["split_kv"]), + ) + +- @flashinfer_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch) ++ @flashinfer_api + def run( + self, + query: torch.Tensor, diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml new file mode 100644 index 000000000000..f49dc129f048 --- /dev/null +++ b/3rdparty/vendor_sources.lock.yaml @@ -0,0 +1,13 @@ +schema_version: 1 +vendors: + flashinfer-prims-ts: + url: https://github.com/yuxianq/flashinfer.git + branch: trtllm-prims-ts + commit: e500966b575ab83db7c0e84e5a0f8fde6a4f3505 + source: flashinfer/attention/prims_ts + destination: tensorrt_llm/_torch/attention/backends/prims_ts + include: + - '**/*.py' + patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch + patch_digest: sha256:0e2f58c6633f57fee03df42049bc78d4d038063b810ad3a2f0ba6d62f8183887 + digest: sha256-tree-v1:e9af5482f6406af3128e711c3fb2d044359fb6e1d1bc9907743dc86d006d23ac diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index aecf0808c365..f657a16809c7 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -1,5 +1,5 @@ task-owned action + 0/1 two peer resource instances + + +----------+ + | GmemQKV | + | Q/K/V | + +-+--+-----+ + | | + +------------+ | + | | LoadTask: TMA Q and TMA K/V + v v + +-----+----+ +--------+ +-----------------+ + | SmemQ | | SmemKV | | TmemStatsDone0/1| + +-----+----+ +---+----+ +-----+-----------+ + | | | + +---------------+-----------------+ + | + +---------------- KV loop -----------------+ + | | + | MmaTask BMM1: SmemQ + SmemKV.K -> S, | + | waiting TmemStatsDone0/1 before S overwrites | + | aliased stats columns. | + v | + +-----+------+ | + | TmemSP0/1 | ---Softmax0/1Task: S -> P, stats--+--+ + | S then P | | + +-----+------+ v + | +--------------+ + | MmaTask BMM2: wait P from the | TmemStats0/1 | + | same TmemSP0/1 + SmemKV.V -> O +--------------+ + v | + +-----+------+ <--- CorrectionTask: stats + O ------+ + | TmemO | rescale O in-place ^ + +-----+------+ | + | | + | CorrectionTask release frees TmemO | + | before next BMM1 writes S in TmemSP. | + +------------------------------------------+ + | + | CorrectionTask tail + v + +-----+----+ + | SmemO0/1 | + +-----+----+ + | + | EpilogueTask: TMA store + v + +-----+----+ + | GmemO0/1 | + +----------+ + +In the illustrated D128 schedule, P readiness is the TmemSP0/1 pipeline state: +Softmax0/1Task stores P and releases the S/P slot, then MmaTask re-acquires the +same slot for BMM2. Staged D256 instead has a separate TmemPResource handoff so +next-tile QK can overlap previous-tile PV. CorrectionTask releases each TmemO +stage after rescaling it. + +TmemStats0/1 in the diagram are TmemStatsResource instances. They store the +correction statistics old_row_max, row_max, row_sum, and pad. + +The optional WorkQueue is scheduler state, not Q/K/V/O dataflow. Persistent +tasks consume it to select work tiles; static hardware scheduling omits it. + +Default D128 tasks and warp ownership (D256 uses the staged 12-warp layout +described by ``FmhaConfig``): + Softmax0Task warps 0-3 softmax on S0 -> P0, produce stats0 + Softmax1Task warps 4-7 softmax on S1 -> P1, produce stats1 + CorrectionTask warps 8-11 rescale O using correction stats + MmaTask warp 12 Q*K and P*V tcgen05 MMA + LoadTask warp 13 TMA loads Q/K/V + EpilogueTask warp 14 TMA stores O + PaddingTask warp 15 warp-group register participation + +Entry points: + - build_fmha_task_manager() -- constructs and validates the runtime task graph + - fmha_ts_kernel() -- @cute.kernel, the GPU kernel + - FmhaTs -- class with @cute.jit __call__ and kernel +""" + +import warnings +from dataclasses import dataclass +from collections.abc import Callable +from typing import Any, Tuple + +import cutlass +import cutlass.experimental.cuda as cuda +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +from cuda.bindings import driver as cuda_drv +from cutlass import Int32 +from ..tensor_map import create_tensor_map_ragged_from_tensor + +from cutlass.experimental.task_scheduling.memory import ( + SmemAllocation, + SmemAllocator, + TmemAllocator, +) +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + PipelineConfig, + TileSchedulerConfig, + WorkQueue, +) +from cutlass.experimental.task_scheduling.task import Task +from cutlass.experimental.task_scheduling.task_manager import TaskManager + +from .fmha_resources import ( + _SUPPORTED_CONTEXT_PAGE_SIZES, + FmhaConfig, + GmemOResource, + GmemQKVResource, + S0S1SequenceResource, + SmemKVResource, + SmemOResource, + SmemPageOffsetsKvResource, + SmemQResource, + TmemOResource, + TmemPResource, + TmemSPResource, + TmemStatsResource, + TmemStatsDoneResource, +) +from .fmha_tasks import ( + PackedContextWorkQueue, + create_correction_task, + create_epilogue_task, + create_load_task, + create_mma_task, + create_padding_task, + create_page_offsets_task, + create_scheduler_task, + create_softmax_task, +) + + +from .helpers import ( + bottom_right_window_max_tiles, + bottom_right_window_tile_start, + variable_window_cta_min_start, +) +from cutlass.experimental import primitives as prims + + +def _as_i32(x: int | Int32) -> Int32: + return Int32(x) if isinstance(x, int) else x + + +def _domain_min(a: int | Int32, b: int | Int32) -> int | Int32: + """``min`` that stays a Python ``int`` for static inputs and defers to + ``cute.math.min`` otherwise.""" + if isinstance(a, int) and isinstance(b, int): + return min(a, b) + return cute.math.min(_as_i32(a), _as_i32(b)) + + +def _domain_max(a: int | Int32, b: int | Int32) -> int | Int32: + """``max`` counterpart of :func:`_domain_min`.""" + if isinstance(a, int) and isinstance(b, int): + return max(a, b) + return cute.math.max(_as_i32(a), _as_i32(b)) + + +def _init_task_with_domain( + task: Task, + kwargs: dict[str, Any], + domain: int | None, + task_init: Callable[..., None] = Task.__init__, +) -> None: + """Initialize a captured Task with its static validation loop domain.""" + schedule = kwargs.get("schedule") + if schedule is None: + raise ValueError("Causal domain tasks require a captured schedule.") + if domain is not None: + schedule.loop_end = domain + task_init(task, **kwargs) + + +def _init_causal_domain_state( + task: Task, + *, + num_kv_tiles: int | Int32, + tile_size_q: int | Int32, + tile_size_kv: int | Int32, + q_offset: int | Int32, + seq_idx: int, + batch_idx: int | None, + cum_seqlen_q: cute.Tensor | None, + cum_seqlen_k: cute.Tensor | None, + seq_lens_kv: cute.Pointer | None, + runtime_kv_tile_multiple: int, + reverse_seq_tiles: int | Int32 | None, + offset: int, + window_size_left: int, + packed_window: bool, + kwargs: dict[str, Any], + task_init: Callable[..., None] = Task.__init__, +) -> None: + """Initialize causal-domain fields and the static validation domain.""" + static_domain = None + if isinstance(num_kv_tiles, int): + assert isinstance(tile_size_q, int) and isinstance(tile_size_kv, int) + if packed_window: + static_domain = min( + num_kv_tiles, + bottom_right_window_max_tiles( + q_tile_m=tile_size_q, + kv_tile_n=tile_size_kv, + window_size_left=window_size_left, + ), + ) + else: + seq_coord = 0 + max_q_row = q_offset + seq_coord * tile_size_q + tile_size_q - 1 + causal_n = max_q_row // tile_size_kv + 1 + static_domain = max( + tile_size_q // tile_size_kv, min(num_kv_tiles, causal_n) + ) + if window_size_left > 0 and not packed_window: + static_domain -= bottom_right_window_tile_start( + seq_coord=0, + q_tile_m=tile_size_q, + kv_tile_n=tile_size_kv, + q_offset=q_offset, + window_size_left=window_size_left, + ) + if offset: + static_domain -= offset + static_domain = max(static_domain, 0) + _init_task_with_domain(task, kwargs, static_domain, task_init=task_init) + task._num_kv_tiles = num_kv_tiles + task._tile_size_q = tile_size_q + task._tile_size_kv = tile_size_kv + task._q_offset = q_offset + # Index of the sequence coordinate in tile_coord. + task._seq_idx = seq_idx + # Mixed packed causal launches derive a request-local safe K-loop extent + # from live metadata. Fixed and uniform-packed launches leave these fields + # unset and retain the static plan-time domain calculation above. + task._batch_idx = batch_idx + task._cum_seqlen_q = cum_seqlen_q + task._cum_seqlen_k = cum_seqlen_k + task._seq_lens_kv = seq_lens_kv + task._runtime_kv_tile_multiple = runtime_kv_tile_multiple + task._reverse_seq_tiles = reverse_seq_tiles + # Offset adjusts the domain count. + task._offset = offset + task._window_size_left = window_size_left + task._packed_window = packed_window + + +class CausalDomainTask(Task): + """Task with causal and optional left-window masking domain. + + For causal FMHA, each Q tile only attends to K tiles where k <= q. + The domain (K-loop iteration count) varies per tile based on the + tile's sequence position. Head-paired sliding-window schedules reuse + the same domain calculation after subtracting skipped left-window KV + tiles. + """ + + def __init__( + self, + num_kv_tiles: int | Int32, + tile_size_q: int | Int32, + tile_size_kv: int | Int32, + q_offset: int | Int32 = 0, + seq_idx: int = 0, + batch_idx: int | None = None, + cum_seqlen_q: cute.Tensor | None = None, + cum_seqlen_k: cute.Tensor | None = None, + seq_lens_kv: cute.Pointer | None = None, + runtime_kv_tile_multiple: int = 1, + reverse_seq_tiles: int | Int32 | None = None, + offset: int = 0, + window_size_left: int = 0, + packed_window: bool = False, + **kwargs: Any, + ) -> None: + """Initialize causal-domain parameters. + + Args: + num_kv_tiles: Total number of K/V tiles in the sequence. + tile_size_q: Number of Q rows covered by one CTA tile. + tile_size_kv: Number of K/V rows covered by one K-loop iteration. + q_offset: Causal row-index shift for S_q < S_kv. + seq_idx: Index of the sequence coordinate in ``tile_coord``. + batch_idx: Index of the request coordinate in ``tile_coord`` for + packed causal launches. + cum_seqlen_q: Per-run packed-Q cumulative offsets. + cum_seqlen_k: Per-run packed-K/V cumulative offsets for contiguous K/V. + seq_lens_kv: Per-run K/V lengths for paged K/V. + runtime_kv_tile_multiple: Round a request-local K/V tile count up + to this multiple. Query-paired zero-offset causal scheduling + uses two to retain its synthetic peer-0 tail slot. + reverse_seq_tiles: Number of Q sequence work tiles when causal + balancing reverses launch order. ``None`` keeps natural order. + offset: Loop-count decrement after the causal/window tile count is + computed. Use 0 for N, 1 for N-1, and 2 for N-2 domains. + window_size_left: Left sliding-window size in tokens. Zero disables + left-window domain trimming. + packed_window: Use an offset-independent maximum window span for a + packed-ragged batch; runtime loads still use each request's offset. + **kwargs: Remaining ``Task`` constructor arguments, including an + optional captured ``schedule``. + """ + _init_causal_domain_state( + self, + num_kv_tiles=num_kv_tiles, + tile_size_q=tile_size_q, + tile_size_kv=tile_size_kv, + q_offset=q_offset, + seq_idx=seq_idx, + batch_idx=batch_idx, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=seq_lens_kv, + runtime_kv_tile_multiple=runtime_kv_tile_multiple, + reverse_seq_tiles=reverse_seq_tiles, + offset=offset, + window_size_left=window_size_left, + packed_window=packed_window, + kwargs=kwargs, + ) + + def get_domain(self, tile_coord: cute.Coord) -> int | Int32: + """Return the per-work-tile K-loop iteration count. + + Folds to a Python ``int`` when the tile coordinate and shape fields are + static, and defers to ``cute.math`` for runtime coordinates. + """ + seq_coord = tile_coord[self._seq_idx] + if cutlass.const_expr(self._reverse_seq_tiles is not None): + seq_coord = self._reverse_seq_tiles - seq_coord - 1 + q_offset = self._q_offset + num_kv_tiles = self._num_kv_tiles + runtime_q_tile_active = None + if cutlass.const_expr(self._cum_seqlen_q is not None): + assert self._batch_idx is not None + batch_coord = Int32(tile_coord[self._batch_idx]) + q_begin = Int32(self._cum_seqlen_q[batch_coord]) + seqlen_q = Int32(self._cum_seqlen_q[batch_coord + Int32(1)]) - q_begin + if cutlass.const_expr(self._seq_lens_kv is not None): + seqlen_k = Int32(self._seq_lens_kv[batch_coord]) + else: + assert self._cum_seqlen_k is not None + k_begin = Int32(self._cum_seqlen_k[batch_coord]) + seqlen_k = Int32(self._cum_seqlen_k[batch_coord + Int32(1)]) - k_begin + runtime_q_tile_active = Int32(seq_coord * self._tile_size_q < seqlen_q) + q_offset = seqlen_k - seqlen_q + num_kv_tiles = cute.ceil_div(seqlen_k, self._tile_size_kv) + if cutlass.const_expr(self._runtime_kv_tile_multiple > 1): + num_kv_tiles = ( + cute.ceil_div(num_kv_tiles, self._runtime_kv_tile_multiple) + * self._runtime_kv_tile_multiple + ) + if self._packed_window: + max_window_tiles = bottom_right_window_max_tiles( + q_tile_m=self._tile_size_q, + kv_tile_n=self._tile_size_kv, + window_size_left=self._window_size_left, + ) + result = _domain_min(num_kv_tiles, max_window_tiles) + else: + max_q_row = q_offset + seq_coord * self._tile_size_q + self._tile_size_q - 1 + causal_n = max_q_row // self._tile_size_kv + 1 + # Softmax assumes there is at least one Q-tile-width of K work. + result = _domain_max( + self._tile_size_q // self._tile_size_kv, + _domain_min(num_kv_tiles, causal_n), + ) + if self._window_size_left > 0 and not self._packed_window: + result -= bottom_right_window_tile_start( + seq_coord=seq_coord, + q_tile_m=self._tile_size_q, + kv_tile_n=self._tile_size_kv, + q_offset=q_offset, + window_size_left=self._window_size_left, + ) + if self._offset: + result = result - self._offset + if cutlass.const_expr(runtime_q_tile_active is not None): + # The outer grid is sized by the planned maximum Q length, so a + # mixed packed request can have trailing work-tile slots with no Q + # rows. Retain the smallest legal N/N-1/N-2 pipeline domains for + # those slots instead of traversing K/V according to a large + # bottom-right offset. Resource-level ragged extents suppress the + # dummy Q/O traffic; the minimum domains preserve task handoffs. + minimum_domain = max( + self._tile_size_q // self._tile_size_kv - self._offset, 0 + ) + result = ( + runtime_q_tile_active * result + + (Int32(1) - runtime_q_tile_active) * minimum_domain + ) + return result + + +class CausalSoftmaxDomainTask(CausalDomainTask): + """Softmax task with causal domain selection.""" + + def __init__( + self, + num_kv_tiles: int | Int32, + tile_size_q: int | Int32, + tile_size_kv: int | Int32, + q_offset: int | Int32 = 0, + seq_idx: int = 0, + batch_idx: int | None = None, + cum_seqlen_q: cute.Tensor | None = None, + cum_seqlen_k: cute.Tensor | None = None, + seq_lens_kv: cute.Pointer | None = None, + runtime_kv_tile_multiple: int = 1, + reverse_seq_tiles: int | Int32 | None = None, + offset: int = 0, + window_size_left: int = 0, + packed_window: bool = False, + **kwargs: Any, + ) -> None: + """Initialize softmax causal-domain parameters. + + Parameters + ---------- + num_kv_tiles : int or Int32 + Total number of K/V tiles in the sequence. + tile_size_q : int or Int32 + Number of Q rows covered by one CTA tile. + tile_size_kv : int or Int32 + Number of K/V rows covered by one K-loop iteration. + q_offset : int or Int32 + Causal row-index shift for S_q < S_kv. + seq_idx : int + Index of the sequence coordinate in ``tile_coord``. + batch_idx : int or None + Index of the request coordinate for packed causal launches. + cum_seqlen_q, cum_seqlen_k : cute.Tensor or None + Per-run cumulative offsets used to derive request-local causal + shifts for contiguous packed input. + seq_lens_kv : cute.Pointer or None + Per-run K/V lengths used for paged input. + runtime_kv_tile_multiple : int + Request-local K/V tile-count alignment for paired-tail scheduling. + reverse_seq_tiles : int or Int32 or None + Number of Q sequence work tiles for reversed causal-balanced order. + offset : int + Loop-count decrement after the causal/window tile count is + computed. + window_size_left : int + Left sliding-window size in tokens. Zero disables left-window + domain trimming. + packed_window : bool + Whether to use an offset-independent packed-ragged window span. + **kwargs : Any + Remaining ``Task`` constructor arguments. + """ + _init_causal_domain_state( + self, + num_kv_tiles=num_kv_tiles, + tile_size_q=tile_size_q, + tile_size_kv=tile_size_kv, + q_offset=q_offset, + seq_idx=seq_idx, + batch_idx=batch_idx, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=seq_lens_kv, + runtime_kv_tile_multiple=runtime_kv_tile_multiple, + reverse_seq_tiles=reverse_seq_tiles, + offset=offset, + window_size_left=window_size_left, + packed_window=packed_window, + kwargs=kwargs, + ) + + +class VariableWindowDomainTask(Task): + """Task whose K-loop domain comes from packed-Q window endpoints.""" + + def __init__( + self, + variable_window_token_starts: cute.Tensor, + variable_window_token_ends: cute.Tensor, + variable_window_cta_starts: cute.Tensor, + num_kv_tiles: int | Int32, + q_stride: int | Int32, + tile_size_q: int, + tile_size_kv: int, + seq_idx: int, + batch_idx: int, + offset: int = 0, + **kwargs: Any, + ) -> None: + """Initialize a variable-window task domain. + + Args: + variable_window_token_starts: Flattened inclusive K-token start for + every Q row, laid out with ``q_stride`` rows per batch. + variable_window_token_ends: Flattened inclusive K-token end for + every Q row, using the same layout as the start tensor. + variable_window_cta_starts: Precomputed minimum K-token start for + each Q CTA, used as the common K/V load origin. + num_kv_tiles: Planned maximum number of K/V tiles. + q_stride: Number of Q rows per batch in the flattened bound tensors. + tile_size_q: Number of Q rows covered by one CTA. + tile_size_kv: Number of K/V tokens covered by one loop tile. + seq_idx: Index of the Q-sequence coordinate in ``tile_coord``. + batch_idx: Index of the batch coordinate in ``tile_coord``. + offset: Domain-count decrement. Zero selects N; one selects N-1. + **kwargs: Remaining ``Task`` arguments, including the required + captured ``schedule``. + """ + if kwargs.get("schedule") is None: + raise ValueError("VariableWindow domain tasks require a captured schedule") + super().__init__(**kwargs) + self._variable_window_token_starts = variable_window_token_starts + self._variable_window_token_ends = variable_window_token_ends + self._variable_window_cta_starts = variable_window_cta_starts + self._num_kv_tiles = num_kv_tiles + self._q_stride = q_stride + self._tile_size_q = tile_size_q + self._tile_size_kv = tile_size_kv + self._seq_idx = seq_idx + self._batch_idx = batch_idx + self._offset = offset + + def get_domain(self, tile_coord: cute.Coord) -> Int32: + """Return the number of K tiles intersecting this Q CTA's bounds.""" + seq_coord = Int32(tile_coord[self._seq_idx]) + batch_coord = Int32(tile_coord[self._batch_idx]) + first_local_q = seq_coord * self._tile_size_q + last_local_q = cute.math.min( + first_local_q + self._tile_size_q - Int32(1), + self._q_stride - Int32(1), + ) + packed_q_base = batch_coord * self._q_stride + first_k = variable_window_cta_min_start( + self._variable_window_cta_starts, + batch_coord=batch_coord, + seq_coord=seq_coord, + q_stride=self._q_stride, + tile_size_q=self._tile_size_q, + ) + last_k = Int32(self._variable_window_token_ends[packed_q_base + last_local_q]) + first_k_tile = first_k // self._tile_size_kv + last_k_tile = (last_k + self._tile_size_kv) // self._tile_size_kv + return last_k_tile - first_k_tile - self._offset + + +DomainPolicyValue = int | bool | Int32 | type[Task] | cute.Tensor | cute.Pointer +DomainKwargs = dict[str, DomainPolicyValue] + + +def resolve_head_paired_mode( + *, + head_paired: bool, + is_causal: bool, + window_size_left: int, +) -> bool: + """Return the effective head-paired mode for a mask/window configuration.""" + if window_size_left < 0: + raise ValueError("window_size_left must be non-negative") + if window_size_left > 0 and not is_causal: + raise ValueError("window_size_left requires is_causal=True") + if window_size_left > 0 and not head_paired: + warnings.warn( + f"window_size_left={window_size_left} requires head_paired=True; " + "got head_paired=False, enabling head_paired", + UserWarning, + stacklevel=2, + ) + return head_paired or window_size_left > 0 + + +def validate_head_paired_head_ratio(*, head_paired: bool, h_r: int) -> None: + """Validate the Q/KV head ratio required by head-paired scheduling.""" + if not head_paired: + return + if h_r <= 1: + raise ValueError( + f"head_paired requires grouped-query attention with h_q > h_kv, got {h_r=}" + ) + if h_r % 2 != 0: + raise ValueError( + f"head_paired requires an even Q/KV head ratio, got h_q / h_kv = {h_r}" + ) + + +@dataclass(frozen=True) +class FmhaDomainPolicy: + """Domain and mask policy selected for one FMHA launch flavor.""" + + domain_n_kwargs: DomainKwargs + domain_n_minus_1_kwargs: DomainKwargs + softmax0_domain_kwargs: DomainKwargs + softmax1_domain_kwargs: DomainKwargs + + +def build_context_task_manager( + *, + cfg: FmhaConfig, + tile_sched_params: ( + utils.PersistentTileSchedulerParams + | utils.ClcDynamicPersistentTileSchedulerParams + | None + ), + tma_q_desc: cutlass.Pointer | None, + tma_k_desc: cutlass.Pointer | None, + tma_v_desc: cutlass.Pointer | None, + tma_o_desc: cutlass.Pointer | None, + cum_seqlen_q: cute.Tensor | None, + cum_seqlen_k: cute.Tensor | None, + variable_window_token_starts: cute.Tensor | None = None, + variable_window_token_ends: cute.Tensor | None = None, + variable_window_cta_starts: cute.Tensor | None = None, + variable_window_q_stride: int | Int32 = 0, + scale_softmax_log2: cute.Tensor | None = None, + output_scale: cute.Tensor | None = None, + g_block_tables: cute.Pointer | None = None, + block_table_row_stride: int | Int32 = 0, + g_seq_lens_kv: cute.Pointer | None = None, + max_seq_len_kv: int | Int32 | None = None, + num_kv_tiles: int | Int32, + q_offset: int | Int32, + domain_n_kwargs: DomainKwargs, + domain_n_minus_1_kwargs: DomainKwargs, + softmax0_domain_kwargs: DomainKwargs, + softmax1_domain_kwargs: DomainKwargs, + tmem_o_extra_kwargs: DomainKwargs | None = None, + is_persistent: bool = True, + is_clc_dynamic: bool = False, + clc_response_ptr: cute.Pointer | None = None, + exhaustive_deadlock_race_check: bool = True, +) -> Tuple[ + TaskManager, + list[MemoryResource], + SmemAllocation, + SmemAllocation, + WorkQueue | None, + SmemAllocation | None, +]: + """Build the FMHA TaskManager with all resources, tasks, and dependency graph. + + ``scale_softmax_log2`` and ``output_scale`` are optional one-element + float32 device tensors. Resource setup loads element 0 once before the K/V + loop and reuses the value through task-local dataflow. ``scale_softmax_log2`` + is the base-2 softmax multiplier, normally + ``softmax_scale * log2(e)``; FP8 callers can fold Q/K dequant scales into + it. ``output_scale`` multiplies the final O store; FP8 output can fold the + V dequant scale and output quant scale into it. + + SMEM buffers are declared via ``SmemAllocation`` and bound from + ``ResourceContext`` by auxiliary resource init work. The ``SmemAllocator`` + is passed to ``TaskManager`` for unified allocation. Infrastructure slots + (``tmem_ptr_i32``, ``tmem_dealloc_mbar``) are included in the unified SMEM + block and returned as allocation descriptors. + + Parameters + ---------- + cfg : FmhaConfig + Kernel-wide configuration. + tile_sched_params : PersistentTileSchedulerParams or ClcDynamicPersistentTileSchedulerParams, optional + Persistent tile scheduler params. ``None`` for static or validation-only builds. + tma_q_desc, tma_k_desc, tma_v_desc, tma_o_desc : cutlass.Pointer, optional + TMA descriptor pointers for Q, K, V, O. + cum_seqlen_q, cum_seqlen_k : cute.Tensor, optional + Cumulative sequence-length metadata for varlen launches. + num_kv_tiles : int or Int32 + Number of KV tiles in the loop domain. + q_offset : int or Int32 + Default right shift of Q rows for causal S_q < S_kv masking. Mixed + packed launches derive the request-local shift from live metadata. + domain_n_kwargs : dict + Domain policy for tasks that process the full KV loop. + domain_n_minus_1_kwargs : dict + Domain policy for tasks whose HEAD handles the first KV tile and whose + LOOP handles the remaining tiles. + softmax0_domain_kwargs, softmax1_domain_kwargs : dict + Domain policies for the two softmax peers. + tmem_o_extra_kwargs : dict, optional + Extra domain/resource policy knobs for TMEM O. + is_persistent : bool + Whether tasks use WorkQueue-backed persistent scheduling. + is_clc_dynamic : bool + Whether to use CLC dynamic persistent scheduling. + clc_response_ptr : cute.Pointer, optional + SMEM pointer for the CLC response buffer, required when + ``is_clc_dynamic`` is true. + exhaustive_deadlock_race_check : bool + Whether TaskManager runs the exhaustive schedule checker during + construction. + Returns + ------- + tuple + ``(TaskManager, tmem_resources, tmem_ptr_alloc, dealloc_mbar_alloc, + work_queue, clc_response_alloc)``. + """ + if cfg.use_paged_kv: + if cfg.num_tokens_per_page not in _SUPPORTED_CONTEXT_PAGE_SIZES: + raise ValueError( + "paged context requires num_tokens_per_page in " + f"{_SUPPORTED_CONTEXT_PAGE_SIZES}; got " + f"{cfg.num_tokens_per_page}" + ) + if cfg.causal_single_kv_tile and (cfg.use_paged_kv or cfg.has_varlen): + raise ValueError("causal_single_kv_tile requires fixed contiguous K/V storage") + # --------------------------------------------------------------------------- + # Cluster / CTA layout + # --------------------------------------------------------------------------- + # CTA layout in CuTe VMNK order. V is the CTA-group dimension used for + # 2-CTA cooperative MMA; MNK are the logical cluster axes. FMHA uses a + # single-CTA cluster here, so (V, M, N, K) = (1, 1, 1, 1). + cluster_shape_vmnk = (1, 1, 1, 1) + + # --------------------------------------------------------------------------- + # Warp counts + # --------------------------------------------------------------------------- + # Warp ownership: 4 softmax, 4 correction, 1 MMA, 1 load, 1 epilogue, and + # one padding/scheduler warp-group participant. + num_mma_warps = 1 + num_softmax_warps = 4 + num_correction_warps = 4 + num_epilogue_warps = 1 + + # --------------------------------------------------------------------------- + # Cooperative groups + # --------------------------------------------------------------------------- + warp_size = cute.arch.WARP_SIZE + Agent = pipeline.Agent + + # For TMA pipelines: elect_one arrival, producer_group = 1 thread. + tma_producer_group = pipeline.CooperativeGroup(Agent.Thread) + + # --------------------------------------------------------------------------- + # Pipeline configs + # --------------------------------------------------------------------------- + + # SmemQ: TmaUmma, topology-derived stages, Load -> MMA. The paired D128 + # schedule uses two Q instances; staged D256 uses one. + # advance_on_wait=True advances consumer_state on ConsumerWait. + smem_q_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.q_stage, + num_bytes=cfg.tma_copy_q_bytes, + producer_group=tma_producer_group, + consumer_group=pipeline.CooperativeGroup(Agent.Thread), + cta_layout_vmnk=cluster_shape_vmnk, + advance_on_wait=True, + ) + # SmemKV: capacity-derived stages, Load -> MMA. + # advance_on_wait=True advances consumer_state at ConsumerWait rather than + # ConsumerRelease. This keeps the previous V tile live while MMA starts the + # next QK tile, giving QK0 -> PV1(previous V) -> QK1 ordering without + # releasing the previous V tile first. + smem_kv_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.kv_stage, + num_bytes=cfg.tma_copy_kv_bytes, + producer_group=tma_producer_group, + consumer_group=pipeline.CooperativeGroup(Agent.Thread), + cta_layout_vmnk=cluster_shape_vmnk, + advance_on_wait=True, + ) + # Page-offsets prefetch (paged-KV only): the auxiliary warp produces and + # the load warp consumes. + # Async-async because both are CUDA-thread groups (no TMA arrival barrier). + smem_page_offsets_pipeline_cfg = None + smem_page_offsets_v_pipeline_cfg = None + split_page_offset_pipelines = False + if cfg.stages_page_offsets_in_smem: + # A staged single-instance schedule keeps independent K and V page-ID + # rings because its K-ahead/V-delayed lifetimes overlap. Split the + # ordinary page-offset stage budget across those two rings, just as + # FMHA decode does for its split-head-dimension page-offset flow. V + # needs one fewer credit because the boundary tile's IDs are retained + # in registers before its SMEM window is released. The shared K/V + # ring used by every other topology retains the full depth. + page_offset_stage_counts = cfg.page_offset_pipeline_stage_counts + page_offsets_k_stages = page_offset_stage_counts[0] + page_offsets_v_stages = None + if len(page_offset_stage_counts) == 2: + split_page_offset_pipelines = True + page_offsets_v_stages = page_offset_stage_counts[1] + + # The load warp (1 warp = 32 threads) is the only consumer; signal-All + # requires the consumer group to size match num_warps × warp_size. + def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: + return PipelineConfig.create_async_async_pipeline_cfg( + num_stages=num_stages, + producer_group=pipeline.CooperativeGroup( + Agent.Thread, + cfg.page_offsets_num_warps * warp_size, + ), + consumer_group=pipeline.CooperativeGroup( + Agent.Thread, + warp_size, + ), + cta_layout_vmnk=cluster_shape_vmnk, + producer_op=pipeline.PipelineOp.AsyncLoad, + ) + + smem_page_offsets_pipeline_cfg = make_page_offsets_pipeline_cfg( + page_offsets_k_stages + ) + if page_offsets_v_stages is not None: + smem_page_offsets_v_pipeline_cfg = make_page_offsets_pipeline_cfg( + page_offsets_v_stages + ) + + softmax_group = pipeline.CooperativeGroup( + Agent.Thread, + num_softmax_warps * warp_size, + ) + correction_group = pipeline.CooperativeGroup( + Agent.Thread, + num_correction_warps * warp_size, + ) + epilogue_group = pipeline.CooperativeGroup( + Agent.Thread, + num_epilogue_warps * warp_size, + ) + umma_hw_group = pipeline.CooperativeGroup(Agent.Thread) + mma_group = pipeline.CooperativeGroup(Agent.Thread, num_mma_warps * warp_size) + + tmem_sp0_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=cfg.mma_softmax_stage, + producer_group=umma_hw_group, + consumer_group=softmax_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + tmem_sp1_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=cfg.mma_softmax_stage, + producer_group=umma_hw_group, + consumer_group=softmax_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + tmem_p0_pipeline_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=cfg.mma_softmax_stage, + producer_group=softmax_group, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + tmem_o_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=cfg.mma_corr_stage, + producer_group=umma_hw_group, + consumer_group=correction_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + + tmem_vec0_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=cfg.softmax_corr_stage, + producer_group=softmax_group, + consumer_group=correction_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + tmem_vec1_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=cfg.softmax_corr_stage, + producer_group=softmax_group, + consumer_group=correction_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + smem_o_0_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=correction_group, + consumer_group=( + correction_group if cfg.fuse_epilogue_into_correction else epilogue_group + ), + cta_layout_vmnk=cluster_shape_vmnk, + ) + smem_o_1_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=correction_group, + consumer_group=epilogue_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + # S0-S1 sequence barrier: PipelineAsync, 1 stage. + # Ensures Softmax0 finishes P store before Softmax1 starts P compute. + # Both groups have 4 warps (128 threads), but they are different warps. + s0s1_seq_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=softmax_group, + consumer_group=softmax_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + # TmemStatsDone barriers: MMA acquires before QK->S that aliases stats TMEM + # columns, and Correction releases after reading stats from TMEM. The + # pipeline starts empty, so MMA's first ProducerAcquire succeeds without + # priming. Separate barriers let MMA start QK->S0 as soon as TmemStats0 is + # read, without waiting for TmemStats1. + tmem_stats_done_0_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=mma_group, + consumer_group=correction_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + tmem_stats_done_1_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=mma_group, + consumer_group=correction_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + + # --------------------------------------------------------------------------- + # Create resource instances + # --------------------------------------------------------------------------- + gmem_qkv = GmemQKVResource( + tma_q_desc=tma_q_desc, + tma_k_desc=tma_k_desc, + tma_v_desc=tma_v_desc, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + q_offset=q_offset, + cfg=cfg, + seqlens_kv=g_seq_lens_kv, + block_table_row_stride=block_table_row_stride, + max_seq_len_kv=max_seq_len_kv, + variable_window_token_starts=variable_window_token_starts, + variable_window_cta_starts=variable_window_cta_starts, + variable_window_q_stride=variable_window_q_stride, + name="gmem_qkv", + ) + smem_q = SmemQResource( + tma_q_desc=tma_q_desc, + pipeline_config=smem_q_pipeline_cfg, + cfg=cfg, + name="smem_q", + ) + smem_page_offsets_kv: SmemPageOffsetsKvResource | None = None + smem_page_offsets_v: SmemPageOffsetsKvResource | None = None + if cfg.stages_page_offsets_in_smem: + smem_page_offsets_kv = SmemPageOffsetsKvResource( + block_tables=g_block_tables, + pipeline_config=smem_page_offsets_pipeline_cfg, + cfg=cfg, + name="smem_page_offsets_kv", + ) + if split_page_offset_pipelines: + # The D256 schedule keeps independent K/V page-window + # stages because its K-ahead/V-delayed pipeline crosses window + # boundaries. D128 consumes K and V together and FlashInfer's + # public paged API supplies one shared page-ID row, so it retains + # one stage for both sides. + smem_page_offsets_v = SmemPageOffsetsKvResource( + block_tables=g_block_tables, + pipeline_config=smem_page_offsets_v_pipeline_cfg, + cfg=cfg, + page_table_is_v=True, + name="smem_page_offsets_v", + ) + smem_kv = SmemKVResource( + tma_k_desc=tma_k_desc, + tma_v_desc=tma_v_desc, + pipeline_config=smem_kv_pipeline_cfg, + cfg=cfg, + page_offsets_kv=smem_page_offsets_kv, + page_offsets_v=smem_page_offsets_v, + block_tables=g_block_tables, + name="smem_kv", + ) + + # WorkQueue: persistent tile scheduler state (static or CLC dynamic), not + # Q/K/V/O dataflow. Non-persistent launches omit it so each CTA executes + # one tile directly from its hardware tile coordinate. + work_queue: WorkQueue | None = None + if not is_persistent and tile_sched_params is not None: + raise ValueError( + "non-persistent FMHA task managers must not receive tile_sched_params" + ) + if is_clc_dynamic: + if not is_persistent: + raise ValueError("CLC dynamic scheduling requires persistent mode") + # CLC dynamic: pipeline + CLC tile scheduler config. FMHA uses + # single-CTA clusters. + cluster_size = 1 + num_consumer_threads = cfg.block_warps * warp_size * cluster_size + work_queue_pipeline_cfg = PipelineConfig.create_clc_fetch_async_pipeline_cfg( + num_stages=1, + num_bytes=16, + producer_group=pipeline.CooperativeGroup(Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + Agent.Thread, + num_consumer_threads, + ), + cta_layout_vmnk=cluster_shape_vmnk, + ) + tile_scheduler_config = ( + TileSchedulerConfig.create_clc_dynamic_persistent_tile_scheduler_params( + tile_scheduler_params=tile_sched_params, + response_ptr=clc_response_ptr, + ) + ) + work_queue_kwargs = { + "tile_scheduler_config": tile_scheduler_config, + "pipeline_config": work_queue_pipeline_cfg, + "name": "work_queue", + } + if ( + cfg.is_causal + and cfg.has_varlen + and not cfg.has_uniform_varlen + and cum_seqlen_q is not None + ): + work_queue = PackedContextWorkQueue( + cfg=cfg, + cum_seqlen_q=cum_seqlen_q, + **work_queue_kwargs, + ) + else: + work_queue = WorkQueue(**work_queue_kwargs) + elif is_persistent: + tile_scheduler_config = ( + TileSchedulerConfig.create_static_persistent_tile_scheduler_params( + tile_scheduler_params=tile_sched_params, + ) + ) + work_queue_kwargs = { + "tile_scheduler_config": tile_scheduler_config, + "name": "work_queue", + } + if ( + cfg.is_causal + and cfg.has_varlen + and not cfg.has_uniform_varlen + and cum_seqlen_q is not None + ): + work_queue = PackedContextWorkQueue( + cfg=cfg, + cum_seqlen_q=cum_seqlen_q, + **work_queue_kwargs, + ) + else: + work_queue = WorkQueue(**work_queue_kwargs) + + tmem_sp0 = TmemSPResource( + pipeline_config=tmem_sp0_pipeline_cfg, + cfg=cfg, + tmem_s_offset=cfg.tmem_s0_offset, + tmem_p_offset=cfg.tmem_p0_offset, + q_half=0, + q_offset=q_offset, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=g_seq_lens_kv, + variable_window_token_starts=variable_window_token_starts, + variable_window_token_ends=variable_window_token_ends, + variable_window_cta_starts=variable_window_cta_starts, + variable_window_q_stride=variable_window_q_stride, + scale_softmax_log2=scale_softmax_log2, + name="tmem_sp0", + ) + tmem_p0: TmemPResource | None = None + if cfg.has_tmem_p_pipeline: + tmem_p0 = TmemPResource( + pipeline_config=tmem_p0_pipeline_cfg, + cfg=cfg, + tmem_p_offset=cfg.tmem_p0_offset, + name="tmem_p0", + ) + tmem_vec0 = TmemStatsResource( + pipeline_config=tmem_vec0_pipeline_cfg, + cfg=cfg, + tmem_vec_offset=cfg.tmem_vec0_offset, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + name="tmem_vec0", + ) + smem_o_0 = SmemOResource( + pipeline_config=smem_o_0_pipeline_cfg, + cfg=cfg, + stage_idx=0, + tmem_vec_resource=tmem_vec0, + name="smem_o_0", + ) + gmem_o_0 = GmemOResource( + tma_o_desc=tma_o_desc, + cum_seqlen_q=cum_seqlen_q, + cfg=cfg, + stage_idx=0, + name="gmem_o_0", + ) + # TmemStatsDone barrier shared between MMA + # (producer) and Correction (consumer). Prevents cross-tile aliasing races. + tmem_stats_done_0 = TmemStatsDoneResource( + pipeline_config=tmem_stats_done_0_pipeline_cfg, + name="tmem_stats_done_0", + ) + + single_qkv_instance = cfg.single_qkv_instance + tmem_sp1: TmemSPResource | None = None + tmem_vec1: TmemStatsResource | None = None + smem_o_1: SmemOResource | None = None + gmem_o_1: GmemOResource | None = None + s0s1_seq: S0S1SequenceResource | None = None + tmem_stats_done_1: TmemStatsDoneResource | None = None + + if not single_qkv_instance: + tmem_sp1 = TmemSPResource( + pipeline_config=tmem_sp1_pipeline_cfg, + cfg=cfg, + tmem_s_offset=cfg.tmem_s1_offset, + tmem_p_offset=cfg.tmem_p1_offset, + q_half=1, + q_offset=q_offset, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=g_seq_lens_kv, + variable_window_token_starts=variable_window_token_starts, + variable_window_token_ends=variable_window_token_ends, + variable_window_cta_starts=variable_window_cta_starts, + variable_window_q_stride=variable_window_q_stride, + scale_softmax_log2=scale_softmax_log2, + name="tmem_sp1", + ) + tmem_vec1 = TmemStatsResource( + pipeline_config=tmem_vec1_pipeline_cfg, + cfg=cfg, + tmem_vec_offset=cfg.tmem_vec1_offset, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + name="tmem_vec1", + ) + smem_o_1 = SmemOResource( + pipeline_config=smem_o_1_pipeline_cfg, + cfg=cfg, + stage_idx=1, + tmem_vec_resource=tmem_vec1, + name="smem_o_1", + ) + gmem_o_1 = GmemOResource( + tma_o_desc=tma_o_desc, + cum_seqlen_q=cum_seqlen_q, + cfg=cfg, + stage_idx=1, + name="gmem_o_1", + ) + # S0-S1 sequence barrier: one sequencing resource shared between Softmax0 + # (producer) and Softmax1 (consumer), like other shared resources. + s0s1_seq = S0S1SequenceResource( + pipeline_config=s0s1_seq_pipeline_cfg, + name="s0s1_seq", + ) + tmem_stats_done_1 = TmemStatsDoneResource( + pipeline_config=tmem_stats_done_1_pipeline_cfg, + name="tmem_stats_done_1", + ) + + tmem_o_kwargs = tmem_o_extra_kwargs or {} + tmem_o = TmemOResource( + pipeline_config=tmem_o_pipeline_cfg, + cfg=cfg, + tmem_o0_offset=cfg.tmem_o0_offset, + tmem_o1_offset=cfg.tmem_o1_offset, + tmem_vec0_resource=tmem_vec0, + tmem_vec1_resource=tmem_vec1, + name="tmem_o", + **tmem_o_kwargs, + ) + + # --------------------------------------------------------------------------- + # Create tasks & dependency graph + # --------------------------------------------------------------------------- + # Per-task domains match the handwritten FMHA schedule's iteration counts: + # - Softmax0/1: N iterations (one per KV tile) + # - MMA/Load/Correction: N-1 iterations + # (HEAD handles first tile, LOOP handles remaining N-1) + # - Epilogue/Padding: TAIL-only (domain does not matter, use N) + # + # With causal masking, the domain depends on the Q tile position: only + # process K tiles where at least one Q row can attend (k <= q). The + # CausalDomainTask subclass overrides get_domain(); seq_idx follows the + # tile-coordinate order selected by the host launch. + def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: + """Append the scheduler resource only for WorkQueue-backed launches.""" + deps = list(resources) + if work_queue is not None: + deps.append(work_queue) + return deps + + mma_domain_kwargs = domain_n_minus_1_kwargs + if single_qkv_instance and not cfg.has_tmem_p_pipeline: + # The non-split single-instance schedule performs every QK/PV pair + # inside its loop; it has no separate HEAD QK or TAIL PV iteration. + mma_domain_kwargs = domain_n_kwargs + load_task = create_load_task( + gmem_qkv, + smem_q, + smem_kv, + work_queue, + smem_page_offsets_kv=smem_page_offsets_kv, + smem_page_offsets_v=smem_page_offsets_v, + **domain_n_kwargs, + ) + mma_task = create_mma_task( + gmem_qkv, + smem_q, + smem_kv, + tmem_sp0, + tmem_sp1, + tmem_p0, + tmem_o, + tmem_stats_done_0, + tmem_stats_done_1, + work_queue, + **mma_domain_kwargs, + ) + + # Causal query-paired moves masked/invalid K iterations from LOOP to TAIL + # with no runtime branch. SP resources derive mask selection from cfg. + softmax0_task = create_softmax_task( + 0, + tmem_sp0, + tmem_vec0, + tmem_p0, + s0s1_seq, + work_queue, + **softmax0_domain_kwargs, + ) + softmax1_task: Task | None = None + if not single_qkv_instance: + if tmem_sp1 is None or tmem_vec1 is None: + raise ValueError("paired softmax scheduling requires peer-1 resources") + softmax1_task = create_softmax_task( + 1, + tmem_sp1, + tmem_vec1, + None, + s0s1_seq, + work_queue, + **softmax1_domain_kwargs, + ) + correction_task = create_correction_task( + tmem_vec0, + tmem_vec1, + tmem_o, + smem_o_0, + smem_o_1, + gmem_o_0, + gmem_o_1, + tmem_stats_done_0, + tmem_stats_done_1, + work_queue, + **domain_n_minus_1_kwargs, + ) + epilogue_task: Task | None = None + freed_epilogue_task: Task | None = None + if cfg.fuse_epilogue_into_correction: + if not is_clc_dynamic: + # Warp 10 remains part of the producer/correction warpgroup and + # must participate in setmaxnreg before it becomes a scheduler. + freed_epilogue_task = create_padding_task( + work_queue, + warp_idx=cfg.epilogue_warp_id, + num_registers=cfg.num_regs_other, + name="EpiloguePaddingTask", + **domain_n_kwargs, + ) + else: + epilogue_task = create_epilogue_task( + smem_o_0, + smem_o_1, + gmem_o_0, + gmem_o_1, + work_queue, + **domain_n_kwargs, + ) + scheduler_task: Task | None = None + if is_clc_dynamic: + scheduler_task = create_scheduler_task( + work_queue, + warp_idx=( + cfg.epilogue_warp_id + if cfg.fuse_epilogue_into_correction + else cfg.empty_warp_id + ), + num_registers=cfg.num_regs_other, + domain=0, + ) + if smem_page_offsets_kv is not None: + # Paged-KV: the auxiliary warp prefetches page-table entries instead + # of padding and still participates in setmaxnreg.sync. + auxiliary_task = create_page_offsets_task( + gmem_qkv, + smem_page_offsets_kv, + work_queue, + num_registers=cfg.num_regs_other, + smem_page_offsets_v=smem_page_offsets_v, + **domain_n_kwargs, + ) + elif scheduler_task is not None and not cfg.fuse_epilogue_into_correction: + # D128 retains the original topology where the scheduler itself is the + # final warp-group participant. + auxiliary_task = scheduler_task + scheduler_task = None + else: + # setmaxnreg.sync requires every warp in the final warp group to + # participate, so the empty warp needs a padding/scheduler task. + auxiliary_task = create_padding_task( + work_queue, + warp_idx=cfg.empty_warp_id, + num_registers=cfg.num_regs_other, + **domain_n_kwargs, + ) + + task_list = [softmax0_task] + if softmax1_task is not None: + task_list.append(softmax1_task) + task_list.extend([correction_task, mma_task, load_task]) + if epilogue_task is not None: + task_list.append(epilogue_task) + if freed_epilogue_task is not None: + task_list.append(freed_epilogue_task) + if scheduler_task is not None: + task_list.append(scheduler_task) + task_list.append(auxiliary_task) + + if single_qkv_instance: + tmem_o_source = tmem_p0 if tmem_p0 is not None else tmem_sp0 + tmem_o_dependencies = scheduler_deps(tmem_o_source) + else: + if ( + tmem_sp1 is None + or tmem_vec1 is None + or smem_o_1 is None + or gmem_o_1 is None + or s0s1_seq is None + or tmem_stats_done_1 is None + ): + raise ValueError("paired resource graph requires peer-1 resources") + tmem_o_dependencies = scheduler_deps(tmem_sp0, tmem_sp1) + + smem_kv_deps: list[MemoryResource] = [gmem_qkv] + if smem_page_offsets_kv is not None: + smem_kv_deps.append(smem_page_offsets_kv) + if smem_page_offsets_v is not None: + smem_kv_deps.append(smem_page_offsets_v) + stats_done_0_deps = [] if cfg.stats_via_smem else [tmem_stats_done_0] + resource_dependency_graph: dict[MemoryResource, list[MemoryResource]] = { + smem_q: scheduler_deps(gmem_qkv), + smem_kv: scheduler_deps(*smem_kv_deps), + tmem_sp0: scheduler_deps(tmem_sp0, smem_q, smem_kv, *stats_done_0_deps), + tmem_vec0: scheduler_deps(tmem_sp0), + tmem_o: tmem_o_dependencies, + smem_o_0: scheduler_deps(tmem_vec0, tmem_o), + gmem_o_0: scheduler_deps(smem_o_0), + } + if not cfg.stats_via_smem: + resource_dependency_graph[tmem_stats_done_0] = [tmem_vec0] + if tmem_p0 is not None: + resource_dependency_graph[tmem_p0] = scheduler_deps(tmem_sp0) + if not single_qkv_instance: + resource_dependency_graph.update( + { + tmem_sp1: scheduler_deps( + tmem_sp1, + smem_q, + smem_kv, + s0s1_seq, + *([] if cfg.stats_via_smem else [tmem_stats_done_1]), + ), + tmem_vec1: scheduler_deps(tmem_sp1), + smem_o_1: scheduler_deps(tmem_vec1, tmem_o), + gmem_o_1: scheduler_deps(smem_o_1), + s0s1_seq: [tmem_sp0], + } + ) + if not cfg.stats_via_smem: + resource_dependency_graph[tmem_stats_done_1] = [tmem_vec1] + if work_queue is not None: + resource_dependency_graph[work_queue] = [work_queue] if is_clc_dynamic else [] + if smem_page_offsets_kv is not None: + resource_dependency_graph[smem_page_offsets_kv] = scheduler_deps(gmem_qkv) + if smem_page_offsets_v is not None: + resource_dependency_graph[smem_page_offsets_v] = scheduler_deps(gmem_qkv) + + smem_allocator = SmemAllocator() + registered_smem_resource_ids: set[int] = set() + + def add_smem_resource(resource: MemoryResource | None) -> None: + """Register one resource's data and barriers exactly once.""" + if resource is None or id(resource) in registered_smem_resource_ids: + return + smem_allocator.add_resource(resource) + registered_smem_resource_ids.add(id(resource)) + + add_smem_resource(smem_q) + add_smem_resource(smem_kv) + if smem_page_offsets_kv is not None: + add_smem_resource(smem_page_offsets_kv) + if smem_page_offsets_v is not None: + add_smem_resource(smem_page_offsets_v) + # Register every pipeline resource, including pipeline-only TMEM handoffs, + # with the unified allocator. Otherwise CUTLASS materializes those + # barriers as separate dynamic-SMEM arrays that the capacity selector + # cannot see. + add_smem_resource(tmem_sp0) + if tmem_p0 is not None: + add_smem_resource(tmem_p0) + add_smem_resource(tmem_vec0) + add_smem_resource(tmem_o) + if not cfg.stats_via_smem: + add_smem_resource(tmem_stats_done_0) + if tmem_sp1 is not None: + add_smem_resource(tmem_sp1) + if tmem_vec1 is not None: + add_smem_resource(tmem_vec1) + if s0s1_seq is not None: + add_smem_resource(s0s1_seq) + if tmem_stats_done_1 is not None and not cfg.stats_via_smem: + add_smem_resource(tmem_stats_done_1) + if work_queue is not None and work_queue.pipeline_config is not None: + add_smem_resource(work_queue) + if single_qkv_instance: + add_smem_resource(smem_o_0) + add_smem_resource(gmem_o_0) + smem_allocator.add_alias_group( + [ + [smem_o_0._alloc], + [gmem_o_0._alloc], + ] + ) + else: + add_smem_resource(smem_o_0) + add_smem_resource(smem_o_1) + add_smem_resource(gmem_o_0) + add_smem_resource(gmem_o_1) + smem_allocator.add_alias_group( + [ + [smem_o_0._alloc], + [gmem_o_0._alloc], + ] + ) + smem_allocator.add_alias_group( + [ + [smem_o_1._alloc], + [gmem_o_1._alloc], + ] + ) + tmem_ptr_alloc = smem_allocator.add_tmem_ptr( + SmemAllocation("tmem_ptr_i32", dtype=cutlass.Int32, count=2, alignment=4) + ) + dealloc_mbar_alloc = smem_allocator.add( + SmemAllocation("tmem_dealloc_mbar", dtype=cutlass.Int64, alignment=8) + ) + clc_response_alloc: SmemAllocation | None = None + if is_clc_dynamic and clc_response_ptr is None: + # Keep the CLC response inside the unified TS allocation. The kernel + # derives its pointer from this descriptor after allocate(), so no + # assumption about its physical offset is required. + assert work_queue is not None + assert work_queue.pipeline_config is not None + clc_response_alloc = smem_allocator.add( + SmemAllocation( + "clc_response", + dtype=cutlass.Int128, + count=work_queue.pipeline_config.num_stages, + alignment=16, + ) + ) + smem_allocator.compute_layout() + expected_barrier_bytes = ( + sum( + _context_pipeline_stage_counts( + cfg, + kv_stages=cfg.kv_stage, + is_clc_dynamic=is_clc_dynamic, + ).values() + ) + * _PIPELINE_BARRIER_BYTES_PER_STAGE + ) + if smem_allocator.barrier_smem_bytes != expected_barrier_bytes: + raise AssertionError( + "context pipeline barrier accounting drifted: allocator has " + f"{smem_allocator.barrier_smem_bytes} bytes, topology requires " + f"{expected_barrier_bytes} bytes" + ) + + tmem_allocator = TmemAllocator() + if single_qkv_instance: + tmem_allocator.add_resource(tmem_o) + tmem_allocator.add_resource(tmem_sp0) + tmem_allocator.add_resource(tmem_vec0) + if cfg.stage_scoped_tmem_stats and not cfg.stats_via_smem: + tmem_allocator.add_alias_group( + [ + [tmem_sp0._alloc], + [tmem_vec0._alloc], + ] + ) + else: + tmem_allocator.add_resource(tmem_sp0) + tmem_allocator.add_resource(tmem_sp1) + tmem_allocator.add_resource(tmem_vec0) + tmem_allocator.add_resource(tmem_vec1) + tmem_allocator.add_resource(tmem_o) + tmem_allocator.add_alias_group( + [ + [tmem_sp0._alloc], + [tmem_vec0._alloc], + ] + ) + tmem_allocator.add_alias_group( + [ + [tmem_sp1._alloc], + [tmem_vec1._alloc], + ] + ) + tmem_allocator.compute_layout() + + skip = not isinstance(num_kv_tiles, int) + task_manager = TaskManager( + tasks=task_list, + resource_dependency_graph=resource_dependency_graph, + skip_validation=skip, + verbose=not skip, + smem_allocator=smem_allocator, + tmem_allocator=tmem_allocator, + exhaustive_deadlock_race_check=exhaustive_deadlock_race_check, + ) + + if single_qkv_instance: + tmem_resources = [tmem_sp0, tmem_vec0, tmem_o, smem_o_0] + else: + tmem_resources = [ + tmem_sp0, + tmem_sp1, + tmem_vec0, + tmem_vec1, + tmem_o, + smem_o_0, + smem_o_1, + ] + return ( + task_manager, + tmem_resources, + tmem_ptr_alloc, + dealloc_mbar_alloc, + work_queue, + clc_response_alloc, + ) + + +def _should_use_tmem_p_pipeline(cfg: FmhaConfig) -> bool: + """Return whether P readiness needs its own TMEM pipeline resource. + + The TMEM P pipeline belongs to the staged, single-QKV-instance topology. + That path stages K/V by 128-wide head-dimension slices and uses a separate + P-ready handoff so the MMA task can overlap next-tile QK with previous-tile + PV. The paired D128 path keeps P on TmemSPResource. + """ + return cfg.single_qkv_instance and cfg.stage_kv_by_head_dim + + +_Q_ROW_SMEM_ALIGNMENT_BYTES = 128 +_PIPELINE_BARRIER_BYTES_PER_STAGE = 2 * cutlass.Int64.width // 8 + + +def _context_pipeline_stage_counts( + cfg: FmhaConfig, + *, + kv_stages: int, + is_clc_dynamic: bool, +) -> dict[str, int]: + """Return every physical pipeline's mbarrier stage count.""" + counts = { + "smem_q": cfg.q_stage, + "smem_kv": kv_stages, + "smem_page_offsets": ( + sum(cfg.page_offset_pipeline_stage_counts) + if cfg.stages_page_offsets_in_smem + else 0 + ), + "tmem_sp": cfg.mma_softmax_stage * cfg.num_qkv_instances, + "tmem_p": cfg.mma_softmax_stage if cfg.has_tmem_p_pipeline else 0, + "tmem_vec": cfg.softmax_corr_stage * cfg.num_qkv_instances, + "tmem_o": cfg.mma_corr_stage, + "smem_o": cfg.num_qkv_instances, + "s0s1_seq": 0 if cfg.single_qkv_instance else 1, + "tmem_stats_done": 0 if cfg.stats_via_smem else cfg.num_qkv_instances, + "work_queue": 1 if is_clc_dynamic else 0, + } + return {name: stages for name, stages in counts.items() if stages} + + +def _infer_single_instance_kv_stages( + cfg: FmhaConfig, + *, + is_clc_dynamic: bool, + page_table_window_entries: int | None = None, + require_cadence: bool = True, +) -> int: + """Return the deepest K/V ring that fits the exact TS SMEM footprint. + + All terms come from resource topology or public CUTLASS hardware metadata: + Q, O, correction statistics, page-ID rings, fixed control records, and one + 16-byte pipeline barrier per physical stage. The task manager remains the + authoritative check and uses the same stage-count policy below. + """ + q_row_bytes = (cfg.q_dtype.width * cfg.qk_mma_tiler[2] + 7) // 8 + q_row_bytes = ( + (q_row_bytes + _Q_ROW_SMEM_ALIGNMENT_BYTES - 1) + // _Q_ROW_SMEM_ALIGNMENT_BYTES + * _Q_ROW_SMEM_ALIGNMENT_BYTES + ) + q_tile_bytes = q_row_bytes * cfg.qk_mma_tiler[0] + + o_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + ) + o_stage_bytes = (cfg.epi_tile[0] * o_head_dim * cfg.o_dtype.width + 7) // 8 + + stats_bytes = 0 + if cfg.stats_via_smem: + stats_rows = len(cfg.softmax0_warp_ids) * cute.arch.WARP_SIZE + stats_values_per_row = 2 + stats_bytes = ( + cfg.softmax_corr_stage + * stats_rows + * stats_values_per_row + * cutlass.Float32.width + // 8 + ) + + page_offset_stage_counts = cfg.page_offset_pipeline_stage_counts + page_offset_stages = ( + sum(page_offset_stage_counts) if cfg.stages_page_offsets_in_smem else 0 + ) + if page_table_window_entries is None: + page_table_window_entries = cfg.page_table_window_entries + page_offsets_bytes = ( + page_offset_stages * page_table_window_entries * cutlass.Int32.width // 8 + ) + + control_bytes = (2 * cutlass.Int32.width + cutlass.Int64.width) // 8 + if is_clc_dynamic: + control_bytes += cutlass.Int128.width // 8 + fixed_barrier_stages = sum( + _context_pipeline_stage_counts( + cfg, + kv_stages=0, + is_clc_dynamic=is_clc_dynamic, + ).values() + ) + fixed_smem_bytes = ( + q_tile_bytes * cfg.q_stage + + o_stage_bytes + + stats_bytes + + page_offsets_bytes + + control_bytes + + fixed_barrier_stages * _PIPELINE_BARRIER_BYTES_PER_STAGE + ) + kv_dtype_width = max(cfg.k_dtype.width, cfg.v_dtype.width) + kv_stage_bytes = ( + cfg.qk_mma_tiler[1] * cfg.head_dim_per_stage_kv * kv_dtype_width // 8 + ) + kv_stage_footprint_bytes = kv_stage_bytes + _PIPELINE_BARRIER_BYTES_PER_STAGE + kv_budget_bytes = utils.get_smem_capacity_in_bytes("sm_100") - fixed_smem_bytes + memory_fit_stages = kv_budget_bytes // kv_stage_footprint_bytes + cadence_stages = cfg.num_head_dim_stages_k + cfg.num_head_dim_stages_v + if require_cadence and memory_fit_stages < cadence_stages: + raise ValueError( + "single-instance context staging requires at least " + f"{cadence_stages} K/V stages, but the shared-memory budget fits " + f"only {memory_fit_stages}" + ) + return memory_fit_stages + + +def _configure_pipeline_stages(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None: + """Set topology- and capacity-derived context pipeline stage counts.""" + cfg.q_stage = cfg.num_qkv_instances + cfg.kv_stage = 3 + cfg.has_tmem_p_pipeline = _should_use_tmem_p_pipeline(cfg) + cfg.stage_scoped_tmem_stats = cfg.has_tmem_p_pipeline + cfg.mma_softmax_stage = 2 if cfg.has_tmem_p_pipeline else 1 + cfg.softmax_corr_stage = 2 if cfg.stage_scoped_tmem_stats else 1 + # SMEM-backed D256 removes the independent StatsDone credit and always + # writes the same physical O0 accumulator. Its MMA->Correction handoff must + # therefore be single-stage so PV(i+1) cannot overwrite O0 before + # Correction consumes PV(i). TMEM-stats schedules retain their established + # two-stage O + StatsDone ordering. + cfg.mma_corr_stage = 1 if cfg.single_qkv_instance and cfg.stats_via_smem else 2 + if cfg.single_qkv_instance: + natural_page_window_entries = cute.arch.WARP_SIZE + cfg.page_table_window_entries = natural_page_window_entries + candidate_page_window_entries = cfg.page_table_window_candidate_entries + if candidate_page_window_entries > natural_page_window_entries: + candidate_kv_stages = _infer_single_instance_kv_stages( + cfg, + is_clc_dynamic=is_clc_dynamic, + page_table_window_entries=candidate_page_window_entries, + require_cadence=False, + ) + cadence_stages = cfg.num_head_dim_stages_k + cfg.num_head_dim_stages_v + if candidate_kv_stages >= cadence_stages: + cfg.page_table_window_entries = candidate_page_window_entries + cfg.kv_stage = _infer_single_instance_kv_stages( + cfg, + is_clc_dynamic=is_clc_dynamic, + ) + + +# Dense work traverses the full K domain for every Q tile, so its persistent +# mainloop keeps more registers on the load/MMA/epilogue/scheduler warpgroup. +# Causal work has a triangular, request-local K domain and retains the +# softmax/correction-heavy allocation. Both policies consume the same complete +# CTA register budget across the 8/4/4 participating warps; neither depends on +# batch size, sequence length, head count, layout, or a measured crossover. +_EARLY_TILE_SUM_DENSE_REGISTER_BUDGET = (176, 80, 80) +_EARLY_TILE_SUM_CAUSAL_REGISTER_BUDGET = (184, 88, 56) + + +def _configure_early_tile_sum_policy( + cfg: FmhaConfig, + *, + is_persistent: bool, +) -> None: + """Couple the concrete early-sum algorithm to its register budget.""" + # The current implementation has the Q2/KV1 paired topology and the D256 + # single-instance topology. Q2/KV1 benefits across scheduler and head + # mappings; keep the persistence condition explicit for parity with the + # upstream policy if another paired geometry is added later. + # D256 staged FP8 retires each probability through conversion and row-sum + # reduction eight values after EXP2. It therefore returns a scalar tile + # sum through the same task-local path as the paired early-sum policy while + # retaining the D256 200/192/112 register split below. + # D128 uses the same early-sum dataflow for every scheduler and storage + # layout; scheduler selection must not silently change its math pipeline. + cfg.enable_early_tile_sum = cfg.uses_d256_fp8_softmax_cadence or ( + cfg.uses_early_tile_sum and (not cfg.single_qkv_instance or is_persistent) + ) + if cfg.single_qkv_instance: + return + if not cfg.enable_early_tile_sum: + # Preserve the established register split for unsupported paired paths. + # The legacy 192/96/32 fallback serves Q1/KV2, which is not implemented. + return + ( + cfg.num_regs_softmax, + cfg.num_regs_correction, + cfg.num_regs_other, + ) = ( + _EARLY_TILE_SUM_CAUSAL_REGISTER_BUDGET + if cfg.is_causal + else _EARLY_TILE_SUM_DENSE_REGISTER_BUDGET + ) + + +def _configure_smem_shapes(cfg: FmhaConfig) -> None: + """Derive per-stage SMEM element counts from the configured tile shapes.""" + kv_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_kv_by_head_dim else cfg.qk_mma_tiler[2] + ) + o_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + ) + cfg.sQ_shape = ( + cfg.q_stage, + cfg.qk_mma_tiler[0] * cfg.qk_mma_tiler[2], + ) + cfg.sK_shape = ( + cfg.kv_stage, + cfg.qk_mma_tiler[1] * kv_head_dim, + ) + cfg.sO_stage_elements = cfg.epi_tile[0] * o_head_dim + + +def _validate_tmem_columns(cfg: FmhaConfig) -> None: + """Reject tile shapes that exceed the selected context TMEM layout.""" + sp_tmem_cols = cfg.num_qkv_instances * cfg.qk_mma_tiler[1] * cfg.mma_softmax_stage + o_tmem_cols = cfg.epi_tile[1] + stats_tmem_cols = 0 + if cfg.single_qkv_instance and not cfg.stage_scoped_tmem_stats: + stats_tmem_cols = cfg.tmem_stats_cols + required_tmem_cols = ( + sp_tmem_cols + cfg.num_qkv_instances * o_tmem_cols + stats_tmem_cols + ) + if required_tmem_cols > cfg.tmem_alloc_cols: + raise ValueError( + f"head dimension {o_tmem_cols} requires {required_tmem_cols} " + f"TMEM columns for the current FMHA context schedule, " + f"but only {cfg.tmem_alloc_cols} are available" + ) + + +def _configure_single_instance_tmem_layout(cfg: FmhaConfig) -> None: + """Select the one-S/P, one-O TMEM layout used for d>128 context FMHA.""" + cfg.tmem_o0_offset = 0 + cfg.tmem_s0_offset = cfg.epi_tile[1] + cfg.tmem_p0_offset = cfg.tmem_s0_offset + cfg.tmem_x_load_s + if cfg.stage_scoped_tmem_stats: + cfg.tmem_vec0_offset = cfg.tmem_s0_offset + else: + cfg.tmem_vec0_offset = cfg.tmem_s0_offset + cfg.qk_mma_tiler[1] + + +def _configure_single_instance_warp_layout(cfg: FmhaConfig) -> None: + """Use one softmax warpgroup and a compact producer/consumer warpgroup.""" + cfg.softmax0_warp_ids = (0, 1, 2, 3) + cfg.softmax1_warp_ids = () + cfg.correction_warp_ids = (4, 5, 6, 7) + cfg.mma_warp_id = 8 + cfg.load_warp_id = 9 + cfg.epilogue_warp_id = 10 + cfg.empty_warp_id = 11 + cfg.block_warps = 12 + # Use the validated D256 register split: four softmax warps, four + # correction warps, and four producer/scheduler warps. + cfg.num_regs_softmax = 200 + cfg.num_regs_correction = 192 + cfg.num_regs_other = 112 + + +def _configure_head_dim_staging(cfg: FmhaConfig) -> None: + """Split d>128 single-instance K/V/O staging into 128-wide slices.""" + cfg.head_dim_per_stage_kv = 0 + cfg.num_head_dim_stages_k = 1 + cfg.num_head_dim_stages_v = 1 + cfg.num_o_head_dim_stages = 1 + cfg.stage_kv_by_head_dim = False + cfg.stage_o_by_head_dim = False + if cfg.num_qkv_instances != 1: + return + cfg.head_dim_per_stage_kv = 128 + cfg.num_head_dim_stages_k = cfg.qk_mma_tiler[2] // cfg.head_dim_per_stage_kv + cfg.num_head_dim_stages_v = cfg.pv_mma_tiler[1] // cfg.head_dim_per_stage_kv + cfg.num_o_head_dim_stages = cfg.epi_tile[1] // cfg.head_dim_per_stage_kv + cfg.stage_kv_by_head_dim = True + cfg.stage_o_by_head_dim = True + # Stage K, V, and O as 128-wide head-dimension slices so the d>128 K/V + # pipeline can run deeper without exceeding Blackwell's SMEM budget. + + +def _configure_head_paired_tilers( + cfg: FmhaConfig, + *, + mma_tiler_mn: tuple[int, int], + d: int, +) -> None: + """Set CTA/MMA/epilogue tile shapes for head-paired FMHA.""" + # Head-paired maps the two peer tiles onto Q heads instead of sequence + # rows, so it keeps a one-tile CTA shape while reusing the same FmhaConfig + # and task-manager entry point as query-paired FMHA. + mma_tiler = (*mma_tiler_mn, d) + cfg.qk_mma_tiler = mma_tiler + cfg.pv_mma_tiler = (mma_tiler[0], mma_tiler[2], mma_tiler[1]) + cfg.epi_tile = cfg.pv_mma_tiler[:2] + + +def _configure_head_paired_tma_copy_metadata( + cfg: FmhaConfig, + *, + q_dtype: type, + k_dtype: type, + o_dtype: type, +) -> None: + """Derive TMA copy granularities for head-paired Q/K/V/O tensors.""" + inner_dim_size = cfg.qk_mma_tiler[2] * q_dtype.width // 8 + cfg.tma_copy_qkv_iters = 1 + if inner_dim_size % 128 == 0: + cfg.tma_copy_qkv_iters = inner_dim_size // 128 + elif inner_dim_size != 64 and inner_dim_size != 32: + raise RuntimeError(f"Unsupported inner dimension size: {inner_dim_size}") + tma_copy_qkv_granu_inner = cfg.qk_mma_tiler[2] // cfg.tma_copy_qkv_iters + + cfg.q_tile_m = cfg.qk_mma_tiler[0] + cfg.tma_copy_q_elements = cfg.sQ_shape[1] + cfg.tma_copy_q_granu_inner = tma_copy_qkv_granu_inner + cfg.tma_copy_q_granu_elems = cfg.tma_copy_q_elements // cfg.tma_copy_qkv_iters + cfg.tma_copy_q_bytes = cfg.tma_copy_q_elements * q_dtype.width // 8 + + cfg.seq_tile_n = cfg.qk_mma_tiler[1] + cfg.kv_tile_n = cfg.qk_mma_tiler[1] + kv_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_kv_by_head_dim else cfg.qk_mma_tiler[2] + ) + cfg.tma_copy_kv_elements = cfg.sK_shape[1] + cfg.tma_copy_kv_granu_inner = tma_copy_qkv_granu_inner + cfg.tma_copy_kv_stage_iters = kv_head_dim // tma_copy_qkv_granu_inner + cfg.tma_copy_kv_granu_elems = ( + cfg.tma_copy_kv_elements // cfg.tma_copy_kv_stage_iters + ) + cfg.tma_copy_kv_bytes = cfg.tma_copy_kv_elements * k_dtype.width // 8 + + output_inner_dim_size = cfg.epi_tile[1] * o_dtype.width // 8 + cfg.tma_copy_o_iters = 1 + if output_inner_dim_size % 128 == 0: + cfg.tma_copy_o_iters = output_inner_dim_size // 128 + elif output_inner_dim_size != 64 and output_inner_dim_size != 32: + raise RuntimeError( + f"Unsupported output inner dimension size: {output_inner_dim_size}" + ) + cfg.tma_copy_o_granu_inner = cfg.epi_tile[1] // cfg.tma_copy_o_iters + o_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + ) + cfg.tma_copy_o_stage_iters = o_head_dim // cfg.tma_copy_o_granu_inner + cfg.tma_copy_o_elements = cfg.epi_tile[0] * o_head_dim + cfg.tma_copy_o_granu_elems = cfg.tma_copy_o_elements // cfg.tma_copy_o_stage_iters + + +def _configure_common_launch_flags( + cfg: FmhaConfig, + *, + d: int, + h_r: int, + is_causal: bool, + balance_causal_workload: bool, + window_size_left: int, + has_variable_window: bool, +) -> None: + """Fill launch flags shared by query-paired and head-paired modes.""" + cfg.h_r = h_r + cfg.is_causal = is_causal + cfg.balance_causal_workload = balance_causal_workload + cfg.window_size_left = window_size_left + cfg.has_variable_window = has_variable_window + + +def _causal_domain_kwargs( + *, + num_kv_tiles: int | Int32, + tile_size_q: int, + tile_size_kv: int, + q_offset: int | Int32, + seq_idx: int, + batch_idx: int | None = None, + cum_seqlen_q: cute.Tensor | None = None, + cum_seqlen_k: cute.Tensor | None = None, + seq_lens_kv: cute.Pointer | None = None, + runtime_kv_tile_multiple: int = 1, + offset: int, + reverse_seq_tiles: int | Int32 | None = None, + window_size_left: int | None = None, + packed_window: bool = False, +) -> DomainKwargs: + """Build kwargs for a causal/window-aware task domain. + + ``offset`` is the loop-domain decrement applied after the causal/window + tile count is computed. The schedule uses it to derive N, N-1, and N-2 + loop domains from the same causal formula without changing the Q/K tile + coordinates or the S_q < S_kv ``q_offset`` mask shift. + """ + result: DomainKwargs = { + "task_class": CausalDomainTask, + "num_kv_tiles": num_kv_tiles, + "tile_size_q": tile_size_q, + "tile_size_kv": tile_size_kv, + "q_offset": q_offset, + "seq_idx": seq_idx, + "offset": offset, + } + if reverse_seq_tiles is not None: + result["reverse_seq_tiles"] = reverse_seq_tiles + if cum_seqlen_q is not None: + if batch_idx is None or (cum_seqlen_k is None and seq_lens_kv is None): + raise ValueError( + "runtime causal domains require batch_idx, cumulative Q offsets, " + "and cumulative K offsets or paged K/V lengths" + ) + result["batch_idx"] = batch_idx + result["cum_seqlen_q"] = cum_seqlen_q + if cum_seqlen_k is not None: + result["cum_seqlen_k"] = cum_seqlen_k + if seq_lens_kv is not None: + result["seq_lens_kv"] = seq_lens_kv + if runtime_kv_tile_multiple > 1: + result["runtime_kv_tile_multiple"] = runtime_kv_tile_multiple + if window_size_left is not None: + result["window_size_left"] = window_size_left + if packed_window: + result["packed_window"] = True + return result + + +def _select_fmha_domain_policy( + cfg: FmhaConfig, + *, + num_kv_tiles: int | Int32, + q_offset: int | Int32, + cum_seqlen_q: cute.Tensor | None, + cum_seqlen_k: cute.Tensor | None, + seq_lens_kv: cute.Pointer | None = None, + variable_window_token_starts: cute.Tensor | None = None, + variable_window_token_ends: cute.Tensor | None = None, + variable_window_cta_starts: cute.Tensor | None = None, + variable_window_q_stride: int | Int32 = 0, +) -> FmhaDomainPolicy: + """Select loop domains and softmax masks for the configured FMHA mode.""" + seq_idx = cfg.work_tile_coord_indices[0] + batch_idx = cfg.work_tile_coord_indices[2] + reverse_seq_tiles = ( + cfg.num_seq_tiles + if cfg.uses_causal_reversed_head_batch_seq_tile_order + else None + ) + if cfg.has_variable_window: + if ( + variable_window_token_starts is None + or variable_window_token_ends is None + or variable_window_cta_starts is None + ): + raise ValueError("VariableWindow domain requires start and end tensors") + base_kwargs: DomainKwargs = { + "task_class": VariableWindowDomainTask, + "variable_window_token_starts": variable_window_token_starts, + "variable_window_token_ends": variable_window_token_ends, + "variable_window_cta_starts": variable_window_cta_starts, + "num_kv_tiles": num_kv_tiles, + "q_stride": variable_window_q_stride, + "tile_size_q": cfg.cta_tiler[0], + "tile_size_kv": cfg.kv_tile_n, + "seq_idx": seq_idx, + "batch_idx": batch_idx, + } + domain_n_kwargs = {**base_kwargs, "offset": 0} + domain_n_minus_1_kwargs = {**base_kwargs, "offset": 1} + return FmhaDomainPolicy( + domain_n_kwargs=domain_n_kwargs, + domain_n_minus_1_kwargs=domain_n_minus_1_kwargs, + softmax0_domain_kwargs=domain_n_kwargs, + softmax1_domain_kwargs=domain_n_kwargs, + ) + # Head-paired causal/window: both peers use the same Q sequence tile from + # adjacent Q heads, so both softmax tasks share the same causal tail domain. + if cfg.head_paired and cfg.is_causal: + causal_n = _causal_domain_kwargs( + num_kv_tiles=num_kv_tiles, + tile_size_q=cfg.q_tile_m, + tile_size_kv=cfg.kv_tile_n, + q_offset=q_offset, + seq_idx=seq_idx, + offset=0, + reverse_seq_tiles=reverse_seq_tiles, + window_size_left=cfg.window_size_left, + packed_window=cfg.has_varlen, + ) + causal_n_minus_1 = _causal_domain_kwargs( + num_kv_tiles=num_kv_tiles, + tile_size_q=cfg.q_tile_m, + tile_size_kv=cfg.kv_tile_n, + q_offset=q_offset, + seq_idx=seq_idx, + offset=1, + reverse_seq_tiles=reverse_seq_tiles, + window_size_left=cfg.window_size_left, + packed_window=cfg.has_varlen, + ) + return FmhaDomainPolicy( + domain_n_kwargs=causal_n, + domain_n_minus_1_kwargs=causal_n_minus_1, + softmax0_domain_kwargs={ + **causal_n_minus_1, + "task_class": CausalSoftmaxDomainTask, + }, + softmax1_domain_kwargs={ + **causal_n_minus_1, + "task_class": CausalSoftmaxDomainTask, + }, + ) + # Head-paired dense: no causal/window trimming is needed, so both peers use + # the full static K/V domain and the same non-tail softmax mask settings. + if cfg.head_paired: + domain_n_kwargs: DomainKwargs = {"domain": num_kv_tiles} + domain_n_minus_1_kwargs: DomainKwargs = {"domain": num_kv_tiles - 1} + return FmhaDomainPolicy( + domain_n_kwargs=domain_n_kwargs, + domain_n_minus_1_kwargs=domain_n_minus_1_kwargs, + softmax0_domain_kwargs=domain_n_kwargs, + softmax1_domain_kwargs=domain_n_kwargs, + ) + # Query-paired causal: peer0 and peer1 cover consecutive Q sequence tiles, + # so peer0 may need one fewer K/V loop tile than peer1. + if cfg.is_causal: + if cfg.causal_single_kv_tile: + # The whole fixed K/V extent is one tile. Static N=1/N-1=0 + # domains balance the existing head/tail protocol without the + # generic synthetic peer0 tail iteration. + domain_n_kwargs: DomainKwargs = {"domain": 1} + domain_n_minus_1_kwargs: DomainKwargs = {"domain": 0} + return FmhaDomainPolicy( + domain_n_kwargs=domain_n_kwargs, + domain_n_minus_1_kwargs=domain_n_minus_1_kwargs, + softmax0_domain_kwargs=domain_n_minus_1_kwargs, + softmax1_domain_kwargs=domain_n_minus_1_kwargs, + ) + runtime_kv_tile_multiple = ( + (cfg.cta_tiler[0] + cfg.kv_tile_n - 1) // cfg.kv_tile_n + if cfg.skip_causal_invalid_peer0 + else 1 + ) + causal_n = _causal_domain_kwargs( + num_kv_tiles=num_kv_tiles, + tile_size_q=cfg.cta_tiler[0], + tile_size_kv=cfg.kv_tile_n, + q_offset=q_offset, + seq_idx=seq_idx, + batch_idx=batch_idx, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=seq_lens_kv, + runtime_kv_tile_multiple=runtime_kv_tile_multiple, + offset=0, + reverse_seq_tiles=reverse_seq_tiles, + ) + causal_n_minus_1 = _causal_domain_kwargs( + num_kv_tiles=num_kv_tiles, + tile_size_q=cfg.cta_tiler[0], + tile_size_kv=cfg.kv_tile_n, + q_offset=q_offset, + seq_idx=seq_idx, + batch_idx=batch_idx, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=seq_lens_kv, + runtime_kv_tile_multiple=runtime_kv_tile_multiple, + offset=1, + reverse_seq_tiles=reverse_seq_tiles, + ) + causal_n_minus_2 = _causal_domain_kwargs( + num_kv_tiles=num_kv_tiles, + tile_size_q=cfg.cta_tiler[0], + tile_size_kv=cfg.kv_tile_n, + q_offset=q_offset, + seq_idx=seq_idx, + batch_idx=batch_idx, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=seq_lens_kv, + runtime_kv_tile_multiple=runtime_kv_tile_multiple, + offset=2, + reverse_seq_tiles=reverse_seq_tiles, + ) + # Query-paired causal maps peer0 one Q tile before peer1. + # Whether peer0 has an extra invalid SP tile depends on the QK MMA + # tile geometry: if peer0 spans fewer K/V tiles than the paired CTA, + # Softmax0 uses the N-2 domain and consumes the remaining extra tile in + # TAIL. + # Otherwise it uses the same N-1 domain as peer1. S_q < S_kv keeps N-1 + # and lets the shifted causal mask handle the right edge. + mma0_has_invalid_tail = cfg.skip_causal_invalid_peer0 + mma0_domain = causal_n_minus_2 if mma0_has_invalid_tail else causal_n_minus_1 + mma1_domain = causal_n_minus_1 + softmax0_domain = { + **mma0_domain, + "task_class": CausalSoftmaxDomainTask, + } + softmax1_domain = { + **mma1_domain, + "task_class": CausalSoftmaxDomainTask, + } + return FmhaDomainPolicy( + domain_n_kwargs=causal_n, + domain_n_minus_1_kwargs=causal_n_minus_1, + softmax0_domain_kwargs=softmax0_domain, + softmax1_domain_kwargs=softmax1_domain, + ) + # Dense query-paired fallback: no head-pairing or causal mask, so both + # peers traverse the full static K/V domain. + domain_n_kwargs = {"domain": num_kv_tiles} + domain_n_minus_1_kwargs = {"domain": num_kv_tiles - 1} + return FmhaDomainPolicy( + domain_n_kwargs=domain_n_kwargs, + domain_n_minus_1_kwargs=domain_n_minus_1_kwargs, + softmax0_domain_kwargs=domain_n_kwargs, + softmax1_domain_kwargs=domain_n_kwargs, + ) + + +def build_fmha_task_manager( + cfg: FmhaConfig, + tile_sched_params: ( + utils.PersistentTileSchedulerParams + | utils.ClcDynamicPersistentTileSchedulerParams + | None + ), + tma_q_desc: cutlass.Pointer | None, + tma_k_desc: cutlass.Pointer | None, + tma_v_desc: cutlass.Pointer | None, + tma_o_desc: cutlass.Pointer | None, + cum_seqlen_q: cute.Tensor | None, + cum_seqlen_k: cute.Tensor | None, + num_kv_tiles: int | Int32, + variable_window_token_starts: cute.Tensor | None = None, + variable_window_token_ends: cute.Tensor | None = None, + variable_window_cta_starts: cute.Tensor | None = None, + variable_window_q_stride: int | Int32 = 0, + scale_softmax_log2: cute.Tensor | None = None, + output_scale: cute.Tensor | None = None, + q_offset: int | Int32 = 0, + g_block_tables: cute.Pointer | None = None, + block_table_row_stride: int | Int32 = 0, + g_seq_lens_kv: cute.Pointer | None = None, + max_seq_len_kv: int | Int32 | None = None, + is_persistent: bool = True, + is_clc_dynamic: bool = False, + clc_response_ptr: cute.Pointer | None = None, + exhaustive_deadlock_race_check: bool = True, +) -> Tuple[ + TaskManager, + list[MemoryResource], + SmemAllocation, + SmemAllocation, + WorkQueue | None, + SmemAllocation | None, +]: + """Build the FMHA TaskManager using the shared context TS graph. + + Runtime scale arguments follow the same contract as + :func:`build_context_task_manager`: one-element float32 device tensors, + indexed at element 0 and cached by resource auxiliary work. + """ + if ( + cfg.causal_single_kv_tile + and isinstance(num_kv_tiles, int) + and num_kv_tiles != 1 + ): + raise ValueError( + "causal_single_kv_tile requires exactly one compile-time K/V tile" + ) + if cfg.causal_single_kv_tile and (cfg.use_paged_kv or cfg.has_varlen): + raise ValueError("causal_single_kv_tile requires fixed contiguous K/V storage") + domain_num_kv_tiles = num_kv_tiles + if cfg.skip_causal_invalid_peer0: + # Query-paired causal skips peer0 work with a constexpr last-loop test. + # Partial final CTAs need the task domain padded so the extra peer0 slot + # remains statically invalid; aligned static domains need no padding. + paired_kv_tiles = (cfg.cta_tiler[0] + cfg.kv_tile_n - 1) // cfg.kv_tile_n + if not isinstance(num_kv_tiles, int) or num_kv_tiles % paired_kv_tiles != 0: + domain_num_kv_tiles = ( + cute.ceil_div(num_kv_tiles, paired_kv_tiles) * paired_kv_tiles + ) + if cfg.reuses_page_table_windows: + # The paged plan's maximum page count is compile-time semantic geometry. + # Its tile count equals ceil(max_seq_len_kv / kv_tile_n), so exposing it + # here gives stock TaskManager a structural stride-window loop without a + # shape-tuned threshold or a custom control-flow feature. + pages_per_kv_tile = cfg.kv_tile_n // cfg.num_tokens_per_page + static_paged_num_kv_tiles = ( + cfg.max_num_pages_per_seq_kv + pages_per_kv_tile - 1 + ) // pages_per_kv_tile + domain_num_kv_tiles = static_paged_num_kv_tiles + # Keep equal-length fixed launches constexpr-zero while preserving the + # runtime per-request offset path for packed/bottom-right attention. + effective_q_offset = q_offset if cfg.has_q_offset else 0 + domain_policy = _select_fmha_domain_policy( + cfg, + num_kv_tiles=domain_num_kv_tiles, + q_offset=effective_q_offset, + # A uniform packed plan remains uniform under its replay contract, so + # keep that specialization free of redundant GMEM indptr loads. Mixed + # packed and paged plans derive their causal domain from per-run Q + # offsets and the corresponding contiguous or paged K/V lengths. + cum_seqlen_q=( + cum_seqlen_q if cfg.has_varlen and not cfg.has_uniform_varlen else None + ), + cum_seqlen_k=( + cum_seqlen_k if cfg.has_varlen and not cfg.has_uniform_varlen else None + ), + seq_lens_kv=(g_seq_lens_kv if cfg.use_paged_kv else None), + variable_window_token_starts=variable_window_token_starts, + variable_window_token_ends=variable_window_token_ends, + variable_window_cta_starts=variable_window_cta_starts, + variable_window_q_stride=variable_window_q_stride, + ) + + return build_context_task_manager( + cfg=cfg, + tile_sched_params=tile_sched_params, + tma_q_desc=tma_q_desc, + tma_k_desc=tma_k_desc, + tma_v_desc=tma_v_desc, + tma_o_desc=tma_o_desc, + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + variable_window_token_starts=variable_window_token_starts, + variable_window_token_ends=variable_window_token_ends, + variable_window_cta_starts=variable_window_cta_starts, + variable_window_q_stride=variable_window_q_stride, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + g_block_tables=g_block_tables, + block_table_row_stride=block_table_row_stride, + g_seq_lens_kv=g_seq_lens_kv, + max_seq_len_kv=max_seq_len_kv, + num_kv_tiles=domain_num_kv_tiles, + q_offset=effective_q_offset, + domain_n_kwargs=domain_policy.domain_n_kwargs, + domain_n_minus_1_kwargs=domain_policy.domain_n_minus_1_kwargs, + softmax0_domain_kwargs=domain_policy.softmax0_domain_kwargs, + softmax1_domain_kwargs=domain_policy.softmax1_domain_kwargs, + is_persistent=is_persistent, + is_clc_dynamic=is_clc_dynamic, + clc_response_ptr=clc_response_ptr, + exhaustive_deadlock_race_check=exhaustive_deadlock_race_check, + ) + + +# --------------------------------------------------------------------------- +# GPU Kernel +# --------------------------------------------------------------------------- + + +class FmhaTs: + """Warp-specialised persistent FMHA kernel using the TS framework. + + Usage:: + + fmha = FmhaTs(qk_acc_dtype=Float32, pv_acc_dtype=Float32, + mma_tiler_mn=(128, 128)) + fmha(q_cute, k_cute, v_cute, o_cute, stream) + + Parameters + ---------- + qk_acc_dtype : type, optional + Accumulator dtype for QK GEMM (default: Float32). + pv_acc_dtype : type, optional + Accumulator dtype for PV GEMM (default: Float32). + in_dtype : type, optional + Input tensor dtype for Q, K, V (default: Float16). + out_dtype : type, optional + Output tensor dtype for O (default: Float16). + mma_tiler_mn : Tuple[int, int], optional + MMA tile shape (M, N) (default: (128, 128)). + d : int, optional + Head dimension (default: 128). + is_persistent : bool, optional + Use persistent scheduling (default: True). + is_causal : bool, optional + Enable causal masking (default: False). + balance_causal_workload : bool, optional + Use TRT-style causal workload balancing: head_batch_seq logical tile + order with reversed Q sequence tiles. Paired causal CLC schedules + enable this automatically; setting the flag also requests it for other + causal scheduler topologies. + is_clc_dynamic : bool, optional + Use CLC dynamic persistent scheduling (default: False). + Requires ``is_persistent=True``. + h_r : int, optional + Number of GQA head repeats (default: 1). Head-paired mode requires + grouped-query attention with an even repeat count. + enable_skip_correction : bool, optional + Enable skip-correction for softmax rescaling (default: True). + use_paged_kv : bool, optional + Read K/V from a physical page pool through a fixed block table. + num_tokens_per_page : int, optional + Number of K/V tokens stored in each physical page (default: 32). + max_kv_len : int, optional + Planned upper bound for each request's K/V length. Paged kernels derive + their static page and tile capacity from this bound (default: 1). + causal_single_kv_tile : bool, optional + Use the fixed causal one-K/V-tile task domains. The context runner + enables this only for query-paired, fixed-length inputs whose K/V + extent fits one 128-token tile (default: False). + """ + + def __init__( + self, + qk_acc_dtype: type | None = None, + pv_acc_dtype: type | None = None, + in_dtype: type | None = None, + out_dtype: type | None = None, + mma_tiler_mn: Tuple[int, int] = (128, 128), + d: int = 128, + is_persistent: bool = True, + is_causal: bool = False, + balance_causal_workload: bool = False, + is_clc_dynamic: bool = False, + head_paired: bool = False, + window_size_left: int = 0, + has_variable_window: bool = False, + h_r: int = 1, + enable_skip_correction: bool = True, + use_paged_kv: bool = False, + num_tokens_per_page: int = 32, + max_kv_len: int = 1, + causal_single_kv_tile: bool = False, + exhaustive_deadlock_race_check: bool = True, + ) -> None: + """Initialize mode-specific tiling, dtype, and schedule configuration.""" + head_paired = resolve_head_paired_mode( + head_paired=head_paired, + is_causal=is_causal, + window_size_left=window_size_left, + ) + if causal_single_kv_tile and (not is_causal or head_paired): + raise ValueError( + "causal_single_kv_tile requires query-paired causal attention" + ) + if causal_single_kv_tile and use_paged_kv: + raise ValueError( + "causal_single_kv_tile requires fixed contiguous K/V storage" + ) + if is_clc_dynamic and not is_persistent: + raise ValueError("CLC dynamic scheduling requires persistent mode") + if head_paired and not is_persistent: + raise ValueError("Head-paired scheduling requires persistent mode") + if has_variable_window and (is_causal or window_size_left > 0): + raise ValueError( + "VariableWindow bounds replace causal and sliding-window masks" + ) + if has_variable_window and use_paged_kv: + raise NotImplementedError( + "variable-window masking is not supported for paged context" + ) + validate_head_paired_head_ratio(head_paired=head_paired, h_r=h_r) + if use_paged_kv: + if num_tokens_per_page not in _SUPPORTED_CONTEXT_PAGE_SIZES: + raise ValueError( + "paged context requires num_tokens_per_page in " + f"{_SUPPORTED_CONTEXT_PAGE_SIZES}; got " + f"{num_tokens_per_page}" + ) + if max_kv_len < 1: + raise ValueError( + f"max_kv_len must be >= 1 for paged context, got {max_kv_len}" + ) + self.is_persistent = is_persistent + self.is_causal = is_causal + self.is_clc_dynamic = is_clc_dynamic + self.exhaustive_deadlock_race_check = exhaustive_deadlock_race_check + + q_dtype = in_dtype or cutlass.Float16 + k_dtype = in_dtype or cutlass.Float16 + v_dtype = in_dtype or cutlass.Float16 + o_dtype = out_dtype or cutlass.Float16 + + cfg = FmhaConfig() + self.cfg = cfg + if d > 128: + cfg.num_qkv_instances = 1 + cfg.use_paged_kv = use_paged_kv + single_instance_persistent = ( + is_persistent and cfg.single_qkv_instance and not head_paired + ) + # Paired Q2 and persistent D256 schedules route correction statistics + # through a compact SMEM ring and omit the per-K StatsDone + # serialization. This keeps the stats payload disjoint from the S/P + # TMEM columns while producer latency varies across storage layouts. + cfg.stats_via_smem = single_instance_persistent or (not cfg.single_qkv_instance) + cfg.fuse_epilogue_into_correction = cfg.single_qkv_instance + cfg.num_tokens_per_page = num_tokens_per_page + cfg.max_num_pages_per_seq_kv = ( + max_kv_len + num_tokens_per_page - 1 + ) // num_tokens_per_page + cfg.causal_single_kv_tile = causal_single_kv_tile + # FP16/BF16 causal attention retains the default 192/96/32 + # softmax/correction/auxiliary split. Other topologies start from + # 184/88/56; the paired D128 early-sum policy is rebalanced below. + # Every selected split totals 2048 registers across the 16 warps. + if not (is_causal and q_dtype.width == 16): + cfg.num_regs_softmax = 184 + cfg.num_regs_correction = 88 + cfg.num_regs_other = 56 + cfg.enable_skip_correction = enable_skip_correction + cfg.qk_acc_dtype = qk_acc_dtype or cutlass.Float32 + cfg.pv_acc_dtype = pv_acc_dtype or cutlass.Float32 + + # Store dtypes as compile-time constants + cfg.q_dtype = q_dtype + cfg.k_dtype = k_dtype + cfg.v_dtype = v_dtype + cfg.o_dtype = o_dtype + cfg.head_paired = head_paired + cfg.is_causal = is_causal + balance_causal_workload = balance_causal_workload or ( + is_causal and is_clc_dynamic and not cfg.single_qkv_instance + ) + + if head_paired: + _configure_head_paired_tilers(cfg, mma_tiler_mn=mma_tiler_mn, d=d) + _configure_head_dim_staging(cfg) + _configure_pipeline_stages(cfg, is_clc_dynamic=is_clc_dynamic) + if cfg.single_qkv_instance: + _configure_single_instance_tmem_layout(cfg) + _configure_single_instance_warp_layout(cfg) + _configure_smem_shapes(cfg) + _validate_tmem_columns(cfg) + _configure_head_paired_tma_copy_metadata( + cfg, + q_dtype=q_dtype, + k_dtype=k_dtype, + o_dtype=o_dtype, + ) + _configure_common_launch_flags( + cfg, + d=d, + h_r=h_r, + is_causal=is_causal, + balance_causal_workload=balance_causal_workload, + window_size_left=window_size_left, + has_variable_window=has_variable_window, + ) + _configure_early_tile_sum_policy(cfg, is_persistent=is_persistent) + return + + # MMA tiler: (M, N, K) = (128, 128, 128) + mma_tiler = (*mma_tiler_mn, d) + cfg.qk_mma_tiler = mma_tiler + cfg.pv_mma_tiler = (mma_tiler[0], mma_tiler[2], mma_tiler[1]) + cfg.epi_tile = cfg.pv_mma_tiler[:2] + _configure_head_dim_staging(cfg) + _configure_pipeline_stages(cfg, is_clc_dynamic=is_clc_dynamic) + if cfg.single_qkv_instance: + _configure_single_instance_tmem_layout(cfg) + _configure_single_instance_warp_layout(cfg) + + # Pipeline stages and SMEM shapes (stages, elements_per_stage) + _configure_smem_shapes(cfg) + _validate_tmem_columns(cfg) + + # TMA copy granularity for Q + qkv_tma_bits = cfg.qk_mma_tiler[2] * q_dtype.width + if qkv_tma_bits % (128 * 8) != 0: + raise ValueError( + "FMHA TS requires a 128-byte aligned Q/K/V inner dimension, " + f"got {qkv_tma_bits // 8} bytes" + ) + cfg.tma_copy_qkv_iters = qkv_tma_bits // (128 * 8) + cfg.q_tile_m = cfg.qk_mma_tiler[0] + cfg.tma_copy_q_granu_inner = cfg.qk_mma_tiler[2] // cfg.tma_copy_qkv_iters + cfg.tma_copy_q_elements = cfg.sQ_shape[1] + cfg.tma_copy_q_granu_elems = cfg.tma_copy_q_elements // cfg.tma_copy_qkv_iters + cfg.tma_copy_q_bytes = cfg.tma_copy_q_elements * q_dtype.width // 8 + + # TMA copy granularity for KV + cfg.kv_tile_n = cfg.qk_mma_tiler[1] + kv_head_dim = ( + cfg.head_dim_per_stage_kv + if cfg.stage_kv_by_head_dim + else cfg.qk_mma_tiler[2] + ) + cfg.tma_copy_kv_granu_inner = cfg.qk_mma_tiler[2] // cfg.tma_copy_qkv_iters + cfg.tma_copy_kv_elements = cfg.sK_shape[1] + cfg.tma_copy_kv_stage_iters = kv_head_dim // cfg.tma_copy_kv_granu_inner + cfg.tma_copy_kv_granu_elems = ( + cfg.tma_copy_kv_elements // cfg.tma_copy_kv_stage_iters + ) + cfg.tma_copy_kv_bytes = cfg.tma_copy_kv_elements * k_dtype.width // 8 + + # TMA copy granularity for O + cfg.tma_copy_o_iters = (cfg.epi_tile[1] * o_dtype.width) // 1024 + cfg.tma_copy_o_granu_inner = cfg.epi_tile[1] // cfg.tma_copy_o_iters + o_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + ) + cfg.tma_copy_o_stage_iters = o_head_dim // cfg.tma_copy_o_granu_inner + cfg.tma_copy_o_elements = cfg.epi_tile[0] * o_head_dim + cfg.tma_copy_o_granu_elems = ( + cfg.tma_copy_o_elements // cfg.tma_copy_o_stage_iters + ) + + _configure_common_launch_flags( + cfg, + d=d, + h_r=h_r, + is_causal=is_causal, + balance_causal_workload=balance_causal_workload, + window_size_left=window_size_left, + has_variable_window=has_variable_window, + ) + _configure_early_tile_sum_policy(cfg, is_persistent=is_persistent) + + # --------------------------------------------------------------------------- + # Host entry point + # --------------------------------------------------------------------------- + @cute.jit + def __call__( + self, + q_cute: cute.Tensor, + k_cute: cute.Tensor, + v_cute: cute.Tensor, + o_cute: cute.Tensor, + scale_softmax_log2: cute.Tensor, + output_scale: cute.Tensor, + max_active_clusters: int, + stream: cuda_drv.CUstream, + cum_seqlen_q: cute.Tensor | None = None, + cum_seqlen_k: cute.Tensor | None = None, + max_seqlen_q: Int32 | None = None, + max_seqlen_k: Int32 | None = None, + block_tables: cute.Tensor | None = None, + seq_lens_kv: cute.Tensor | None = None, + variable_window_token_starts: cute.Tensor | None = None, + variable_window_token_ends: cute.Tensor | None = None, + variable_window_cta_starts: cute.Tensor | None = None, + ) -> None: + """Set up TMA descriptors, compute grid, and launch the kernel. + + ``scale_softmax_log2`` and ``output_scale`` must be one-element float32 + device tensors. Example host values are + ``[math.log2(math.e) / math.sqrt(d)]`` for the softmax scale and + ``[1.0]`` for FP16 output. FP8 callers may fold Q/K/V and output + quantization scales into these runtime tensors. + """ + cfg = self.cfg + if cutlass.const_expr( + cfg.use_paged_kv + and ( + cum_seqlen_q is None + or block_tables is None + or seq_lens_kv is None + or max_seqlen_q is None + or max_seqlen_k is None + ) + ): + raise ValueError( + "paged context requires qo_indptr, block_tables, seq_lens_kv, " + "max_seqlen_q, and max_seqlen_k" + ) + if cutlass.const_expr(cfg.has_variable_window): + if cutlass.const_expr( + variable_window_token_starts is None + or variable_window_token_ends is None + or variable_window_cta_starts is None + ): + raise ValueError("VariableWindow requires start and end tensors") + + # Create TMA descriptors. Both query-paired and head-paired modes use + # the same logical Q/K/V/O tensor-map boxes after FmhaConfig lowers the + # mode-specific CTA shape. Fixed-shape launches include the batch rank; + # varlen launches drop it because the flattened tensors are indexed via + # cum_seqlen metadata. + tma_qkv_swizzle = cuda.TensorMapSwizzle.s128b + # Paged K/V issues one 128-byte tensor-map fragment per physical page. + # Promote that exact fragment width so the Q-head tiles sharing each + # logical K/V tile reuse the same cache line, matching the reference + # tensor-map descriptor policy. + tma_kv_l2_promotion = ( + cuda.TensorMapL2Promotion.l2_128b + if cutlass.const_expr(cfg.use_paged_kv) + else cuda.TensorMapL2Promotion.none + ) + if cutlass.const_expr(cfg.head_paired): + inner_dim_size = cfg.qk_mma_tiler[2] * cfg.q_dtype.width // 8 + tma_qkv_swizzle = cuda.TensorMapSwizzle.none + if cutlass.const_expr(inner_dim_size % 128 == 0): + tma_qkv_swizzle = cuda.TensorMapSwizzle.s128b + tma_kv_l2_promotion = cuda.TensorMapL2Promotion.l2_128b + elif cutlass.const_expr(inner_dim_size == 64): + tma_qkv_swizzle = cuda.TensorMapSwizzle.s64b + tma_kv_l2_promotion = cuda.TensorMapL2Promotion.l2_64b + elif cutlass.const_expr(inner_dim_size == 32): + tma_qkv_swizzle = cuda.TensorMapSwizzle.s32b + + output_inner_dim_size = cfg.tma_copy_o_granu_inner * cfg.o_dtype.width // 8 + tma_o_swizzle = cuda.TensorMapSwizzle.none + if cutlass.const_expr(output_inner_dim_size % 128 == 0): + tma_o_swizzle = cuda.TensorMapSwizzle.s128b + elif cutlass.const_expr(output_inner_dim_size == 64): + tma_o_swizzle = cuda.TensorMapSwizzle.s64b + elif cutlass.const_expr(output_inner_dim_size == 32): + tma_o_swizzle = cuda.TensorMapSwizzle.s32b + + q_box_dims = (1, cfg.qk_mma_tiler[0], 1, cfg.tma_copy_q_granu_inner) + kv_box_dims = (1, cfg.qk_mma_tiler[1], 1, cfg.tma_copy_kv_granu_inner) + v_box_dims = ( + 1, + cfg.pv_mma_tiler[2], + 1, + cfg.pv_mma_tiler[1] // cfg.tma_copy_qkv_iters, + ) + o_box_dims = (1, cfg.epi_tile[0], 1, cfg.tma_copy_o_granu_inner) + stride_order = (3, 2, 1, 0) + # K/V descriptors track the K/V tensor rank, which paged-KV pins to 4 + # (page pool) regardless of var-len. Q/O follow the var-len decision. + kv_stride_order = stride_order + if cutlass.const_expr(cum_seqlen_q is not None): + q_box_dims = (cfg.qk_mma_tiler[0], 1, cfg.tma_copy_q_granu_inner) + o_box_dims = (cfg.epi_tile[0], 1, cfg.tma_copy_o_granu_inner) + stride_order = (2, 1, 0) + if cutlass.const_expr(not cfg.use_paged_kv): + kv_box_dims = ( + cfg.qk_mma_tiler[1], + 1, + cfg.tma_copy_kv_granu_inner, + ) + v_box_dims = ( + cfg.pv_mma_tiler[2], + 1, + cfg.pv_mma_tiler[1] // cfg.tma_copy_qkv_iters, + ) + kv_stride_order = stride_order + if cutlass.const_expr(cum_seqlen_q is not None): + # Q/O use the packed sequence axis as a ragged TMA dimension. + # Dense TMA bounds only see the full sum_seqlen tensor and cannot + # prevent a partial final tile from crossing into the next packed + # sequence. Paged K/V keep native page-pool descriptors. Invalid K + # scores are masked before softmax; invalid V rows are overwritten + # in SMEM after their TMA stage completes and before PV MMA. + tma_q_desc = create_tensor_map_ragged_from_tensor( + q_cute, + box_dims=q_box_dims, + ragged_dim=0, + stride_order=stride_order, + swizzle=tma_qkv_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.none, + ) + else: + tma_q_desc = cuda.create_tensor_map_tiled_from_view( + q_cute, + box_dims=q_box_dims, + stride_order=stride_order, + swizzle=tma_qkv_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.none, + ) + + if cutlass.const_expr(cfg.use_paged_kv): + # Paged-KV path: K/V are compact pool tensors with shape + # (total_pages, h_kv, num_tokens_per_page, d). The TMA box covers + # one page × one d-fragment; the loader stitches pages and d-halves + # together via per-fragment coordinates. Keep both descriptors on + # the native rank-4 view. Synthetic ragged V maps with numeric-zero + # OOB fill have triggered mixed-specialization SM100 aborts; their + # special-NaN alternative is also invalid because tcgen05.mma + # propagates the NaN instead of treating it as a zero operand. + paged_kv_box_dims = ( + 1, + 1, + cfg.num_tokens_per_page, + cfg.tma_copy_kv_granu_inner, + ) + tma_k_desc = cuda.create_tensor_map_tiled_from_view( + k_cute, + box_dims=paged_kv_box_dims, + stride_order=kv_stride_order, + swizzle=tma_qkv_swizzle, + l2_promotion=tma_kv_l2_promotion, + ) + tma_v_desc = cuda.create_tensor_map_tiled_from_view( + v_cute, + box_dims=paged_kv_box_dims, + stride_order=kv_stride_order, + swizzle=tma_qkv_swizzle, + l2_promotion=tma_kv_l2_promotion, + ) + else: + tma_k_desc = cuda.create_tensor_map_tiled_from_view( + k_cute, + box_dims=kv_box_dims, + stride_order=kv_stride_order, + swizzle=tma_qkv_swizzle, + l2_promotion=tma_kv_l2_promotion, + ) + tma_v_desc = cuda.create_tensor_map_tiled_from_view( + v_cute, + box_dims=v_box_dims, + stride_order=kv_stride_order, + swizzle=tma_qkv_swizzle, + l2_promotion=tma_kv_l2_promotion, + ) + + if cutlass.const_expr(cum_seqlen_q is not None): + tma_o_desc = create_tensor_map_ragged_from_tensor( + o_cute, + box_dims=o_box_dims, + ragged_dim=0, + stride_order=stride_order, + swizzle=tma_o_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.none, + ) + else: + tma_o_desc = cuda.create_tensor_map_tiled_from_view( + o_cute, + box_dims=o_box_dims, + stride_order=stride_order, + swizzle=tma_o_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.none, + ) + + # Compute tile scheduler and grid + if cutlass.const_expr(cum_seqlen_q is None): + b = o_cute.shape[0] + s_q = o_cute.shape[1] + h_q = o_cute.shape[2] + if cutlass.const_expr(cfg.use_paged_kv): + # k_cute is the page pool — its sequence axis is per-page, + # not the logical max_seq_len_kv. Caller must pass it in. + s_k = max_seqlen_k + else: + s_k = k_cute.shape[1] + else: + b = cute.size(cum_seqlen_q) - 1 + s_q = max_seqlen_q + h_q = o_cute.shape[1] + s_k = max_seqlen_k + + num_seq_tiles = cute.ceil_div(s_q, cfg.cta_tiler[0]) + num_kv_tiles = cute.ceil_div(s_k, cfg.kv_tile_n) + num_head_tiles = h_q // cfg.work_tile_q_heads + q_offset = Int32(0) if cutlass.const_expr(cfg.head_paired) else 0 + if cutlass.const_expr(self.is_causal): + q_offset = s_k - s_q + + # Tile scheduling order: + # + # Causal defaults to seq_head_batch to keep adjacent Q sequence tiles + # close for K/V locality. Paired FP8 uses head_batch_seq, and paired + # causal CLC additionally reverses its sequence axis so the dynamic + # queue retires the heaviest causal tiles first. An explicit + # balance_causal_workload request uses that same reversed order. + # + # Dense GQA uses head_batch_seq so Q-head groups that share one K/V + # head stay adjacent. Dense MHA retains seq_head_batch because it has + # no cross-head K/V reuse. + if cutlass.const_expr(cfg.uses_head_batch_seq_tile_order): + problem_shape = (num_head_tiles, b, num_seq_tiles) + else: + problem_shape = (num_seq_tiles, num_head_tiles, b) + + if cutlass.const_expr(self.is_clc_dynamic): + tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + problem_shape, + cfg.cluster_shape_mn + (1,), + ) + grid = tile_sched_params.get_grid_shape() + elif cutlass.const_expr(self.is_persistent): + tile_sched_params = utils.PersistentTileSchedulerParams( + problem_shape, + cfg.cluster_shape_mn + (1,), + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, + max_active_clusters, + ) + else: + # Non-persistent launch: one CTA per tile. No tile scheduler params + # are passed into the task manager, so tasks use the hardware tile + # coordinate directly instead of a WorkQueue. + tile_sched_params = None + grid = problem_shape + + block_size = cfg.block_warps * 32 + + # Paged-KV side-channel data: one fixed row-strided page table plus + # per-batch K/V lengths. None in the contiguous path; the kernel + # branches on ``cfg.use_paged_kv`` so passing None is safe. + block_tables_iter = ( + block_tables.iterator + if cutlass.const_expr(block_tables is not None) + else None + ) + block_table_row_stride = ( + Int32(block_tables.stride[0]) + if cutlass.const_expr(block_tables is not None) + else Int32(0) + ) + seq_lens_kv_iter = ( + seq_lens_kv.iterator + if cutlass.const_expr(seq_lens_kv is not None) + else None + ) + + self.kernel( + tma_q_desc, + tma_k_desc, + tma_v_desc, + tma_o_desc, + tile_sched_params, + scale_softmax_log2, + output_scale, + num_kv_tiles, + num_seq_tiles, + q_offset, + cum_seqlen_q, + cum_seqlen_k, + block_tables_iter, + block_table_row_stride, + seq_lens_kv_iter, + Int32(s_k), + variable_window_token_starts, + variable_window_token_ends, + variable_window_cta_starts, + Int32(s_q), + self.is_persistent, + self.is_clc_dynamic, + ).launch( + grid=grid, + block=[block_size, 1, 1], + cluster=cfg.cluster_shape_mn + (1,), + stream=stream, + min_blocks_per_mp=1, + ) + + # --------------------------------------------------------------------------- + # Device kernel + # --------------------------------------------------------------------------- + @cute.kernel + def kernel( + self, + tma_q_desc: cutlass.GridConstant[cuda.TensorMap], + tma_k_desc: cutlass.GridConstant[cuda.TensorMap], + tma_v_desc: cutlass.GridConstant[cuda.TensorMap], + tma_o_desc: cutlass.GridConstant[cuda.TensorMap], + tile_sched_params: ( + utils.PersistentTileSchedulerParams + | utils.ClcDynamicPersistentTileSchedulerParams + | None + ), + scale_softmax_log2: cute.Tensor, + output_scale: cute.Tensor, + num_kv_tiles: Int32, + num_seq_tiles: Int32, + q_offset: Int32, + cum_seqlen_q: cute.Tensor | None, + cum_seqlen_k: cute.Tensor | None, + block_tables: cute.Pointer | None, + block_table_row_stride: Int32, + seq_lens_kv: cute.Pointer | None, + max_seq_len_kv: Int32 | None, + variable_window_token_starts: cute.Tensor | None, + variable_window_token_ends: cute.Tensor | None, + variable_window_cta_starts: cute.Tensor | None, + variable_window_q_stride: Int32, + is_persistent: cutlass.Constexpr[bool] = True, + is_clc_dynamic: cutlass.Constexpr[bool] = False, + ) -> None: + """Warp-specialised persistent FMHA TS kernel. + + ``scale_softmax_log2`` and ``output_scale`` are one-element float32 + device tensors passed through to the TS resources. The resources load + element 0 in auxiliary setup before the K/V loop, so the hot loop uses + the cached task-local values rather than reloading global memory. + + Structure: + 1. Prefetch TMA descriptors + 2. Build TaskManager (resources + tasks + dependency graph) + 3. setup_resources_and_tasks() — unified SMEM alloc + pipeline barriers + 4. Derive infrastructure ptrs (tmem_ptr, dealloc mbar) from SMEM block + 5. Init tmem dealloc barrier + fence + cluster sync + 6. TMEM allocation + immediate permit relinquish (MMA warp) + named barrier sync + 7. task_manager.run() — persistent execution + 8. TMEM deallocation + """ + cfg = self.cfg + cfg.num_seq_tiles = num_seq_tiles + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # 1. Prefetch TMA descriptors on load warp + if warp_idx == cfg.load_warp_id: + prims.prefetch_tensormap(tma_q_desc.get_ptr()) + prims.prefetch_tensormap(tma_k_desc.get_ptr()) + prims.prefetch_tensormap(tma_v_desc.get_ptr()) + prims.prefetch_tensormap(tma_o_desc.get_ptr()) + + # 2. CLC dynamic: the response buffer is declared by the builder and + # bound from the unified task-manager SMEM allocation below. + clc_response_ptr = None + + # 3. Build TaskManager (infrastructure slots via SmemAllocator) + ( + task_manager, + tmem_resources, + tmem_ptr_alloc, + dealloc_mbar_alloc, + work_queue, + clc_response_alloc, + ) = build_fmha_task_manager( + cfg=cfg, + tile_sched_params=tile_sched_params, + tma_q_desc=tma_q_desc.get_ptr(), + tma_k_desc=tma_k_desc.get_ptr(), + tma_v_desc=tma_v_desc.get_ptr(), + tma_o_desc=tma_o_desc.get_ptr(), + cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + g_block_tables=block_tables, + block_table_row_stride=block_table_row_stride, + g_seq_lens_kv=seq_lens_kv, + max_seq_len_kv=max_seq_len_kv, + num_kv_tiles=num_kv_tiles, + variable_window_token_starts=variable_window_token_starts, + variable_window_token_ends=variable_window_token_ends, + variable_window_cta_starts=variable_window_cta_starts, + variable_window_q_stride=variable_window_q_stride, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + q_offset=q_offset, + is_persistent=is_persistent, + is_clc_dynamic=is_clc_dynamic, + clc_response_ptr=clc_response_ptr, + exhaustive_deadlock_race_check=self.exhaustive_deadlock_race_check, + ) + + # Bind the CLC scheduler to its exact suballocation before WorkQueue + # create() materializes the concrete scheduler during normal setup. + if cutlass.const_expr(is_clc_dynamic): + assert work_queue is not None + assert clc_response_alloc is not None + smem_allocator = task_manager.smem_allocator + assert smem_allocator is not None + smem_allocator.allocate() + clc_response = smem_allocator.get(clc_response_alloc) + clc_response_ptr = cute.make_ptr( + cutlass.Int128, + clc_response.data_ptr(), + mem_space=cutlass.AddressSpace.smem, + assumed_align=16, + ) + work_queue.tile_scheduler_config = ( + TileSchedulerConfig.create_clc_dynamic_persistent_tile_scheduler_params( + tile_scheduler_params=tile_sched_params, + response_ptr=clc_response_ptr, + ) + ) + + # 4. Initialize all pipeline barriers + allocate unified SMEM + task_manager.setup_resources_and_tasks() + + # Derive infrastructure pointers from the unified SMEM block. + # tmem_ptr_i32 for ResourceContext is auto-populated by TaskManager + # in setup_resources_and_tasks() via SmemAllocator.tmem_ptr_alloc. + smem_allocator = task_manager.smem_allocator + tmem_ptr_i32 = smem_allocator.get(tmem_ptr_alloc) + tmem_dealloc_mbar = smem_allocator.get(dealloc_mbar_alloc) + + # 5. Initialize tmem dealloc barrier + num_tmem_consumer_threads = cute.arch.WARP_SIZE * ( + len(cfg.softmax0_warp_ids) + + len(cfg.softmax1_warp_ids) + + len(cfg.correction_warp_ids) + ) + if warp_idx == cfg.empty_warp_id: + if prims.elect_sync(): + prims.mbarrier_init(tmem_dealloc_mbar, num_tmem_consumer_threads) + + # Fence barrier inits from setup_resources_and_tasks() — pipeline + # barriers need fencing before first use. + prims.fence_mbarrier_init() + + # Cluster sync before TMEM allocation / task execution + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + + # 6. TMEM allocation (MMA warp) + if warp_idx == cfg.mma_warp_id: + tmem_alloc_cols = Int32(cfg.tmem_alloc_cols) + prims.tcgen05_alloc(tmem_ptr_i32, tmem_alloc_cols) + prims.tcgen05_relinquish_alloc_permit() + + # All-warp barrier to ensure TMEM allocation is visible. + # Previously only correction+MMA warps synced; expanded to all + # warps so we can cache tmem_ptr_i32 once and avoid repeated + # ld.shared in resource work methods (SMEM is volatile to LLVM). + prims.barrier_cta_sync( + cfg.tmem_bar_id, + thread_count=cfg.block_warps * cute.arch.WARP_SIZE, + ) + + # Cache TMEM address: one ld.shared per warp. This replaces + # ~6 ld.shared calls spread across resource work methods (each + # cascading into shr/and/shl address arithmetic that the compiler + # cannot hoist because SMEM is volatile). + tmem_addr_cached = tmem_ptr_i32.load() + for r in tmem_resources: + r.tmem_addr_cached = tmem_addr_cached + + # 7. Execute all tasks via the TS scheduler + task_manager.run() + + # 8. TMEM deallocation + is_softmax_warp = ( + warp_idx >= cfg.softmax0_warp_ids[0] + and warp_idx <= cfg.softmax0_warp_ids[-1] + ) + if cutlass.const_expr(len(cfg.softmax1_warp_ids) > 0): + is_softmax_warp = is_softmax_warp or ( + warp_idx >= cfg.softmax1_warp_ids[0] + and warp_idx <= cfg.softmax1_warp_ids[-1] + ) + is_correction_warp = ( + warp_idx >= cfg.correction_warp_ids[0] + and warp_idx <= cfg.correction_warp_ids[-1] + ) + if is_softmax_warp or is_correction_warp: + prims.mbarrier_arrive(tmem_dealloc_mbar) + + if warp_idx == cfg.mma_warp_id: + while not prims.mbarrier_try_wait_parity(tmem_dealloc_mbar, 0): + pass + tmem_alloc_cols = Int32(cfg.tmem_alloc_cols) + tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) + prims.tcgen05_dealloc(tmem_ptr, tmem_alloc_cols) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py new file mode 100644 index 000000000000..7c3c1417db23 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py @@ -0,0 +1,4971 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resource definitions for the TS FMHA kernel. + +Maps the FMHA data-flow onto ``MemoryResource`` subclasses: + +Each resource owns the work attached to one live buffer: producer methods fill +the buffer, consumer methods drain it, and the pipeline state records when the +next task may use the data. Task files only order these resource work calls. + +Schedule phase terms follow TS schedule-builder naming. HEAD is the one-time +schedule before the repeated K/V tile loop, LOOP is the repeated K/V tile body, +and TAIL is the one-time cleanup and drain after LOOP exits. + +SMEM resources (TMA pipelines) +------------------------------ +- SmemQResource : SMEM Q buffer, TmaUmma pipeline. Load -> MMA. +- SmemKVResource : SMEM KV buffer, TmaUmma pipeline. Load -> MMA. The D256 + depth is derived from the public SM100 SMEM capacity. +- SmemOResource : One SMEM O buffer per Q/O instance, AsyncAsync pipeline. + Correction -> Epilogue. + +TMEM resources (split from former TmemComputeResource) +------------------------------------------------------ +- TmemSPResource : S/P buffer, UmmaAsync. D256 uses a two-stage S/P ring and + an independent TmemPResource readiness handoff. + MMA writes S (Q*K scores). Softmax reads S, computes P, + and writes P back into the same slot for BMM2. + Self-edge in dependency graph enables ping-pong validation. +- TmemStatsResource : Correction statistics, AsyncAsync pipeline. + Softmax writes [old_max, new_max, row_sum], Correction reads. +- TmemOResource : O accumulation, UmmaAsync pipeline. + MMA writes P*V -> O. Correction waits for O, rescales it + in-place, then releases the stage for the interleaved MMA. + +Sequencing resources +-------------------- +- S0S1SequenceResource : PipelineAsync (1 stage), Softmax0 → Softmax1. + Ensures S0 finishes P store to TMEM before S1 starts + P computation. Prevents TMEM write contention. + Operations are inlined in TmemSPResource.consumer_work. + +GMEM resources (no pipeline) +----------------------------- +- GmemQKVResource : TMA descriptors + per-tile coordinate resolution. +- GmemOResource : TMA descriptor for O stores. +""" + +import math +from dataclasses import dataclass, field, replace +from typing import Any, Optional, TypeAlias + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from ..tensor_map import transform_ragged_coords + +from cutlass.experimental.task_scheduling.enums import WorkAttr +from ..stage import FmhaStage +from cutlass.experimental.task_scheduling.memory import SmemAllocation, TmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + PipelineConfig, + StageInfo, + TaskLocalVariable, +) +from cutlass.experimental.task_scheduling.resources import consumer_work, producer_work + + +_SUPPORTED_CONTEXT_PAGE_SIZES = (16, 32, 64, 128) +from cutlass.pipeline import PipelineAsync, PipelineState +from cutlass.cutlass_dsl import Boolean, Constexpr, dsl_user_op, if_generate + +from ..placeholder_helpers import _placeholder_smem_array, _placeholder_tmem_ptr +from .helpers import ( + bottom_right_window_left_bound, + bottom_right_window_tile_start, + freeze_smem_descriptor, + variable_window_cta_min_start, +) +from cutlass.experimental import primitives as prims + +SmemDescOffsets: TypeAlias = tuple[int, int] +TmemAddr: TypeAlias = int | Int32 +TmemPtr: TypeAlias = cutlass.Array + +SoftmaxScalar: TypeAlias = Float32 +SoftmaxChunk: TypeAlias = cutlass.Vector +SoftmaxChunks: TypeAlias = list[SoftmaxChunk] +SoftmaxRowSumContribution: TypeAlias = SoftmaxChunks | SoftmaxScalar + +# Trace-time storage for TmemSP s_data vectors (S/P chunks). +# Stored at module level (not on self) to avoid adding a non-dynamic-expression +# field to the dataclass, which breaks the framework's scf.if handling. +_tmem_sp_sdata: dict[int, list] = {} + + +@cute.jit +def _bmsk_clamp(start: Int32, width: Int32) -> Int32: + """Create a contiguous 32-bit mask with clamped bounds.""" + return cute.arch.inline_ptx( + "bmsk.clamp.b32 {$w0}, {$r0}, {$r1};", + write_only_types=[Int32], + read_only_args=[start, width], + ) + + +@cute.jit +def _mask_score_quad( + valid_bits: Int32, + score0: Float32, + score1: Float32, + score2: Float32, + score3: Float32, +) -> tuple[Float32, Float32, Float32, Float32]: + """Expand four bitmap bits with setp and replace invalid scores.""" + return cute.arch.inline_ptx( + """ + { + .reg .pred valid<4>; + .reg .b32 bit; + mov.b32 {$w0}, {$r1}; + mov.b32 {$w1}, {$r2}; + mov.b32 {$w2}, {$r3}; + mov.b32 {$w3}, {$r4}; + and.b32 bit, {$r0}, 0x1; + setp.ne.u32 valid0, bit, 0; + and.b32 bit, {$r0}, 0x2; + setp.ne.u32 valid1, bit, 0; + and.b32 bit, {$r0}, 0x4; + setp.ne.u32 valid2, bit, 0; + and.b32 bit, {$r0}, 0x8; + setp.ne.u32 valid3, bit, 0; + @!valid0 mov.b32 {$w0}, 0xff800000; + @!valid1 mov.b32 {$w1}, 0xff800000; + @!valid2 mov.b32 {$w2}, 0xff800000; + @!valid3 mov.b32 {$w3}, 0xff800000; + } + """, + write_only_types=[Float32, Float32, Float32, Float32], + read_only_args=[valid_bits, score0, score1, score2, score3], + ) + + +@cute.jit +def _pack_float4_to_fp8_e4m3( + v0: Float32, + v1: Float32, + v2: Float32, + v3: Float32, +) -> Int32: + """Pack four FP32 values with the public packed-conversion primitive.""" + lo = prims.cvt_packfloat_f32( + v1, + v0, + Int32(0), + prims.CVTPackFloat.E4M3X2, + rnd=prims.FPRoundingMode.RN, + sat=prims.SaturationModeKind.SATFINITE, + extract_hi=False, + ) + return prims.cvt_packfloat_f32( + v3, + v2, + lo, + prims.CVTPackFloat.E4M3X2, + rnd=prims.FPRoundingMode.RN, + sat=prims.SaturationModeKind.SATFINITE, + extract_hi=True, + ) + + +def _placeholder_softmax_chunks(cfg: Any) -> SoftmaxChunks: + """Build zero P chunks with the same structure as runtime softmax chunks.""" + try: + tmem_x = cfg.tmem_x_load_s + num_chunks = cfg.qk_mma_tiler[1] // tmem_x + chunks = [] + for _ in range(num_chunks): + zeros = tuple(cfg.qk_acc_dtype(0.0) for _ in range(tmem_x)) + chunks.append(cutlass.Vector.from_elements(zeros, cfg.qk_acc_dtype)) + return chunks + except RuntimeError: + return [] + + +# --------------------------------------------------------------------------- +# FmhaConfig -- kernel-wide configuration +# --------------------------------------------------------------------------- + + +@dataclass +class FmhaConfig: + """Compile-time and runtime configuration for the FMHA kernel. + + Mirrors the attributes from BlackwellFusedMultiHeadAttentionForward.__init__ + and _setup_attributes, collected into a single portable dataclass. + + All fields are marked Constexpr so that tree_flatten does not try to + recursively extract MLIR values from dtype classes or plain Python ints/tuples. + """ + + # Data types + q_dtype: type | None = None + k_dtype: type | None = None + v_dtype: type | None = None + o_dtype: type | None = None + qk_acc_dtype: type | None = None + pv_acc_dtype: type | None = None + + # Tile shapes + qk_mma_tiler: tuple[int, int, int] = (128, 128, 64) + pv_mma_tiler: tuple[int, int, int] = (128, 64, 128) + epi_tile: tuple[int, int] = (128, 64) + # Number of interleaved Q/KV/O instances per CTA: two selects the paired + # schedule, while one selects the single-instance schedule used for D>128. + num_qkv_instances: int = 2 + + # Pipeline stages + q_stage: int = 2 + kv_stage: int = 3 + mma_softmax_stage: int = 1 + # Use the two-stage loop-carried S/P schedule and an independent P-ready + # handoff, allowing QK(i+1) and PV(i) to operate on opposite S/P stages. + has_tmem_p_pipeline: bool = False + stats_via_smem: bool = False + stage_scoped_tmem_stats: bool = False + softmax_corr_stage: int = 1 + mma_corr_stage: int = 2 + + # TMA copy granularity + tma_copy_qkv_iters: int = 1 + tma_copy_q_granu_inner: int = 128 + tma_copy_q_elements: int = 0 + tma_copy_q_granu_elems: int = 0 + tma_copy_q_bytes: int = 0 + tma_copy_kv_granu_inner: int = 128 + tma_copy_kv_elements: int = 0 + tma_copy_kv_granu_elems: int = 0 + tma_copy_kv_bytes: int = 0 + tma_copy_o_iters: int = 1 + tma_copy_o_granu_inner: int = 0 + tma_copy_o_elements: int = 0 + tma_copy_o_granu_elems: int = 0 + q_tile_m: int = 128 + kv_tile_n: int = 128 + + # Warp assignments + softmax0_warp_ids: tuple[int, int, int, int] = (0, 1, 2, 3) + softmax1_warp_ids: tuple[int, int, int, int] = (4, 5, 6, 7) + correction_warp_ids: tuple[int, int, int, int] = (8, 9, 10, 11) + mma_warp_id: int = 12 + load_warp_id: int = 13 + epilogue_warp_id: int = 14 + empty_warp_id: int = 15 + + # Register budgets + num_regs_softmax: int = 192 + num_regs_correction: int = 96 + num_regs_other: int = 32 + + # TMEM layout + tmem_alloc_cols: int = 512 + tmem_stats_cols: int = 4 + tmem_s0_offset: int = 0 + tmem_s1_offset: int = 128 + tmem_o0_offset: int = 256 + tmem_o1_offset: int = 384 + tmem_p0_offset: int = 32 + tmem_p1_offset: int = 160 + tmem_vec0_offset: int = 0 + tmem_vec1_offset: int = 128 + + # SMEM shapes (set during __init__ of FmhaTs) + sO_stage_elements: int = 0 + sQ_shape: tuple[int, int] = (2, 0) + sK_shape: tuple[int, int] = (3, 0) + + # Misc + buffer_align_bytes: int = 1024 + tmem_bar_id: int = 2 + cluster_shape_mn: tuple[int, int] = (1, 1) + block_warps: int = 16 + + # GQA: head ratio h_q // h_kv (1 = MHA, >1 = GQA) + h_r: int = 1 + + # Causal masking: when True, mask out positions where k_idx > q_idx + is_causal: bool = False + # Explicit packed-Q inclusive [start, end] bounds replace static masks. + has_variable_window: bool = False + # Causal balancing uses head_batch_seq logical tile order and reverses Q + # sequence tiles. + balance_causal_workload: bool = False + num_seq_tiles: int | Int32 = 0 + # Skip correction optimization: when True, skip rescale if old_max == new_max + enable_skip_correction: bool = True + + # Variable sequence length mode stores Q/K/V/O as flattened + # [sum_seqlen, head, dim] tensors and uses cum_seqlen_* for per-batch + # sequence offsets. + has_varlen: bool = False + # Uniform packed plans retain the ragged tensor-map ABI but derive their + # cumulative offsets arithmetically, avoiding dependent indptr loads on + # every persistent work tile. + has_uniform_varlen: bool = False + uniform_seq_len_q: int = 0 + uniform_seq_len_k: int = 0 + + # When true, map each work tile to two Q heads sharing one K/V head. + # This enables the grouped-query/sliding-window context flavor while + # reusing the unified FMHA context implementation. + head_paired: bool = False + + # Concrete kernel policy derived from paired geometry and V dtype. The + # resource consumes this flag directly so the row-sum algorithm and the + # register budget selected by FmhaTs cannot diverge. + enable_early_tile_sum: bool = False + + seq_tile_n: int = 128 + tmem_x_load_s: int = 32 + # Causal S_q < S_kv shifts Q rows right by q_offset = S_kv - S_q. + # This flag selects the shifted causal mask; there is no second causal mode. + has_q_offset: bool = False + # Fixed causal attention with exactly one K/V tile does not need the + # synthetic peer0 tail slot used by the general query-paired schedule. + causal_single_kv_tile: bool = False + window_size_left: int = 0 + # Number of valid K/V rows in the final fixed-length dense tile. Zero + # means that the K/V extent is tile-aligned (or that this specialization + # does not use the fixed dense-tail mask). + fixed_dense_k_tail: int = 0 + # Packed-contiguous and paged dense attention normally mask scores past + # each request's logical K length. Plans with uniform, tile-aligned K + # lengths can compile that mask away because every K tile is fully valid. + packed_dense_k_mask: bool = True + + # ------------------------------------------------------------------ + # Paged KV cache (vLLM-style logical->physical page indirection) + # + # When use_paged_kv is True, K/V live in a fixed-size page pool + # [num_pages_in_pool, h_kv, num_tokens_per_page, d] and the kernel follows + # a fixed row-strided block table to resolve logical (b, s) -> physical page + # id at TMA-issue time. + # + # Staged D256 assigns page-offset prefetch to its empty/padding warp. + # Paired D128 reads page IDs directly from the page table in its load task. + # ------------------------------------------------------------------ + use_paged_kv: bool = False + # D256 uses a single Q/KV instance and can issue the final O TMA store + # from one correction warp after the four-warp correction group has + # staged O. This frees the standalone epilogue warp for scheduling. + fuse_epilogue_into_correction: bool = False + num_tokens_per_page: int = 32 + # Static upper bound derived from max_kv_len during kernel construction. + # Runtime active-page bounds come from seq_lens_kv. + max_num_pages_per_seq_kv: int = 1 + page_offsets_num_warps: int = 1 + # Selected internally from the staged topology, static page geometry, and + # exact SMEM capacity. It is derived during kernel construction and is not + # a public tuning input. + page_table_window_entries: int = 32 + + # Work-tile mapping for the two peer Q/O tiles handled by each CTA: + # query-paired maps peers to two sequence tiles in one Q head, while + # head-paired mode maps peers to two Q heads at one sequence tile. + @property + def single_qkv_instance(self) -> bool: + """Return whether one work tile carries a single Q/KV/O instance.""" + return self.num_qkv_instances == 1 + + @property + def uses_early_tile_sum(self) -> bool: + """Return whether paired M128 geometry supports early V-tile reduction.""" + return ( + not self.single_qkv_instance + and self.q_tile_m == 128 + and self.v_dtype + in ( + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + ) + ) + + @property + def uses_d256_fp8_softmax_cadence(self) -> bool: + """Return whether staged D256 FP8 uses interleaved softmax retirement.""" + return ( + self.single_qkv_instance + and self.has_tmem_p_pipeline + and self.stage_kv_by_head_dim + and self.qk_mma_tiler == (128, 128, 256) + and self.q_dtype == cutlass.Float8E4M3FN + and self.k_dtype == cutlass.Float8E4M3FN + and self.v_dtype == cutlass.Float8E4M3FN + ) + + @property + def uses_d128_fp8_softmax_cadence(self) -> bool: + """Return whether paired D128 FP8 uses interleaved softmax retirement.""" + return ( + not self.stage_kv_by_head_dim + and self.enable_early_tile_sum + and self.q_dtype == cutlass.Float8E4M3FN + and self.k_dtype == cutlass.Float8E4M3FN + and self.v_dtype == cutlass.Float8E4M3FN + ) + + @property + def reuses_page_table_windows(self) -> bool: + """Whether dense paged-KV admits structural page-ID windows.""" + if ( + not self.use_paged_kv + or self.is_causal + or self.num_tokens_per_page <= 0 + or self.kv_tile_n % self.num_tokens_per_page != 0 + ): + return False + pages_per_tile = self.kv_tile_n // self.num_tokens_per_page + window_entries = self.page_table_window_entries + if ( + pages_per_tile <= 0 + or pages_per_tile > window_entries + or window_entries % pages_per_tile != 0 + ): + return False + window_period = window_entries // pages_per_tile + if self.single_qkv_instance and self.has_tmem_p_pipeline and window_period < 3: + # The K-ahead/V-delayed HEAD requires distinct K0, K1, and tail + # positions. Smaller periods retain the ordinary per-tile path. + return False + static_num_kv_tiles = ( + self.max_num_pages_per_seq_kv + pages_per_tile - 1 + ) // pages_per_tile + return ( + static_num_kv_tiles >= window_period + and static_num_kv_tiles % window_period == 0 + ) + + @property + def stages_page_offsets_in_smem(self) -> bool: + """Whether a dedicated warp stages page IDs for the load warp. + + Staged D256 uses the coalesced SMEM page-window path. Paired D128 loads + page IDs directly in its K/V producer. + """ + return self.use_paged_kv and self.single_qkv_instance + + @property + def needs_paged_v_tail_clear(self) -> bool: + """Whether paged V tiles can contain request-invalid rows. + + Only exact-full causal grids omit the clear. Requiring complete Q + work tiles also excludes a padded final query-paired domain. + """ + if not self.use_paged_kv: + return False + q_work_tile_m = self.q_tile_m * self.work_tile_q_seq_tiles + return not ( + self.has_uniform_varlen + and self.is_causal + and not self.has_q_offset + and self.uniform_seq_len_k % self.kv_tile_n == 0 + and self.uniform_seq_len_q % q_work_tile_m == 0 + ) + + @property + def page_table_window_candidate_entries(self) -> int: + """Return the widest page-ID window admitted by static topology. + + A split-D K/V schedule consumes two head-dimension stages for each + logical tile. When the static domain can cover a complete window, + let each producer lane fetch one ID per D stage so one page-window + handoff spans both stages. Short domains admit only the natural + one-warp window. The kernel's capacity pass makes the final selection. + """ + natural_entries = cute.arch.WARP_SIZE + if ( + not self.use_paged_kv + or self.is_causal + or self.num_tokens_per_page <= 0 + or not (self.single_qkv_instance and self.has_tmem_p_pipeline) + ): + return natural_entries + pages_per_tile = self.kv_tile_n // self.num_tokens_per_page + staged_entries = natural_entries * self.num_head_dim_stages_k + staged_period = staged_entries // pages_per_tile + static_num_kv_tiles = ( + self.max_num_pages_per_seq_kv + pages_per_tile - 1 + ) // pages_per_tile + if ( + staged_entries % pages_per_tile == 0 + and static_num_kv_tiles >= staged_period + and static_num_kv_tiles % staged_period == 0 + ): + return staged_entries + return natural_entries + + @property + def page_offset_pipeline_stage_counts(self) -> tuple[int, ...]: + """Return the physical page-ID ring depths for this topology.""" + if not self.use_paged_kv or not self.single_qkv_instance: + return () + # The staged schedule holds one credit for every K/V head-dimension + # slice plus one K-ahead boundary credit. A reused page-table window + # needs independent K-ahead and V-delayed rings; the ordinary path + # shares the same total number of credits. Paired D128 loads page IDs + # directly in its load task and therefore has no physical page ring. + k_stages = self.num_head_dim_stages_k + 1 + v_stages = self.num_head_dim_stages_v + if self.reuses_page_table_windows: + return (k_stages, v_stages) + return (k_stages + v_stages,) + + @property + def cta_tiler(self) -> tuple[int, int, int]: + """Derive the CTA tile from the MMA tile and work-tile mapping.""" + if self.single_qkv_instance or self.head_paired: + return self.qk_mma_tiler + return ( + self.num_qkv_instances * self.qk_mma_tiler[0], + self.qk_mma_tiler[1], + self.qk_mma_tiler[2], + ) + + @property + def uses_causal_reversed_head_batch_seq_tile_order(self) -> bool: + """Return whether causal head_batch_seq tiles reverse Q sequence order.""" + return self.is_causal and self.balance_causal_workload + + @property + def uses_paired_fp8_head_batch_seq_tile_order(self) -> bool: + """Return whether paired FP8 benefits from head-local tile order. + + Causal work uses this order for load balancing. Dense GQA uses it to + keep Q-head groups that share the same K/V head adjacent. Dense MHA + has no cross-head K/V reuse and retains its sequence-local order. + """ + return ( + (self.is_causal or self.h_r > 1) + and not self.single_qkv_instance + and self.q_dtype is not None + and self.k_dtype is not None + and self.v_dtype is not None + and self.q_dtype.width == 8 + and self.k_dtype.width == 8 + and self.v_dtype.width == 8 + ) + + @property + def uses_head_batch_seq_tile_order(self) -> bool: + """Return whether work tiles use head_batch_seq coordinates.""" + return self.uses_paired_fp8_head_batch_seq_tile_order or ( + self.is_causal and self.balance_causal_workload + ) + + @property + def work_tile_coord_indices(self) -> tuple[int, int, int]: + """Return work-tile indices for logical ``(seq, head, batch)``.""" + if self.uses_head_batch_seq_tile_order: + return 2, 0, 1 + return 0, 1, 2 + + @property + def pv_p_scale(self) -> float: + """Return the P scale applied before PV MMA.""" + if self.v_dtype is not None and self.v_dtype.width == 8: + # FP8 E4M3 has max finite magnitude 448; scaling P to that range + # before PV MMA preserves dynamic range. + return 448.0 + # Non-FP8 V uses P directly, so the PV-side P scale is identity. + return 1.0 + + @property + def pv_p_scale_log2(self) -> float: + """Return log2(P scale) for folding into exp2 softmax P.""" + return math.log2(self.pv_p_scale) + + @property + def work_tile_q_heads(self) -> int: + """Return the number of Q heads represented by one work tile.""" + if self.single_qkv_instance: + return 1 + return 2 if self.head_paired else 1 + + @property + def work_tile_q_seq_tiles(self) -> int: + """Return the number of Q sequence tiles represented by one work tile.""" + if self.single_qkv_instance: + return 1 + return 1 if self.head_paired else 2 + + @property + def has_tile_aligned_uniform_q_offset(self) -> bool: + """Whether a uniform causal shift preserves K/V tile boundaries. + + Uniform packed plans retain fixed Q/K lengths under their replay + contract. When their bottom-right shift is an exact K/V-tile + multiple, every query tile's causal diagonal has the same placement + as the zero-offset schedule: query-paired peer 0 can use its explicit + diagonal/invalid-tail protocol, and all other diagonals remain in + TAIL. No LOOP iteration then needs a causal right mask. + """ + return ( + self.has_q_offset + and self.has_uniform_varlen + and self.uniform_seq_len_q % self.cta_tiler[0] == 0 + and (self.uniform_seq_len_k - self.uniform_seq_len_q) % self.kv_tile_n == 0 + ) + + @property + def peer_q_head_stride(self) -> int: + """Return the Q-head stride between the two peer Q/O tiles.""" + return 1 if self.head_paired and not self.single_qkv_instance else 0 + + @property + def peer_q_seq_tile_stride(self) -> int: + """Return the Q-sequence tile stride between the two peer Q/O tiles.""" + return 0 if self.head_paired or self.single_qkv_instance else 1 + + @property + def gmem_o_store_wait_after_write(self) -> bool: + """Return whether each O store must wait for the matching SMEM write.""" + return self.head_paired or self.stage_o_by_head_dim + + @property + def skip_causal_invalid_peer0(self) -> bool: + """Return whether query-paired causal peer0 may skip extra loop work.""" + if ( + not self.is_causal + or self.head_paired + or (self.has_q_offset and not self.has_tile_aligned_uniform_q_offset) + or self.single_qkv_instance + or self.causal_single_kv_tile + ): + return False + peer0_kv_tiles = (self.q_tile_m + self.kv_tile_n - 1) // self.kv_tile_n + paired_kv_tiles = (self.cta_tiler[0] + self.kv_tile_n - 1) // self.kv_tile_n + extra_peer1_kv_tiles = paired_kv_tiles - peer0_kv_tiles + if extra_peer1_kv_tiles > 2: + raise ValueError( + "query-paired causal scheduling supports peer1 at most two " + "K/V tiles ahead of peer0; got " + f"{extra_peer1_kv_tiles} extra K/V tiles" + ) + return extra_peer1_kv_tiles > 0 + + @property + def kv_tile_start_window_size_left(self) -> int: + """Return the left-window width used to compute the first K/V tile.""" + return self.window_size_left if self.head_paired else 0 + + +# --------------------------------------------------------------------------- +# S0S1SequenceResource -- S0-S1 sequence barrier (PipelineAsync, 1 stage) +# --------------------------------------------------------------------------- + + +@cute.jit +def _resolve_work_tile_coords( + cfg: Constexpr[FmhaConfig], + tile_idx: cute.Coord, +) -> tuple[Int32, Int32, Int32]: + """Return ``(seq, head, batch)`` for the configured tile order.""" + seq_idx, head_idx, batch_idx = cfg.work_tile_coord_indices + seq_coord = tile_idx[seq_idx] + head_coord = tile_idx[head_idx] + batch_coord = tile_idx[batch_idx] + if cutlass.const_expr(cfg.uses_causal_reversed_head_batch_seq_tile_order): + seq_coord = cfg.num_seq_tiles - seq_coord - Int32(1) + return seq_coord, head_coord, batch_coord + + +@dataclass(frozen=True) +class _StructuredWaitPipelineAsync(PipelineAsync): + """PipelineAsync with an explicit public-primitive retry loop.""" + + @cute.jit + def _retry_wait( + self, + sync_object: object, + state: PipelineState, + *, + loc: Any = None, + ip: Any = None, + ) -> None: + while not sync_object.try_wait( + state.index, + state.phase, + loc=loc, + ip=ip, + ): + pass + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc: Any = None, + ip: Any = None, + ) -> None: + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self._retry_wait(self.sync_object_empty, state, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + @dsl_user_op + def consumer_wait( + self, + state: PipelineState, + try_wait_token: Optional[Boolean] = None, + *, + loc: Any = None, + ip: Any = None, + ) -> None: + if_generate( + try_wait_token is None or try_wait_token == 0, + lambda: self._retry_wait(self.sync_object_full, state, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + +@dataclass(kw_only=True) +class S0S1SequenceResource(MemoryResource): + """Sequence barrier between Softmax0 (producer) and Softmax1 (consumer). + + Prevents both softmax groups from writing P to TMEM simultaneously. + Matches the handwritten FMHA kernel's ``s0_s1_sequence_mbar`` + PipelineAsync. + + Single resource instance shared across tasks: + - Softmax0's dst_resource (ProducerAcquire/Commit) + - Softmax1's src_resource (ConsumerWait/Release) + """ + + is_barrier: cutlass.Constexpr[bool] = True + + def create_pipeline(self, pipeline_config: PipelineConfig) -> object: + base = super().create_pipeline(pipeline_config) + assert isinstance(base, PipelineAsync) + return _StructuredWaitPipelineAsync( + base.sync_object_full, + base.sync_object_empty, + base.num_stages, + base.producer_mask, + base.consumer_mask, + ) + + +# --------------------------------------------------------------------------- +# TmemStatsDoneResource -- cross-tile stats-read notification +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class TmemStatsDoneResource(MemoryResource): + """Notification barrier: Correction signals after reading stats from TMEM. + + Prevents cross-tile aliasing race where next tile's QK→S UMMA can + overwrite TMEM columns overlapping TmemStats0/1 before correction reads them. + + Single resource instance shared across tasks: + - MMA's dst_resource (ProducerAcquire/Commit) + - Correction's src_resource (ConsumerWait/Release) + """ + + is_barrier: cutlass.Constexpr[bool] = True + + +# --------------------------------------------------------------------------- +# GmemQKVResource -- global memory Q/K/V source (no pipeline) +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class GmemQKVResource(MemoryResource): + """Provides TMA descriptors and per-tile coordinates for Q/K/V loads. + + Consumer side resolves batch/head/seq coordinates from the work tile + so that downstream SmemQ/SmemKV producer_work can issue TMA loads. + """ + + tma_q_desc: cutlass.Pointer | None = field(init=False, default=None) + tma_k_desc: cutlass.Pointer | None = field(init=False, default=None) + tma_v_desc: cutlass.Pointer | None = field(init=False, default=None) + cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) + cum_seqlen_k: cute.Tensor | None = field(init=False, default=None) + variable_window_token_starts: cute.Tensor | None = field(init=False, default=None) + variable_window_cta_starts: cute.Tensor | None = field(init=False, default=None) + variable_window_q_stride: int | Int32 = field(init=False, default=0) + q_offset_default: int | Int32 = field(init=False, default=0) + seqlens_kv: cute.Pointer | None = field(init=False, default=None) + block_table_row_stride: int | Int32 = field(init=False, default=0) + max_seq_len_kv: Optional[Int32 | int] = field(init=False, default=None) + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + seq_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + head_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + kv_head_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + head_coord_kv: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + batch_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + seq_coord_q: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + cuseqlen_q: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + cuseqlen_k: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + seqlen_q: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + seqlen_k: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + kv_tile_start: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + kv_request_begin: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + kv_page_idx_ub: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + tma_q_desc: cutlass.Pointer | None, + tma_k_desc: cutlass.Pointer | None, + tma_v_desc: cutlass.Pointer | None, + cum_seqlen_q: cute.Tensor | None, + cum_seqlen_k: cute.Tensor | None, + q_offset: int | Int32, + cfg: FmhaConfig, + seqlens_kv: cute.Pointer | None = None, + block_table_row_stride: int | Int32 = 0, + max_seq_len_kv: Int32 | int | None = None, + variable_window_token_starts: cute.Tensor | None = None, + variable_window_cta_starts: cute.Tensor | None = None, + variable_window_q_stride: int | Int32 = 0, + **kwargs: Any, + ) -> None: + """Bind Q/K/V descriptors, optional varlen metadata, and FMHA config.""" + super().__init__(**kwargs) + self.tma_q_desc = tma_q_desc + self.tma_k_desc = tma_k_desc + self.tma_v_desc = tma_v_desc + self.cum_seqlen_q = cum_seqlen_q + self.cum_seqlen_k = cum_seqlen_k + self.q_offset_default = q_offset + self.seqlens_kv = seqlens_kv + self.block_table_row_stride = block_table_row_stride + self.max_seq_len_kv = max_seq_len_kv + self.variable_window_token_starts = variable_window_token_starts + self.variable_window_cta_starts = variable_window_cta_starts + self.variable_window_q_stride = variable_window_q_stride + self.cfg = cfg + self.seq_coord = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Q/K/V tile sequence coordinate.", + ) + self.head_coord = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Q/O head coordinate for the current work tile.", + ) + self.kv_head_coord = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="K/V head coordinate for the current work tile.", + ) + self.head_coord_kv = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="K/V head coordinate mirrored for master FMHA context schedules.", + ) + self.batch_coord = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Batch coordinate for the current work tile.", + ) + self.seq_coord_q = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Q/O row coordinate for the current work tile.", + ) + self.cuseqlen_q = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Q sequence cumulative offset for variable-length FMHA.", + ) + self.cuseqlen_k = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="K sequence cumulative offset for variable-length FMHA.", + ) + self.seqlen_q = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Q sequence length for variable-length FMHA.", + ) + self.seqlen_k = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="K sequence length for variable-length FMHA.", + ) + self.kv_tile_start = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="First K/V loop tile for sliding-window FMHA.", + ) + self.kv_request_begin = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Element offset of the request's block-table row.", + ) + self.kv_page_idx_ub = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Inclusive logical-page upper bound for the request.", + ) + + @consumer_work( + returns=( + seq_coord, + head_coord, + kv_head_coord, + head_coord_kv, + batch_coord, + seq_coord_q, + cuseqlen_q, + cuseqlen_k, + seqlen_q, + seqlen_k, + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) + ) + @cute.jit + def compute_coords( + self, stage_info: StageInfo + ) -> tuple[ + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + ]: + """Resolve per-tile coordinates from work_tile for downstream use. + + Populates consumer variables with the batch/head/seq coordinates + that SmemQ, SmemK, and SmemV producer_work methods need for TMA loads. + GQA: head_coord indexes Q/O heads (h_q), kv_head_coord indexes K/V + heads (h_kv). For MHA (h_r=1) they are identical. + """ + seq_coord, head_coord, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + + kv_head_coord = (head_coord * self.cfg.work_tile_q_heads) // self.cfg.h_r + seq_coord_q = seq_coord * self.cfg.q_tile_m * self.cfg.work_tile_q_seq_tiles + head_coord_kv = kv_head_coord + cuseqlen_q = Int32(0) + cuseqlen_k = Int32(0) + seqlen_q = Int32(0) + seqlen_k = Int32(0) + window_q_offset = Int32(self.q_offset_default) + kv_tile_start = Int32(0) + kv_request_begin = Int32(0) + kv_page_idx_ub = Int32(0) + if cutlass.const_expr(self.cfg.has_varlen): + if cutlass.const_expr(self.cfg.has_uniform_varlen): + seqlen_q = Int32(self.cfg.uniform_seq_len_q) + cuseqlen_q = batch_coord * seqlen_q + else: + cuseqlen_q = Int32(self.cum_seqlen_q[batch_coord]) + next_cuseqlen_q = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) + seqlen_q = next_cuseqlen_q - cuseqlen_q + if cutlass.const_expr(self.cfg.use_paged_kv): + # Paged K/V is addressed through a block table rather than a + # packed token buffer, so it has no cumulative token offset. + if cutlass.const_expr(self.cfg.has_uniform_varlen): + seqlen_k = Int32(self.cfg.uniform_seq_len_k) + else: + from .helpers_paged import _load_runtime_seq_len_kv + + seqlen_k = _load_runtime_seq_len_kv( + self.seqlens_kv, self.max_seq_len_kv, batch_coord + ) + elif cutlass.const_expr(self.cfg.has_uniform_varlen): + seqlen_k = Int32(self.cfg.uniform_seq_len_k) + cuseqlen_k = batch_coord * seqlen_k + else: + cuseqlen_k = Int32(self.cum_seqlen_k[batch_coord]) + next_cuseqlen_k = Int32(self.cum_seqlen_k[batch_coord + Int32(1)]) + seqlen_k = next_cuseqlen_k - cuseqlen_k + seq_coord_q = cuseqlen_q + seq_coord_q + # Each packed request uses its own bottom-right window origin. For + # mixed causal plans the task manager also derives the request's + # K-loop extent from these live Q/K lengths. + window_q_offset = seqlen_k - seqlen_q + if cutlass.const_expr( + self.cfg.use_paged_kv and not self.cfg.stages_page_offsets_in_smem + ): + if cutlass.const_expr(self.cfg.has_uniform_varlen): + kv_request_begin = batch_coord * Int32(self.block_table_row_stride) + kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1) + else: + from .helpers_paged import _load_block_table_row_bounds + + kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( + Int32(self.block_table_row_stride), + self.cfg, + seqlen_k, + batch_coord, + ) + if cutlass.const_expr(self.cfg.kv_tile_start_window_size_left > 0): + if cutlass.const_expr(self.cfg.has_varlen or self.cfg.has_q_offset): + kv_tile_start = bottom_right_window_tile_start( + seq_coord=seq_coord, + q_tile_m=self.cfg.q_tile_m, + kv_tile_n=self.cfg.seq_tile_n, + q_offset=window_q_offset, + window_size_left=self.cfg.kv_tile_start_window_size_left, + ) + else: + # Preserve the minimal fixed equal-length specialization: its + # bottom-right offset is statically zero. + kv_tile_start = cute.math.max( + Int32(0), + ( + seq_coord * self.cfg.q_tile_m + - self.cfg.kv_tile_start_window_size_left + ) + // self.cfg.seq_tile_n, + ) + if cutlass.const_expr(self.cfg.has_variable_window): + min_window_start = variable_window_cta_min_start( + self.variable_window_cta_starts, + batch_coord=batch_coord, + seq_coord=seq_coord, + q_stride=self.variable_window_q_stride, + tile_size_q=self.cfg.cta_tiler[0], + ) + kv_tile_start = min_window_start // self.cfg.kv_tile_n + return ( + seq_coord, + head_coord, + kv_head_coord, + head_coord_kv, + batch_coord, + seq_coord_q, + cuseqlen_q, + cuseqlen_k, + seqlen_q, + seqlen_k, + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) + + @consumer_work( + returns=( + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) + ) + @cute.jit + def compute_page_coords(self, stage_info: StageInfo) -> tuple[Int32, Int32, Int32]: + """Resolve only the coordinates needed by paged-KV prefetch. + + The page-offset warp does not consume Q/head coordinates. Keeping its + coordinate path narrow avoids materializing the full Q/K/V coordinate + tuple once per persistent work tile. + """ + seq_coord, _head_coord, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + + if cutlass.const_expr(self.cfg.has_uniform_varlen): + cached_seqlen_kv = Int32(self.cfg.uniform_seq_len_k) + kv_request_begin = batch_coord * Int32(self.block_table_row_stride) + kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1) + else: + from .helpers_paged import _load_runtime_seq_len_kv + + cached_seqlen_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, self.max_seq_len_kv, batch_coord + ) + from .helpers_paged import _load_block_table_row_bounds + + kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( + Int32(self.block_table_row_stride), + self.cfg, + cached_seqlen_kv, + batch_coord, + ) + window_q_offset = Int32(self.q_offset_default) + kv_tile_start = Int32(0) + if cutlass.const_expr(self.cfg.kv_tile_start_window_size_left > 0): + if cutlass.const_expr(self.cfg.has_varlen): + if cutlass.const_expr(self.cfg.has_uniform_varlen): + seqlen_q = Int32(self.cfg.uniform_seq_len_q) + else: + cuseqlen_q = Int32(self.cum_seqlen_q[batch_coord]) + next_cuseqlen_q = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) + seqlen_q = next_cuseqlen_q - cuseqlen_q + window_q_offset = cached_seqlen_kv - seqlen_q + if cutlass.const_expr(self.cfg.has_varlen or self.cfg.has_q_offset): + kv_tile_start = bottom_right_window_tile_start( + seq_coord=seq_coord, + q_tile_m=self.cfg.q_tile_m, + kv_tile_n=self.cfg.seq_tile_n, + q_offset=window_q_offset, + window_size_left=self.cfg.kv_tile_start_window_size_left, + ) + else: + kv_tile_start = cute.math.max( + Int32(0), + ( + seq_coord * self.cfg.q_tile_m + - self.cfg.kv_tile_start_window_size_left + ) + // self.cfg.seq_tile_n, + ) + + return ( + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) + + +def _qkv_inner_dim_size_bytes(cfg: FmhaConfig) -> int: + """Return the byte width of one Q/K/V tile inner dimension.""" + return cfg.qk_mma_tiler[2] * cfg.q_dtype.width // 8 + + +def _qkv_smem_layout(cfg: FmhaConfig) -> int: + """Return the tcgen05 descriptor layout selector for Q/K/V SMEM tiles.""" + inner_dim_size = _qkv_inner_dim_size_bytes(cfg) + if inner_dim_size % 128 == 0: + return 2 + if inner_dim_size == 64: + return 4 + if inner_dim_size == 32: + return 6 + raise RuntimeError(f"Unsupported inner dimension size: {inner_dim_size}") + + +def _qk_smem_desc_offsets(cfg: FmhaConfig) -> SmemDescOffsets: + """Return Q/K descriptor leading and stride byte offsets.""" + leading_byte_offset = 0 if cfg.head_paired else 16 + stride_byte_offset = ( + cfg.qk_mma_tiler[2] * cfg.q_dtype.width // cfg.tma_copy_qkv_iters + ) + return leading_byte_offset, stride_byte_offset + + +def _pv_smem_desc_offsets(cfg: FmhaConfig) -> SmemDescOffsets: + """Return V descriptor leading and stride byte offsets for PV MMA.""" + leading_byte_offset = 0 + if cfg.tma_copy_qkv_iters != 1: + tma_copy_kv_iters = ( + cfg.tma_copy_kv_stage_iters + if cfg.stage_kv_by_head_dim + else cfg.tma_copy_qkv_iters + ) + leading_byte_offset = cfg.tma_copy_kv_bytes // tma_copy_kv_iters + stride_byte_offset = ( + cfg.pv_mma_tiler[1] * cfg.v_dtype.width // cfg.tma_copy_qkv_iters + ) + return leading_byte_offset, stride_byte_offset + + +def _qkv_smem_swizzle(cfg: FmhaConfig) -> cutlass.Swizzle: + """Return the physical TMA swizzle used by Q/K/V SMEM fragments.""" + inner_dim_size = _qkv_inner_dim_size_bytes(cfg) + if inner_dim_size % 128 == 0: + return cutlass.Swizzle(3, 4, 3) + if inner_dim_size == 64: + return cutlass.Swizzle(2, 4, 3) + if inner_dim_size == 32: + return cutlass.Swizzle(1, 4, 3) + raise RuntimeError(f"Unsupported inner dimension size: {inner_dim_size}") + + +def _smem_o_swizzle(cfg: FmhaConfig) -> cutlass.Swizzle: + """Return the shared-memory swizzle used when staging O for TMA store.""" + return _qkv_smem_swizzle(cfg) + + +# --------------------------------------------------------------------------- +# SmemQResource -- SMEM Q tile buffer with TMA pipeline +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class SmemQResource(MemoryResource): + """SMEM buffer for Q tiles with a topology-derived TmaUmma pipeline. + + Producer: LoadTask (TMA loads Q0 and Q1 in the first K-loop iteration). + Consumer: MmaTask (builds SMEM descriptors, holds Q across K-loop). + """ + + sQ_array: cutlass.Array = field(init=False, default=None) + tma_q_desc: cutlass.Pointer | None = field(init=False, default=None) + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + desc_q0_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + desc_q1_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + tma_q_desc: cutlass.Pointer | None, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + **kwargs: Any, + ) -> None: + """Bind the Q TMA descriptor and reserve SMEM for staged Q tiles.""" + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.tma_q_desc = tma_q_desc + self.cfg = cfg + total_elements = cfg.sQ_shape[0] * cfg.sQ_shape[1] + size_bytes = total_elements * cfg.q_dtype.width // 8 + self._alloc = SmemAllocation( + "smem_q", size_bytes, alignment=cfg.buffer_align_bytes + ) + self.sQ_array = _placeholder_smem_array(cfg.q_dtype) + self.desc_q0_base = TaskLocalVariable( + dtype=cutlass.Int64, + default=cutlass.Int64(0), + docs="SMEM descriptor base for the first Q half.", + ) + self.desc_q1_base = TaskLocalVariable( + dtype=cutlass.Int64, + default=cutlass.Int64(0), + docs="SMEM descriptor base for the second Q half.", + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Return the SMEM allocation required for staged Q tiles.""" + return [self._alloc] + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + """Materialize the Q SMEM array and descriptor dataflow slots.""" + smem_base = stage_info.context.smem_base + total_elements = self.cfg.sQ_shape[0] * self.cfg.sQ_shape[1] + self.sQ_array = cutlass.Array( + smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.q_dtype, + shape=(total_elements,), + addrspace=3, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @producer_work + @cute.jit + def tma_load( + self, + stage_info: StageInfo, + *, + seq_coord_q: Int32, + head_coord: Int32, + batch_coord: Int32, + cuseqlen_q: Int32, + seqlen_q: Int32, + inst_idx: cutlass.Constexpr[int], + ) -> None: + """TMA load one Q tile (Q0 or Q1) from GMEM to SMEM. + + inst_idx 0 = Q0, inst_idx 1 = Q1. + Uses seq_coord_q from producer variables (forwarded from GmemQKV). + """ + q_head_coord = ( + head_coord * self.cfg.work_tile_q_heads + + inst_idx * self.cfg.peer_q_head_stride + ) + q_seq_offset = ( + seq_coord_q + inst_idx * self.cfg.peer_q_seq_tile_stride * self.cfg.q_tile_m + ) + q_seq_extent = Int32(0) + if cutlass.const_expr(self.cfg.has_varlen): + q_seq_extent = cuseqlen_q + seqlen_q - q_seq_offset + smem_stage_elements = self.cfg.tma_copy_q_elements + d_granu_inner = self.cfg.tma_copy_q_granu_inner + + sQ_curr = self.sQ_array.subview(stage_info.stage_idx * smem_stage_elements) + if prims.elect_sync(): + for i in cutlass.range_constexpr(self.cfg.tma_copy_qkv_iters): + d_offset = i * d_granu_inner + q_coords = (d_offset, q_head_coord, q_seq_offset, batch_coord) + if cutlass.const_expr(self.cfg.has_varlen): + q_coords = (d_offset, q_head_coord, q_seq_offset) + q_coords = transform_ragged_coords( + q_coords, + ragged_dim_idx=2, + ragged_box_size=self.cfg.qk_mma_tiler[0], + ragged_extent=q_seq_extent, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + sQ_curr.subview(i * self.cfg.tma_copy_q_granu_elems), + self.tma_q_desc, + q_coords, + stage_info.barrier, + ) + + def _build_q_descriptor(self, inst_idx: int) -> prims.Tcgen05SmemDesc: + """Build SMEM descriptor for the current Q tile. + + Uses inst_idx (not stage_idx) to compute the SMEM offset because + Q is consumed twice in HEAD without an intervening ConsumerRelease, + which would otherwise advance consumer_state. + """ + sQ_curr = self.sQ_array.subview(inst_idx * self.cfg.tma_copy_q_elements) + leading_byte_offset, stride_byte_offset = _qk_smem_desc_offsets(self.cfg) + return prims.Tcgen05SmemDesc.build( + sQ_curr, + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_layout(self.cfg), + ) + + @consumer_work(returns=desc_q0_base) + @cute.jit + def q0_desc( + self, stage_info: StageInfo, *, inst_idx: cutlass.Constexpr[int] + ) -> prims.Tcgen05SmemDesc: + """Build Q0 SMEM descriptor -> desc_q0_base.""" + return self._build_q_descriptor(inst_idx) + + @consumer_work(returns=desc_q1_base) + @cute.jit + def q1_desc( + self, stage_info: StageInfo, *, inst_idx: cutlass.Constexpr[int] + ) -> prims.Tcgen05SmemDesc: + """Build Q1 SMEM descriptor -> desc_q1_base.""" + return self._build_q_descriptor(inst_idx) + + +# --------------------------------------------------------------------------- +# SmemPageOffsetsKvResource -- paged-KV page-table cache in SMEM +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class SmemPageOffsetsKvResource(MemoryResource): + """Paged-KV logical-to-physical page IDs staged in SMEM (context kernel). + + The staged D256 path uses a dedicated warp to prefetch page-table + entries for the next K/V tile so the TMA load warp can read SMEM-cached + offsets. Each pipeline stage holds one topology-derived page-ID window + from the request's fixed-table row; all 32 lanes co-load it. ``page_ids`` slices + ``pages_per_tile`` entries for the current tile. + + Differences from decode: + - Single ``load_k`` / ``load_v`` producer pair (context has no + ``num_insts_kv > 1`` four-way split). + - Consumer release labels bind to ``{"k_load", "v_load"}`` (matching + ``SmemKVResource`` producer names). + - Driven by the staged D256 ``cfg.empty_warp_id`` (warp 11). + Paired D128 does not instantiate this resource. + """ + + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + block_tables: cute.Pointer | None = field(init=False, default=None) + page_table_is_v: Constexpr[bool] = field(init=False, default=False) + _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + _smem_page_offsets: cutlass.Array = field(init=False, default=None) + cached_page_ids: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + block_tables: cute.Pointer | None, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + page_table_is_v: bool = False, + **kwargs: Any, + ) -> None: + # ``page_ids`` runs from the downstream K/V resource after this + # resource's ConsumerWait. Preserve the waited stage so the nested + # lookup reads the matching page-table data. + pipeline_config = replace(pipeline_config, advance_on_wait=True) + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.block_tables = block_tables + self.page_table_is_v = page_table_is_v + num_stages = pipeline_config.num_stages + total_entries = num_stages * cfg.page_table_window_entries + self._alloc = SmemAllocation( + "smem_page_offsets_v" if page_table_is_v else "smem_page_offsets_k", + size_bytes=total_entries * 4, + # Page-size 16 consumes eight page IDs per 128-token K/V tile. + # Align only that specialization for one 32-byte vector load; + # preserve the established layout for larger page sizes. + alignment=(32 if cfg.kv_tile_n // cfg.num_tokens_per_page == 8 else 16), + ) + self._smem_page_offsets = _placeholder_smem_array(Int32, total_entries) + self.cached_page_ids = TaskLocalVariable( + dtype=cutlass.Array, + default=None, + docs="Page IDs retained while a delayed V tile crosses a page window.", + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + return [self._alloc] + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + smem_base = stage_info.context.smem_base + num_stages = self.pipeline_config.num_stages + total_entries = num_stages * self.cfg.page_table_window_entries + self._smem_page_offsets = cutlass.Array( + smem_base.data_ptr() + self._alloc.offset, + dtype=cutlass.Int32, + shape=(total_entries,), + addrspace=3, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_read_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=cached_page_ids) + @cute.jit + def init_cached_read_state(self, stage_info: StageInfo) -> cutlass.Array: + """Initialize the SMEM view and register cache for one tile's page IDs.""" + self._init_smem_state(stage_info) + return cutlass.Array( + Int32, + self.cfg.kv_tile_n // self.cfg.num_tokens_per_page, + space=cutlass.AddressSpace.rmem, + ) + + @cute.jit + def page_ids(self, tile_idx: Int32) -> cutlass.Array: + """Slice ``pages_per_tile`` entries from the cached page-ID stage. + + ``tile_idx`` is the runtime-resolved K/V tile index (same expression + the page-offsets producer uses). The window-aligned base is implicit + in the stage's contents; this LDS picks the per-tile entries. + """ + cfg = self.cfg + pages_per_tile = cfg.kv_tile_n // cfg.num_tokens_per_page + window_entries = cfg.page_table_window_entries + stage_idx = self.state_src.consumer_work_stage + group_page_idx = (tile_idx * Int32(pages_per_tile)) & Int32(window_entries - 1) + offset = stage_idx * Int32(window_entries) + group_page_idx + if cutlass.const_expr(pages_per_tile == 8): + return self._smem_page_offsets.load(offset, vector_size=8, alignment=32) + if cutlass.const_expr(pages_per_tile == 4): + return self._smem_page_offsets.load(offset, vector_size=4, alignment=16) + if cutlass.const_expr(pages_per_tile == 2): + return self._smem_page_offsets.load(offset, vector_size=2, alignment=8) + return self._smem_page_offsets.load(offset, vector_size=1, alignment=4) + + @cute.jit + def _producer_load_page_offsets( + self, + stage_info: StageInfo, + tile_offset: cutlass.Constexpr[int] = 0, + *, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + from .helpers_paged import _resolve_kv_tile_idx_context + + cfg = self.cfg + # Context's K/V tile index is kv_tile_start + loop_offset; for V the + # producer runs one tile ahead in TAIL, so reuse the same expression + # but consult kv_tile_start with the loop's stage_info. + tile_idx = _resolve_kv_tile_idx_context( + stage_info, kv_tile_start, tile_offset=tile_offset + ) + pages_per_tile = Int32(cfg.kv_tile_n // cfg.num_tokens_per_page) + + block_tables = self.block_tables + smem_page_offsets = self._smem_page_offsets + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + # Lanes cooperatively fetch one topology-derived aligned window. A + # split-D window gives each lane one scalar per D stage; consumers then + # slice per-tile entries via ``page_ids``. + window_entries = cfg.page_table_window_entries + grouped_base_page_idx = ( + (tile_idx * pages_per_tile) // Int32(window_entries) + ) * Int32(window_entries) + grouped_smem_base = stage_info.stage_idx * Int32(window_entries) + entries_per_lane = window_entries // cute.arch.WARP_SIZE + for lane_group in cutlass.range_constexpr(entries_per_lane): + lane_offset = lane_idx + Int32(lane_group * cute.arch.WARP_SIZE) + grouped_logical_page_idx = cute.math.min( + grouped_base_page_idx + lane_offset, kv_page_idx_ub + ) + prims.cp_async_shared_global( + smem_page_offsets.data_ptr() + grouped_smem_base + lane_offset, + block_tables + kv_request_begin + grouped_logical_page_idx, + 4, + "ca", + ) + + @producer_work + @cute.jit + def load_k( + self, + stage_info: StageInfo, + *, + tile_offset: cutlass.Constexpr[int] = 0, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """Prefetch K-side page IDs for the current K tile.""" + self._producer_load_page_offsets( + stage_info, + tile_offset=tile_offset, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + + @producer_work + @cute.jit + def load_v( + self, + stage_info: StageInfo, + *, + previous: cutlass.Constexpr[bool] = False, + tile_offset: cutlass.Constexpr[int] = 0, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """Prefetch V-side page IDs for the current V tile.""" + self._producer_load_page_offsets( + stage_info, + tile_offset=tile_offset + (-1 if previous else 0), + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + + @consumer_work + @cute.jit + def read_offsets(self, stage_info: StageInfo) -> None: + return + + @consumer_work(returns=cached_page_ids) + @cute.jit + def cache_tile_page_ids( + self, + stage_info: StageInfo, + *, + cached_page_ids: cutlass.Array, + kv_tile_start: Int32, + tile_offset: cutlass.Constexpr[int] = 0, + ) -> cutlass.Array: + """Retain one tile's page IDs after its SMEM window is released.""" + from .helpers_paged import _resolve_kv_tile_idx_context + + tile_idx = _resolve_kv_tile_idx_context( + stage_info, kv_tile_start, tile_offset=tile_offset + ) + page_ids = self.page_ids(tile_idx) + pages_per_tile = self.cfg.kv_tile_n // self.cfg.num_tokens_per_page + for page_frag in cutlass.range_constexpr(pages_per_tile): + cached_page_ids[page_frag] = Int32(page_ids[page_frag]) + return cached_page_ids + + def dma_consumer_release_labels_for( + self, downstream: MemoryResource + ) -> set[str] | None: + """Bind page-offset releases to the K/V TMA loads that consumed them.""" + if isinstance(downstream, SmemKVResource): + if self.cfg.reuses_page_table_windows: + if self.page_table_is_v: + return {"v_load", "v_load_stage", "v_load_stage_cached"} + return {"k_load", "k_load_stage"} + return { + "k_load", + "v_load", + "k_load_stage", + "v_load_stage", + "v_load_stage_cached", + } + return None + + +# --------------------------------------------------------------------------- +# SmemKVResource -- SMEM K/V tile buffer with TMA pipeline +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class SmemKVResource(MemoryResource): + """SMEM buffer for K and V tiles with a capacity-derived TmaUmma pipeline. + + K and V tiles alternate in the pipeline stages: K0, V0, K1, V1, ... + Producer: LoadTask (TMA loads K/V tiles). + Consumer: MmaTask (builds SMEM descriptors for QK and PV MMAs). + """ + + sK_array: cutlass.Array = field(init=False, default=None) + tma_k_desc: cutlass.Pointer | None = field(init=False, default=None) + tma_v_desc: cutlass.Pointer | None = field(init=False, default=None) + page_offsets_kv: Optional["SmemPageOffsetsKvResource"] = field( + init=False, default=None + ) + page_offsets_v: Optional["SmemPageOffsetsKvResource"] = field( + init=False, default=None + ) + block_tables: cute.Pointer | None = field(init=False, default=None) + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + desc_k_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + desc_v_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + tma_k_desc: cutlass.Pointer | None, + tma_v_desc: cutlass.Pointer | None, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + page_offsets_kv: Optional["SmemPageOffsetsKvResource"] = None, + page_offsets_v: Optional["SmemPageOffsetsKvResource"] = None, + block_tables: cute.Pointer | None = None, + **kwargs: Any, + ) -> None: + """Bind K/V TMA descriptors and reserve shared SMEM staging.""" + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.tma_k_desc = tma_k_desc + self.tma_v_desc = tma_v_desc + self.page_offsets_kv = page_offsets_kv + self.page_offsets_v = ( + page_offsets_v if page_offsets_v is not None else page_offsets_kv + ) + self.block_tables = block_tables + self.cfg = cfg + total_elements = cfg.sK_shape[0] * cfg.sK_shape[1] + size_bytes = total_elements * cfg.k_dtype.width // 8 + self._alloc = SmemAllocation( + "smem_kv", size_bytes, alignment=cfg.buffer_align_bytes + ) + self.sK_array = _placeholder_smem_array(cfg.k_dtype) + self.desc_k_base = TaskLocalVariable( + dtype=cutlass.Int64, + default=cutlass.Int64(0), + docs="SMEM descriptor base for the current K tile.", + ) + self.desc_v_base = TaskLocalVariable( + dtype=cutlass.Int64, + default=cutlass.Int64(0), + docs="SMEM descriptor base for the current V tile.", + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Return the SMEM allocation required for staged K/V tiles.""" + return [self._alloc] + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + """Materialize K/V SMEM storage and descriptor dataflow slots.""" + smem_base = stage_info.context.smem_base + total_elements = self.cfg.sK_shape[0] * self.cfg.sK_shape[1] + self.sK_array = cutlass.Array( + smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.k_dtype, + shape=(total_elements,), + addrspace=3, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @property + def loop_offset_sensitive(self) -> bool: + """Return true because K/V loads index the current loop tile.""" + # producer_work uses loop_offset to compute seq_coord_kv. + return True + + @cute.jit + def _tma_load( + self, + stage_info: StageInfo, + tma_desc: cutlass.Pointer | None, + is_v: cutlass.Constexpr[bool] = False, + tile_offset: cutlass.Constexpr[int] = 0, + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + cached_page_ids: cutlass.Array | None = None, + *, + kv_head_coord: Int32, + batch_coord: Int32, + cuseqlen_k: Int32, + seqlen_k: Int32, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """Issue TMA bulk-copy for one K or V tile.""" + seq_offset = ( + kv_tile_start + stage_info.loop_offset + tile_offset + ) * self.cfg.kv_tile_n + smem_stage_elements = self.cfg.tma_copy_kv_elements + sK_curr = self.sK_array.subview(stage_info.stage_idx * smem_stage_elements) + + if cutlass.const_expr(self.cfg.use_paged_kv): + # Paged-KV path: read pre-staged page IDs and issue one TMA per + # (page fragment, d fragment). K and V share the native rank-4 + # descriptor coordinates (d_off, 0, kv_head_coord, page_id). + tile_idx = kv_tile_start + stage_info.loop_offset + tile_offset + pages_per_tile = self.cfg.kv_tile_n // self.cfg.num_tokens_per_page + d_granu_inner = self.cfg.tma_copy_kv_granu_inner + page_d_elems = self.cfg.num_tokens_per_page * d_granu_inner + d_iter_elems = self.cfg.tma_copy_kv_granu_elems + if prims.elect_sync(): + # Only the elected TMA-issuing lane consumes page IDs. Loading + # the vector outside this guard made every lane perform the + # same SMEM read for every K and V tile. + page_ids = cached_page_ids + if cutlass.const_expr(page_ids is None): + page_offsets = ( + self.page_offsets_v + if cutlass.const_expr(is_v) + else self.page_offsets_kv + ) + if cutlass.const_expr(page_offsets is not None): + page_ids = page_offsets.page_ids(tile_idx) + else: + # The paired K/V schedule consumes K and V together. + # Reading its four contiguous page IDs directly avoids + # a producer warp spinning on an always-full auxiliary + # pipeline and leaves that warp available for CLC. + # K and V share the same fixed logical-to-physical page + # row. Clamp both to the pages covered by the request's + # runtime sequence length so padding IDs are untouched. + logical_page_idx = tile_idx * Int32(pages_per_tile) + page_ids = cutlass.Array( + Int32, + pages_per_tile, + space=cutlass.AddressSpace.rmem, + ) + for frag in cutlass.range_constexpr(pages_per_tile): + clamped_page_idx = cute.math.min( + logical_page_idx + Int32(frag), kv_page_idx_ub + ) + page_ids[frag] = Int32( + self.block_tables[kv_request_begin + clamped_page_idx] + ) + for frag in cutlass.range_constexpr(pages_per_tile): + page_id = Int32(page_ids[frag]) + for i in cutlass.range_constexpr(self.cfg.tma_copy_kv_stage_iters): + d_offset = Int32( + head_dim_stage_idx * self.cfg.head_dim_per_stage_kv + + i * d_granu_inner + ) + smem_offset = Int32(i * d_iter_elems + frag * page_d_elems) + prims.cp_async_bulk_tensor_shared_cta_global( + sK_curr.subview(smem_offset), + tma_desc, + (d_offset, Int32(0), kv_head_coord, page_id), + stage_info.barrier, + ) + return + + if prims.elect_sync(): + d_granu_inner = self.cfg.tma_copy_kv_granu_inner + seq_coord_kv = cuseqlen_k + seq_offset + for i in cutlass.range_constexpr(self.cfg.tma_copy_kv_stage_iters): + d_offset = ( + head_dim_stage_idx * self.cfg.head_dim_per_stage_kv + + i * d_granu_inner + ) + kv_coords = (d_offset, kv_head_coord, seq_coord_kv, batch_coord) + if cutlass.const_expr(self.cfg.has_varlen): + kv_coords = (d_offset, kv_head_coord, seq_coord_kv) + prims.cp_async_bulk_tensor_shared_cta_global( + sK_curr.subview(i * self.cfg.tma_copy_kv_granu_elems), + tma_desc, + kv_coords, + stage_info.barrier, + ) + + @producer_work + @cute.jit + def k_load( + self, + stage_info: StageInfo, + *, + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + tile_offset: cutlass.Constexpr[int] = 0, + kv_head_coord: Int32, + batch_coord: Int32, + cuseqlen_k: Int32, + seqlen_k: Int32, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """TMA load K tile from GMEM to SMEM.""" + self._tma_load( + stage_info, + self.tma_k_desc, + tile_offset=tile_offset, + head_dim_stage_idx=head_dim_stage_idx, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + + @producer_work + @cute.jit + def k_load_stage( + self, + stage_info: StageInfo, + *, + stage_id: Constexpr[int], + tile_offset: Constexpr[int] = 0, + kv_head_coord: Int32, + batch_coord: Int32, + cuseqlen_k: Int32, + seqlen_k: Int32, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """TMA load one K head-dim stage for split D scheduling.""" + self._tma_load( + stage_info, + self.tma_k_desc, + False, + tile_offset=tile_offset, + head_dim_stage_idx=stage_id, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + + @producer_work + @cute.jit + def v_load( + self, + stage_info: StageInfo, + *, + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + tile_offset: cutlass.Constexpr[int] = 0, + kv_head_coord: Int32, + batch_coord: Int32, + cuseqlen_k: Int32, + seqlen_k: Int32, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """TMA load V tile from GMEM to SMEM.""" + self._tma_load( + stage_info, + self.tma_v_desc, + True, + tile_offset=tile_offset, + head_dim_stage_idx=head_dim_stage_idx, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + + @producer_work + @cute.jit + def v_load_stage( + self, + stage_info: StageInfo, + *, + stage_id: Constexpr[int], + previous: Constexpr[bool] = False, + tile_offset: Constexpr[int] = 0, + kv_head_coord: Int32, + batch_coord: Int32, + cuseqlen_k: Int32, + seqlen_k: Int32, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """TMA load one current or previous V head-dim stage.""" + self._tma_load( + stage_info, + self.tma_v_desc, + True, + tile_offset=tile_offset + (-1 if previous else 0), + head_dim_stage_idx=stage_id, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + + @producer_work + @cute.jit + def v_load_stage_cached( + self, + stage_info: StageInfo, + *, + cached_v_page_ids: cutlass.Array, + stage_id: Constexpr[int], + tile_offset: Constexpr[int] = 0, + kv_head_coord: Int32, + batch_coord: Int32, + cuseqlen_k: Int32, + seqlen_k: Int32, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + ) -> None: + """Load one V head-dimension stage using register-cached page IDs.""" + self._tma_load( + stage_info, + self.tma_v_desc, + True, + tile_offset=tile_offset, + head_dim_stage_idx=stage_id, + cached_page_ids=cached_v_page_ids, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + + @consumer_work(returns=desc_k_base) + @cute.jit + def k_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Build K SMEM descriptor (K-major layout for QK MMA) -> desc_k_base.""" + smem_stage_elements = self.cfg.tma_copy_kv_elements + sK_curr = self.sK_array.subview(stage_info.stage_idx * smem_stage_elements) + leading_byte_offset, stride_byte_offset = _qk_smem_desc_offsets(self.cfg) + desc_k_base = prims.Tcgen05SmemDesc.build( + sK_curr, + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_layout(self.cfg), + ) + return desc_k_base + + @cute.jit + def _zero_paged_v_tail( + self, + stage_info: StageInfo, + *, + section: cutlass.Constexpr[FmhaStage], + tile_offset: cutlass.Constexpr[int], + seqlen_k: Int32, + kv_tile_start: Int32, + ) -> None: + """Overwrite request-invalid V rows after TMA completion.""" + if cutlass.const_expr(section == FmhaStage.Head): + domain_tile_idx = stage_info.loop_start + elif cutlass.const_expr(section == FmhaStage.Tail): + domain_tile_idx = stage_info.loop_end + else: + domain_tile_idx = stage_info.loop_offset + logical_v_tile_idx = kv_tile_start + domain_tile_idx + tile_offset + valid_rows = cute.math.min( + cute.math.max( + seqlen_k - logical_v_tile_idx * Int32(self.cfg.kv_tile_n), + Int32(0), + ), + Int32(self.cfg.kv_tile_n), + ) + + if valid_rows < Int32(self.cfg.kv_tile_n): + # Each paged TMA transaction writes one swizzled + # (D-fragment, page-token) box. Pages are concatenated within a D + # iteration, and D iterations are concatenated within the stage. + # Mirror that exact physical layout: a flat row-major clear would + # target the wrong bytes under the s128b swizzle. + d_granu_inner = self.cfg.tma_copy_kv_granu_inner + chunks_per_d_iter = d_granu_inner // 16 + chunks_per_v_row = self.cfg.tma_copy_kv_stage_iters * chunks_per_d_iter + page_d_elems = self.cfg.num_tokens_per_page * d_granu_inner + d_iter_elems = self.cfg.tma_copy_kv_granu_elems + invalid_chunks = (Int32(self.cfg.kv_tile_n) - valid_rows) * Int32( + chunks_per_v_row + ) + zero_vec = cutlass.vector.full( + [16], self.cfg.v_dtype(0.0), dtype=self.cfg.v_dtype + ) + sV_curr = self.sK_array.subview( + stage_info.stage_idx * self.cfg.tma_copy_kv_elements + ) + lane_idx = cute.arch.lane_idx() + for tail_chunk in cutlass.range( + lane_idx, + invalid_chunks, + Int32(cute.arch.WARP_SIZE), + unroll=1, + ): + invalid_row = tail_chunk // Int32(chunks_per_v_row) + d_chunk = tail_chunk - invalid_row * Int32(chunks_per_v_row) + d_iter = d_chunk // Int32(chunks_per_d_iter) + d_chunk_in_iter = d_chunk - d_iter * Int32(chunks_per_d_iter) + logical_row = valid_rows + invalid_row + page_frag = logical_row // Int32(self.cfg.num_tokens_per_page) + row_in_page = logical_row - page_frag * Int32( + self.cfg.num_tokens_per_page + ) + smem_offset = ( + d_iter * Int32(d_iter_elems) + + page_frag * Int32(page_d_elems) + + row_in_page * Int32(d_granu_inner) + + d_chunk_in_iter * Int32(16) + ) + sV_curr.subview(smem_offset).data_ptr().store_swizzled( + zero_vec, + alignment=16, + swizzle=_qkv_smem_swizzle(self.cfg), + ) + + # v_desc is called only after this stage's skv.wait(), which makes + # the TMA writes visible. Converge the one MMA warp after its + # generic stores, then publish them to the async SMEM proxy before + # tcgen05 consumes the descriptor. + cute.arch.sync_warp() + prims.fence_proxy( + kind=prims.Proxy.ASYNC_SHARED, + space=prims.SharedSpace.shared_cta, + ) + + def _build_v_descriptor(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Build the current stage's V descriptor after any required clear.""" + smem_stage_elements = self.cfg.tma_copy_kv_elements + sK_curr = self.sK_array.subview(stage_info.stage_idx * smem_stage_elements) + leading_byte_offset, stride_byte_offset = _pv_smem_desc_offsets(self.cfg) + return prims.Tcgen05SmemDesc.build( + sK_curr, + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_layout(self.cfg), + ) + + @consumer_work(returns=desc_v_base) + @cute.jit + def v_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Build a V SMEM descriptor that needs no paged-tail clear.""" + return self._build_v_descriptor(stage_info) + + @consumer_work(returns=desc_v_base) + @cute.jit + def v_desc_paged( + self, + stage_info: StageInfo, + *, + section: cutlass.Constexpr[FmhaStage], + tile_offset: cutlass.Constexpr[int] = 0, + seqlen_k: Int32, + kv_tile_start: Int32, + ) -> prims.Tcgen05SmemDesc: + """Clear invalid paged-V rows, then build its SMEM descriptor.""" + self._zero_paged_v_tail( + stage_info, + section=section, + tile_offset=tile_offset, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + ) + return self._build_v_descriptor(stage_info) + + +# --------------------------------------------------------------------------- +# TmemSPResource -- TMEM S/P ping-pong buffer with UmmaAsync pipeline +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class TmemSPResource(MemoryResource): + """TMEM S/P ping-pong buffer with UmmaAsync pipeline. + + Producer: MMA warp writes S = Q*K scores, then reads P for P*V. + Consumer: Softmax warp reads S, computes P = softmax(S), writes P back. + + Self-edge in dependency graph enables ping-pong validation: + MMA acquires -> writes S -> commits -> Softmax waits -> reads S, + writes P -> releases -> MMA re-acquires -> reads P for P*V -> commits. + """ + + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + tmem_s_offset: Constexpr[int] = field(init=False, default=None) + tmem_p_offset: Constexpr[int] = field(init=False, default=None) + # 0 for SP0 (uses Q0), 1 for SP1 (uses Q1). + q_half: Constexpr[int] = 0 + enable_early_tile_sum: Constexpr[bool] = False + q_offset_default: int | Int32 = field(init=False, default=0) + cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) + cum_seqlen_k: cute.Tensor | None = field(init=False, default=None) + seq_lens_kv: cute.Pointer | None = field(init=False, default=None) + variable_window_token_starts: cute.Tensor | None = field(init=False, default=None) + variable_window_token_ends: cute.Tensor | None = field(init=False, default=None) + variable_window_cta_starts: cute.Tensor | None = field(init=False, default=None) + variable_window_q_stride: int | Int32 = field(init=False, default=0) + scale_softmax_log2: cute.Tensor | None = field(init=False, default=None) + tmem_addr_cached: TmemAddr | None = field(init=False, default=None) + # Precomputed TMEM pointers/addresses (set by auxiliary work). Avoids + # per-iteration inttoptr + address math. + # MMA warp pointer for QK to S. + tmem_ptr_s_cached: TmemPtr | None = field(init=False, default=None) + # Softmax warp per-warp S address. + tmem_s_addr_cached: TmemAddr | None = field(init=False, default=None) + # Softmax warp per-warp P address. + tmem_p_addr_cached: TmemAddr | None = field(init=False, default=None) + _alloc: Constexpr[Optional[TmemAllocation]] = field(init=False, default=None) + old_row_max: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + row_max: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + row_sum: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + p_chunk: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + q_offset: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + seqlen_k: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + variable_window_start: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + variable_window_end: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def __init__( + self, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + tmem_s_offset: int, + tmem_p_offset: int, + q_half: int = 0, + q_offset: int | Int32 = 0, + cum_seqlen_q: cute.Tensor | None = None, + cum_seqlen_k: cute.Tensor | None = None, + seq_lens_kv: cute.Pointer | None = None, + variable_window_token_starts: cute.Tensor | None = None, + variable_window_token_ends: cute.Tensor | None = None, + variable_window_cta_starts: cute.Tensor | None = None, + variable_window_q_stride: int | Int32 = 0, + scale_softmax_log2: cute.Tensor | None = None, + **kwargs: Any, + ) -> None: + """Bind S/P TMEM offsets, Q peer index, and optional varlen metadata.""" + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.tmem_s_offset = tmem_s_offset + self.tmem_p_offset = tmem_p_offset + self.q_half = q_half + self.enable_early_tile_sum = cfg.enable_early_tile_sum + self.q_offset_default = q_offset + self.cum_seqlen_q = cum_seqlen_q + self.cum_seqlen_k = cum_seqlen_k + self.seq_lens_kv = seq_lens_kv + self.variable_window_token_starts = variable_window_token_starts + self.variable_window_token_ends = variable_window_token_ends + self.variable_window_cta_starts = variable_window_cta_starts + self.variable_window_q_stride = variable_window_q_stride + self.scale_softmax_log2 = scale_softmax_log2 + self._alloc = TmemAllocation( + f"tmem_sp_q{q_half}", + cfg.qk_mma_tiler[1] * cfg.mma_softmax_stage, + ) + self.tmem_addr_cached = Int32(0) + self.tmem_ptr_s_cached = _placeholder_tmem_ptr() + self.tmem_s_addr_cached = Int32(0) + self.tmem_p_addr_cached = Int32(0) + self.old_row_max = TaskLocalVariable( + dtype=Float32, + default=Float32(-Float32.inf), + docs="Softmax row maximum from the previous K/V tile.", + ) + self.row_max = TaskLocalVariable( + dtype=Float32, + default=Float32(-Float32.inf), + docs="Softmax row maximum for the current K/V tile.", + ) + self.row_sum = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="Accumulated softmax denominator for the current row.", + ) + if self.enable_early_tile_sum: + self.p_chunk = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="FP32 sum of the current probability tile.", + ) + else: + self.p_chunk = TaskLocalVariable( + dtype=list, + default_factory=lambda: _placeholder_softmax_chunks(cfg), + docs="P fragments retained for post-release row-sum reduction.", + ) + self.q_offset = TaskLocalVariable( + dtype=Int32, + default=Int32(self.q_offset_default), + docs="Causal Q/K sequence offset for the current work tile.", + ) + self.seqlen_k = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Request-local K/V sequence length for packed dense masking.", + ) + self.variable_window_start = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Inclusive first K position for this Q row.", + ) + self.variable_window_end = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Inclusive last K position for this Q row.", + ) + self.scale_softmax_log2_value = TaskLocalVariable( + dtype=Float32, + # Placeholder before load_scale_softmax_log2 reads the runtime tensor. + default=Float32(0.0), + docs="Softmax scale cached from the runtime scale tensor.", + ) + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Return the TMEM allocation for this S/P ping-pong resource.""" + return [self._alloc] + + @property + def loop_offset_sensitive(self) -> bool: + """Return true because MMA and masking decisions use loop_offset.""" + # producer_work and head-paired masking use loop_offset for K tile indices. + return True + + @property + def uses_left_window_loop_mask(self) -> bool: + """Return whether loop iterations need the left sliding-window mask.""" + return self.cfg.kv_tile_start_window_size_left > 0 + + @property + def uses_varlen_loop_right_mask(self) -> bool: + """Return whether loop iterations need mixed-varlen right masking.""" + return self.cfg.head_paired and self.cfg.has_varlen and self.cfg.has_q_offset + + @property + def uses_varlen_q_offset_cache(self) -> bool: + """Return whether masks need a per-work-tile varlen Q/K offset.""" + return self.cfg.has_varlen and self.cfg.has_q_offset + + @property + def uses_variable_window(self) -> bool: + """Return whether softmax consumes explicit packed-Q row bounds.""" + return self.cfg.has_variable_window + + @property + def uses_fixed_dense_k_tail_mask(self) -> bool: + """Return whether fixed dense attention has a partial final K/V tile.""" + return ( + not self.cfg.is_causal + and not self.cfg.has_varlen + and self.cfg.fixed_dense_k_tail > 0 + ) + + @property + def uses_packed_dense_k_mask(self) -> bool: + """Return whether packed or paged dense attention needs local K bounds.""" + return ( + self.cfg.has_varlen + and not self.cfg.is_causal + and self.cfg.packed_dense_k_mask + ) + + @property + def uses_query_paired_q_offset_loop_mask(self) -> bool: + """Return whether query-paired loop iterations need q-offset masking. + + Mixed packed batches use a request-local domain, but paired-tail + alignment and partial Q tiles can conservatively retain a K/V tile + that crosses the causal right edge. Either peer can therefore need + the right mask inside LOOP rather than only in peer0 TAIL. A uniform + tile-aligned shift preserves the ordinary tail placement and compiles + this per-iteration mask away. + """ + return ( + self.cfg.has_q_offset + and not self.cfg.head_paired + and not self.cfg.has_tile_aligned_uniform_q_offset + ) + + @property + def uses_head_paired_causal_tail_mask(self) -> bool: + """Return whether TAIL should use head-paired causal masking.""" + return self.cfg.is_causal and self.cfg.head_paired + + @property + def needs_window_tail_left_mask(self) -> bool: + """Return whether a sliding-window TAIL can cross its left edge. + + For fixed equal-length 128x128 tiling, a window of at least M-1 + tokens places the entire final causal tile on or to the right of the + left bound. Packed and bottom-right-offset inputs retain the general + two-sided mask because their runtime tile origin can shift. + """ + return self.cfg.window_size_left > 0 and ( + self.cfg.has_varlen + or self.cfg.has_q_offset + or self.cfg.q_tile_m != self.cfg.kv_tile_n + or self.cfg.window_size_left < self.cfg.q_tile_m - 1 + ) + + @property + def uses_query_paired_causal_tail_mask(self) -> bool: + """Return whether TAIL should use query-paired causal masking.""" + return self.cfg.is_causal and not self.cfg.head_paired and self.q_half == 0 + + @property + def uses_query_paired_invalid_tail(self) -> bool: + """Return whether peer0 needs the extra wholly-invalid tail slot.""" + return self.cfg.skip_causal_invalid_peer0 and self.q_half == 0 + + @cute.jit + def _stage_col_offset(self, stage_info: StageInfo) -> Int32 | int: + """Return the TMEM column offset for a pipelined S/P stage.""" + stage_col_offset = Int32(0) + if cutlass.const_expr(self.cfg.mma_softmax_stage > 1): + stage_col_offset = stage_info.stage_idx * self.cfg.qk_mma_tiler[1] + return stage_col_offset + + @producer_work + @cute.jit + def qk_mma( + self, + stage_info: StageInfo, + *, + desc_q_base: prims.Tcgen05SmemDesc, + desc_k_base: prims.Tcgen05SmemDesc, + section: cutlass.Constexpr[FmhaStage], + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """QK MMA: compute Q*K -> S in TMEM. + + The schedule aliases either desc_q0_base or desc_q1_base into the + logical desc_q_base producer arg based on which SP instance is being + driven. + + In causal mode with no Q right offset, skips QK0→S0 MMA in the last + LOOP iteration, since Softmax0's domain is N-2 but MMA's domain is N-1. + The task domain pads partial final CTAs so this slot is always outside + peer0's causal reach. + """ + skip_qk0_invalid = False + if cutlass.const_expr(self.cfg.skip_causal_invalid_peer0 and self.q_half == 0): + if cutlass.const_expr(section == FmhaStage.Loop): + if not is_tail: + skip_qk0_invalid = stage_info.loop_offset == ( + stage_info.loop_end - 1 + ) + + if not skip_qk0_invalid: + tmem_ptr_s = self.tmem_ptr_s_cached.subview( + self._stage_col_offset(stage_info) + ) + + if cutlass.const_expr(self.cfg.q_dtype.width == 8): + mma_kind = prims.Tcgen05MMAKind.F8F6F4 + # E4M3 operands use the Float16 encoding handle. + ab_format = cutlass.Float16 + else: + mma_kind = prims.Tcgen05MMAKind.F16 + if cutlass.const_expr(self.cfg.q_dtype == cutlass.BFloat16): + ab_format = cutlass.BFloat16 + else: + ab_format = cutlass.Float16 + + idesc_qk = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=ab_format, + b_dtype=ab_format, + n_dim=self.cfg.qk_mma_tiler[1], + m_dim=self.cfg.qk_mma_tiler[0], + ) + + k_dim_per_mma = 16 + if cutlass.const_expr(self.cfg.q_dtype.width != 16): + k_dim_per_mma = 32 + num_kphases_qk = self.cfg.qk_mma_tiler[2] // k_dim_per_mma + inc_bytes_qk = k_dim_per_mma * self.cfg.q_dtype.width // 8 + + num_kphases_per_tma = num_kphases_qk // self.cfg.tma_copy_qkv_iters + chunk_bytes_qk = inc_bytes_qk * num_kphases_per_tma + if cutlass.const_expr(self.cfg.tma_copy_qkv_iters != 1): + chunk_bytes_qk = ( + self.cfg.tma_copy_kv_bytes // self.cfg.tma_copy_kv_stage_iters + ) + num_tma_iters_qk = self.cfg.tma_copy_qkv_iters + if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): + num_tma_iters_qk = self.cfg.tma_copy_kv_stage_iters + + # Prevent LLVM from rematerializing descriptor + # computations inside each elect_sync basic block. + # Without this, NVPTX recomputes shr+and+cvt+or from + # __dynamic_shmem__0 inside every elect BB (~5 extra + # instructions per MMA call). + desc_q_base_ = freeze_smem_descriptor(desc_q_base) + desc_k_base_ = freeze_smem_descriptor(desc_k_base) + + scale_d = False + if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): + scale_d = head_dim_stage_idx != 0 + for tma_iter in cutlass.range_constexpr(num_tma_iters_qk): + q_tma_iter = head_dim_stage_idx * num_tma_iters_qk + tma_iter + q_tma_iter_offset = chunk_bytes_qk * q_tma_iter + k_tma_iter_offset = chunk_bytes_qk * tma_iter + for k_idx in cutlass.range_constexpr(num_kphases_per_tma): + local_increment = inc_bytes_qk * k_idx + dq = desc_q_base_ + ((local_increment + q_tma_iter_offset) >> 4) + dk = desc_k_base_ + ((local_increment + k_tma_iter_offset) >> 4) + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + prims.CTAGroup.CTA_1, + tmem_ptr_s, + dq, + dk, + idesc_qk, + scale_d, + ) + scale_d = True + + @producer_work + @cute.jit + def p_read(self, stage_info: StageInfo) -> None: + """P-read sync: no-op. SP handle held from QK, consumed by softmax.""" + pass + + @cute.jit + def _init_function_state(self, stage_info: StageInfo) -> None: + """Precompute TMEM pointers/addresses (once, before persistent loop). + + Runs on all warps via init_variables, after tmem_addr_cached is set. + Only the MMA-warp pointer is computed here (needed ungated). + Softmax-warp addresses are deferred to per-work-tile auxiliary work + so they are computed after setmaxnreg and avoid crossing the register + budget boundary. The fields are initialized to Int32(0) here so the + DSL sees a consistent type structure before the scf.while loop. + + Emits the softmax-side state variables (old_row_max, row_max, + row_sum, p_chunk, q_offset) consumed by Softmax tasks; producer-side + desc_q_base / desc_k_base slots are auto-mirrored from + SmemQ / SmemKV by Task.init_variables (with explicit aliasing + from desc_q0_base / desc_q1_base). + """ + # MMA warp: tmem_ptr_s for QK→S producer_work + self.tmem_ptr_s_cached = prims.make_tmem_ptr( + self.tmem_addr_cached, cutlass.Int8 + ).subview(self.tmem_s_offset) + # Initialize to establish DSL type; real values are set per work tile. + self.tmem_s_addr_cached = Int32(0) + self.tmem_p_addr_cached = Int32(0) + _ = stage_info + + @cute.jit + def _default_p_chunk(self) -> SoftmaxRowSumContribution: + if cutlass.const_expr(self.enable_early_tile_sum): + return Float32(0.0) + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + + # PERF NOTE: These P chunk vectors become iter_args of the scf.while + # persistent loop. The MLIR compiler materializes zero-initialization + # HEAD code (~128 add.rn.f32x2 instructions adding 0+0) and + # TAIL finalization code (runs once after LOOP exit). The K-loop + # body instruction count is unaffected — identical with or without + # these iter_args. The HEAD/TAIL overhead may affect performance + # through i-cache pressure (~+3% PTX footprint), register allocation + # changes (ptxas sees more live values at scf.while boundary), and + # pipeline warm-up timing shifts. + p_chunk = [] + for _chunk_idx in cutlass.range_constexpr(num_chunks): + zeros = tuple(self.cfg.qk_acc_dtype(0.0) for _ in range(tmem_x)) + p_chunk.append(cutlass.Vector.from_elements(zeros, self.cfg.qk_acc_dtype)) + return p_chunk + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_state(self, stage_info: StageInfo) -> None: + self._init_function_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_softmax_state_early(self, stage_info: StageInfo) -> None: + """Initialize softmax TMEM state without a function-lifetime P value.""" + self._init_function_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_chunk) + @cute.jit + def init_softmax_state(self, stage_info: StageInfo) -> SoftmaxRowSumContribution: + self._init_function_state(stage_info) + return self._default_p_chunk() + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns="scale_softmax_log2_value") + @cute.jit + def load_scale_softmax_log2(self, stage_info: StageInfo) -> Float32: + """Load the runtime softmax scale once before the K/V loop.""" + _ = stage_info + if cutlass.const_expr(self.scale_softmax_log2 is None): + # Safe fallback for validation-only resource construction. + return Float32(0.0) + return self.scale_softmax_log2[0] + + @cute.jit + def _init_work_tile_state(self, stage_info: StageInfo) -> None: + """Reset softmax state and recompute per-warp TMEM addresses each tile. + + Softmax-warp addresses are computed here (inside the persistent loop, + after setmaxnreg) to avoid spilling them across the register-budget + boundary in ungated HEAD. The returned q_offset defaults to the + uniform kernel argument; varlen causal masks overwrite it once per + work tile via cache_q_offset(). + """ + num_softmax_warps = 4 + warp_id_in_sg = cute.arch.warp_idx() % num_softmax_warps + tmem_raw_addr = self.tmem_addr_cached + tmem_base_row = tmem_raw_addr >> 16 + tmem_base_col = tmem_raw_addr & Int32(0xFFFF) + row_id = tmem_base_row + warp_id_in_sg * cute.arch.WARP_SIZE + self.tmem_s_addr_cached = (row_id << 16) | (tmem_base_col + self.tmem_s_offset) + self.tmem_p_addr_cached = (row_id << 16) | (tmem_base_col + self.tmem_p_offset) + _ = stage_info + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(old_row_max, row_max, row_sum, q_offset), + ) + @cute.jit + def init_softmax_work_tile_state( + self, stage_info: StageInfo + ) -> tuple[Float32, Float32, Float32, Int32]: + self._init_work_tile_state(stage_info) + return ( + Float32(-Float32.inf), + Float32(-Float32.inf), + Float32(0.0), + Int32(self.q_offset_default), + ) + + @cute.jit + def _varlen_batch_coord(self, stage_info: StageInfo) -> Int32: + """Return the batch coordinate for the active tile-order policy.""" + _, _, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + return batch_coord + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=q_offset) + @cute.jit + def cache_q_offset(self, stage_info: StageInfo) -> Int32: + """Cache the per-work-tile causal Q/K sequence offset for masks. + + Mixed-varlen batches cannot use the uniform kernel q_offset because + each batch can have a different S_kv - S_q. This pre-wait hook runs + once in the softmax task HEAD, before the K/V loop, so loop and tail + masks reuse the cached offset instead of rereading the request metadata. + """ + if cutlass.const_expr(self.cfg.has_uniform_varlen): + return Int32(self.cfg.uniform_seq_len_k - self.cfg.uniform_seq_len_q) + batch_coord = self._varlen_batch_coord(stage_info) + if cutlass.const_expr(self.cfg.has_uniform_varlen): + seqlen_q = Int32(self.cfg.uniform_seq_len_q) + else: + cuseqlen_q = Int32(self.cum_seqlen_q[batch_coord]) + seqlen_q = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) - cuseqlen_q + if cutlass.const_expr(self.cfg.use_paged_kv): + seqlen_k = Int32(self.seq_lens_kv[batch_coord]) + elif cutlass.const_expr(self.cfg.has_uniform_varlen): + seqlen_k = Int32(self.cfg.uniform_seq_len_k) + else: + cuseqlen_k = Int32(self.cum_seqlen_k[batch_coord]) + seqlen_k = Int32(self.cum_seqlen_k[batch_coord + Int32(1)]) - cuseqlen_k + return seqlen_k - seqlen_q + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=seqlen_k) + @cute.jit + def cache_seqlen_k(self, stage_info: StageInfo) -> Int32: + """Cache the request-local K/V extent once per work tile.""" + if cutlass.const_expr(self.cfg.has_uniform_varlen): + return Int32(self.cfg.uniform_seq_len_k) + if cutlass.const_expr(self.cfg.use_paged_kv): + batch_coord = self._varlen_batch_coord(stage_info) + return Int32(self.seq_lens_kv[batch_coord]) + batch_coord = self._varlen_batch_coord(stage_info) + cuseqlen_k = Int32(self.cum_seqlen_k[batch_coord]) + return Int32(self.cum_seqlen_k[batch_coord + Int32(1)]) - cuseqlen_k + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(variable_window_start, variable_window_end), + ) + @cute.jit + def cache_variable_window_bounds( + self, stage_info: StageInfo + ) -> tuple[Int32, Int32]: + """Load this lane's bounds relative to the CTA's first K/V tile.""" + seq_coord, _, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + warp_id_in_sg = cute.arch.warp_idx() % 4 + row_in_tile = warp_id_in_sg * cute.arch.WARP_SIZE + cute.arch.lane_idx() + local_q = ( + seq_coord * self.cfg.q_tile_m * self.cfg.work_tile_q_seq_tiles + + self.q_half * self.cfg.peer_q_seq_tile_stride * self.cfg.q_tile_m + + row_in_tile + ) + local_q = cute.math.min( + local_q, + self.variable_window_q_stride - Int32(1), + ) + packed_q = batch_coord * self.variable_window_q_stride + local_q + min_window_start = variable_window_cta_min_start( + self.variable_window_cta_starts, + batch_coord=batch_coord, + seq_coord=seq_coord, + q_stride=self.variable_window_q_stride, + tile_size_q=self.cfg.cta_tiler[0], + ) + kv_base = (min_window_start // self.cfg.kv_tile_n) * self.cfg.kv_tile_n + return ( + Int32(self.variable_window_token_starts[packed_q]) - kv_base, + Int32(self.variable_window_token_ends[packed_q]) - kv_base, + ) + + @cute.jit + def _load_s_chunks(self, stage_info: StageInfo) -> SoftmaxChunks: + """Load ALL S chunks from TMEM into register vectors.""" + tmem_s_addr = self.tmem_s_addr_cached + self._stage_col_offset(stage_info) + tmem_shape = "32x32b" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + s_data = [None] * num_chunks + for chunk_idx in cutlass.range_constexpr(num_chunks): + _chunk = cutlass.Array(self.cfg.qk_acc_dtype, tmem_x) + _chunk[0:tmem_x] = prims.tcgen05_ld( + tmem_shape, + prims.make_tmem_ptr( + tmem_s_addr + chunk_idx * tmem_x, self.cfg.qk_acc_dtype + ), + num=tmem_x, + ) + s_data[chunk_idx] = _chunk + cute.arch.fence_view_async_tmem_load() + for chunk_idx in cutlass.range_constexpr(num_chunks): + s_data[chunk_idx] = s_data[chunk_idx][0:tmem_x] + return s_data + + @cute.jit + def _reduce_row_max( + self, + s_data: SoftmaxChunks, + row_max: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Reduce per-chunk maximums into row_max, stash s_data.""" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + old_row_max = row_max + if cutlass.const_expr( + self.cfg.uses_d128_fp8_softmax_cadence + or self.cfg.uses_d256_fp8_softmax_cadence + ): + max_0 = row_max + max_1 = row_max + max_2 = row_max + max_3 = row_max + for chunk_idx in cutlass.range_constexpr(num_chunks): + for elem_idx in cutlass.range_constexpr(0, tmem_x, 4): + max_0 = cute.math.max(max_0, s_data[chunk_idx][elem_idx], ftz=True) + max_1 = cute.math.max( + max_1, s_data[chunk_idx][elem_idx + 1], ftz=True + ) + max_2 = cute.math.max( + max_2, s_data[chunk_idx][elem_idx + 2], ftz=True + ) + max_3 = cute.math.max( + max_3, s_data[chunk_idx][elem_idx + 3], ftz=True + ) + max_0 = cute.math.max(max_0, max_2, ftz=True) + max_1 = cute.math.max(max_1, max_3, ftz=True) + row_max = cute.math.max(max_0, max_1, ftz=True) + else: + row_values: tuple[Any, ...] = () + for chunk_idx in cutlass.range_constexpr(num_chunks): + for elem_idx in cutlass.range_constexpr(tmem_x): + row_values += (s_data[chunk_idx][elem_idx],) + row_vector = cutlass.Vector.from_elements(row_values, self.cfg.qk_acc_dtype) + tile_row_max = row_vector.reduce("max") + row_max = cute.math.max(row_max, tile_row_max) + _tmem_sp_sdata[id(self)] = s_data + row_max_safe = row_max + if row_max == -Float32.inf: + row_max_safe = Float32(0.0) + return old_row_max, row_max_safe + + @cute.jit + def _exp2_p_store( + self, + stage_col_offset: TmemAddr, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + ) -> SoftmaxRowSumContribution: + """Apply exp2 softmax P, fold the PV P scale, and store P to TMEM.""" + tmem_p_addr = self.tmem_p_addr_cached + stage_col_offset + tmem_shape = "32x32b" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + p_packing_ratio = self.cfg.qk_acc_dtype.width // self.cfg.v_dtype.width + scale = scale_softmax_log2 + if cutlass.const_expr(self.cfg.uses_d256_fp8_softmax_cadence): + return self._exp2_p_store_d256_fp8_cadence( + tmem_p_addr, + row_max, + scale, + ) + if cutlass.const_expr(self.cfg.uses_d128_fp8_softmax_cadence): + return self._exp2_p_store_d128_fp8_cadence( + tmem_p_addr, + row_max, + scale, + ) + p_data_f32 = cutlass.Array(self.cfg.qk_acc_dtype, tmem_x, alignment=16) + p_data_packed = cutlass.Array( + p_data_f32.data_ptr(), + shape=(tmem_x * p_packing_ratio,), + dtype=self.cfg.v_dtype, + ) + p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) + minus_row_max_scale = (Float32(0.0) - row_max) * scale + p_scale_log2 + s_data = _tmem_sp_sdata.pop(id(self)) + if cutlass.const_expr(self.enable_early_tile_sum): + # Keep four independent scalar dependency chains while expressing + # them as two packed float2 values. The explicit packed primitive + # lowers to FADD2 for D128 instead of two scalar FADDs per pair. + local_sum_pair_0 = (Float32(0.0), Float32(0.0)) + local_sum_pair_1 = (Float32(0.0), Float32(0.0)) + for chunk_idx in cutlass.range_constexpr(num_chunks): + p_vals = () + for elem_idx in cutlass.range_constexpr(0, tmem_x, 2): + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) + p0 = cute.math.exp2(fma_pair[0], fastmath=True) + p1 = cute.math.exp2(fma_pair[1], fastmath=True) + if cutlass.const_expr(self.enable_early_tile_sum): + pair_idx = chunk_idx * (tmem_x // 2) + elem_idx // 2 + if cutlass.const_expr(pair_idx % 2 == 0): + local_sum_pair_0 = cute.arch.add_packed_f32x2( + local_sum_pair_0, + (p0, p1), + rnd="rn", + ftz=False, + ) + else: + local_sum_pair_1 = cute.arch.add_packed_f32x2( + local_sum_pair_1, + (p0, p1), + rnd="rn", + ftz=False, + ) + p_vals += (p0, p1) + s_data[chunk_idx] = cutlass.Vector.from_elements( + p_vals, self.cfg.qk_acc_dtype + ) + use_fused_d128_fp8x4_pack = ( + not self.cfg.stage_kv_by_head_dim + and self.cfg.v_dtype == cutlass.Float8E4M3FN + ) + for pair_idx in cutlass.range_constexpr(num_chunks // p_packing_ratio): + if cutlass.const_expr(use_fused_d128_fp8x4_pack): + # Match the handwritten D128 pack: merge both FP8x2 + # conversions in one side-effecting block so ptxas can retain + # the 32-bit word without a PRMT between temporary vectors. + packed_words: tuple[Any, ...] = () + for word_idx in cutlass.range_constexpr(tmem_x): + flat_idx = word_idx * 4 + chunk_idx = pair_idx * p_packing_ratio + flat_idx // tmem_x + elem_idx = flat_idx % tmem_x + packed_word = _pack_float4_to_fp8_e4m3( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + s_data[chunk_idx][elem_idx + 2], + s_data[chunk_idx][elem_idx + 3], + ) + packed_words += (packed_word,) + store_fragment = cutlass.Vector.from_elements(packed_words, Int32) + else: + for slice_idx in cutlass.range_constexpr(p_packing_ratio): + chunk_idx = pair_idx * p_packing_ratio + slice_idx + p_chunk_dtype = s_data[chunk_idx].to(self.cfg.v_dtype) + if cutlass.const_expr(self.cfg.v_dtype.width == 8): + p_chunk_i8 = p_chunk_dtype.bitcast(cutlass.Int8) + p_data_packed[slice_idx * tmem_x : tmem_x] = p_chunk_i8 + else: + p_data_packed[slice_idx * tmem_x : tmem_x] = p_chunk_dtype + store_fragment = p_data_f32[0:tmem_x] + prims.tcgen05_st( + tmem_shape, + prims.make_tmem_ptr(tmem_p_addr + pair_idx * tmem_x, cutlass.Int8), + store_fragment, + ) + if cutlass.const_expr(self.enable_early_tile_sum): + local_sum_pair = cute.arch.add_packed_f32x2( + local_sum_pair_0, + local_sum_pair_1, + rnd="rn", + ftz=False, + ) + tile_sum = local_sum_pair[0] + local_sum_pair[1] + if cutlass.const_expr( + self.enable_early_tile_sum or self.cfg.has_tmem_p_pipeline + ): + # Publish TMEM store through the task-pipeline barrier without a blocking + # store wait. The P-ready consumer pipeline orders the UMMA warp + # after every store in the staged D256 path. + cute.arch.fence_view_async_tmem_store() + else: + # Preserve the legacy publication sequence for paths that retain + # P fragments until after SP release. + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + if cutlass.const_expr(self.enable_early_tile_sum): + return tile_sum + result = [] + for chunk_idx in cutlass.range_constexpr(num_chunks): + result.append(s_data[chunk_idx]) + return result + + @cute.jit + def _exp2_p_store_d128_fp8_cadence( + self, + tmem_p_addr: TmemAddr, + row_max: SoftmaxScalar, + scale: SoftmaxScalar, + ) -> Float32: + """Use TRT's D128 FP8 softmax arithmetic cadence. + + Prefetch eight FMA values and retire FP8 conversions and scalar sums + eight values behind EXP2. Four independent sum chains preserve the + dependency depth of the reference implementation. + """ + tmem_x = self.cfg.tmem_x_load_s + num_values = self.cfg.qk_mma_tiler[1] + fma_lookahead = 8 + retirement_delay = 8 + store_group_words = 4 + words_per_chunk = 2 * store_group_words + p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) + minus_row_max_scale = (Float32(0.0) - row_max) * scale + p_scale_log2 + s_data = _tmem_sp_sdata.pop(id(self)) + local_sum_chains = cutlass.Array( + Float32, + 4, + space=cutlass.AddressSpace.rmem, + ) + for chain_idx in cutlass.range_constexpr(4): + local_sum_chains[chain_idx] = Float32(0.0) + + fma_ring = cute.make_rmem_tensor((fma_lookahead,), Float32) + exp_ring = cute.make_rmem_tensor((retirement_delay,), Float32) + p_output_words_lo = cute.make_rmem_tensor((store_group_words,), Int32) + p_output_words_hi = cute.make_rmem_tensor((store_group_words,), Int32) + num_chunks = num_values // tmem_x + for chunk_idx in cutlass.range_constexpr(num_chunks): + for local_idx in cutlass.range_constexpr(0, fma_lookahead, 2): + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][local_idx], + s_data[chunk_idx][local_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) + fma_ring[local_idx] = fma_pair[0] + fma_ring[local_idx + 1] = fma_pair[1] + + for local_idx in cutlass.range_constexpr(0, tmem_x, 2): + fma_idx = local_idx % fma_lookahead + exp_idx = local_idx % retirement_delay + if cutlass.const_expr( + local_idx >= retirement_delay + and (local_idx - retirement_delay) % 4 == 0 + ): + delayed_idx = local_idx - retirement_delay + word_idx = delayed_idx // 4 + if cutlass.const_expr(word_idx < store_group_words): + p_output_words_lo[word_idx] = _pack_float4_to_fp8_e4m3( + exp_ring[exp_idx], + exp_ring[exp_idx + 1], + exp_ring[exp_idx + 2], + exp_ring[exp_idx + 3], + ) + else: + p_output_words_hi[word_idx - store_group_words] = ( + _pack_float4_to_fp8_e4m3( + exp_ring[exp_idx], + exp_ring[exp_idx + 1], + exp_ring[exp_idx + 2], + exp_ring[exp_idx + 3], + ) + ) + + p_0 = cute.math.exp2(fma_ring[fma_idx], fastmath=True) + # Preserve the odd value from the current pair before the + # circular lookahead slot is refilled with the future pair. + fma_1 = fma_ring[fma_idx + 1] + if cutlass.const_expr(local_idx + fma_lookahead < tmem_x): + future_idx = local_idx + fma_lookahead + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][future_idx], + s_data[chunk_idx][future_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) + fma_ring[fma_idx] = fma_pair[0] + fma_ring[fma_idx + 1] = fma_pair[1] + p_1 = cute.math.exp2(fma_1, fastmath=True) + + if cutlass.const_expr(local_idx >= retirement_delay): + delayed_idx = local_idx - retirement_delay + chain_base = ((delayed_idx // 2) % 2) * 2 + local_sum_chains[chain_base] += exp_ring[exp_idx] + local_sum_chains[chain_base + 1] += exp_ring[exp_idx + 1] + exp_ring[exp_idx] = p_0 + exp_ring[exp_idx + 1] = p_1 + + for delayed_idx in cutlass.range_constexpr( + tmem_x - retirement_delay, + tmem_x, + 4, + ): + exp_idx = delayed_idx % retirement_delay + word_idx = delayed_idx // 4 + if cutlass.const_expr(word_idx < store_group_words): + p_output_words_lo[word_idx] = _pack_float4_to_fp8_e4m3( + exp_ring[exp_idx], + exp_ring[exp_idx + 1], + exp_ring[exp_idx + 2], + exp_ring[exp_idx + 3], + ) + else: + p_output_words_hi[word_idx - store_group_words] = ( + _pack_float4_to_fp8_e4m3( + exp_ring[exp_idx], + exp_ring[exp_idx + 1], + exp_ring[exp_idx + 2], + exp_ring[exp_idx + 3], + ) + ) + for delayed_idx in cutlass.range_constexpr( + tmem_x - retirement_delay, + tmem_x, + 2, + ): + exp_idx = delayed_idx % retirement_delay + chain_base = ((delayed_idx // 2) % 2) * 2 + local_sum_chains[chain_base] += exp_ring[exp_idx] + local_sum_chains[chain_base + 1] += exp_ring[exp_idx + 1] + + prims.tcgen05_st( + "32x32b", + prims.make_tmem_ptr( + tmem_p_addr + chunk_idx * words_per_chunk, + cutlass.Int8, + ), + p_output_words_lo.load(), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + prims.tcgen05_st( + "32x32b", + prims.make_tmem_ptr( + tmem_p_addr + chunk_idx * words_per_chunk + store_group_words, + cutlass.Int8, + ), + p_output_words_hi.load(), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + + local_sum_pair = cute.arch.add_packed_f32x2( + (local_sum_chains[0], local_sum_chains[1]), + (local_sum_chains[2], local_sum_chains[3]), + rnd="rn", + ftz=False, + ) + return local_sum_pair[0] + local_sum_pair[1] + + @cute.jit + def _exp2_p_store_d256_fp8_cadence( + self, + tmem_p_addr: TmemAddr, + row_max: SoftmaxScalar, + scale: SoftmaxScalar, + ) -> Float32: + """Interleave D256 FP8 EXP2, conversion, and tile-sum retirement. + + The eight-value lookahead mirrors the handwritten D256 FMHA cadence + while retaining immutable SSA values. This avoids the long all-EXP2 + burst without reintroducing the mutable cross-typed fragment that was + nondeterministic under Task Scheduling control flow. + """ + tmem_x = self.cfg.tmem_x_load_s + num_values = self.cfg.qk_mma_tiler[1] + p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) + minus_row_max_scale = (Float32(0.0) - row_max) * scale + p_scale_log2 + s_data = _tmem_sp_sdata.pop(id(self)) + + fma_values: tuple[Any, ...] = () + for flat_idx in cutlass.range_constexpr(0, 8, 2): + chunk_idx = flat_idx // tmem_x + elem_idx = flat_idx % tmem_x + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) + fma_values += (fma_pair[0], fma_pair[1]) + + p_values: tuple[Any, ...] = () + packed_words: tuple[Any, ...] = () + local_sum_pair_0 = (Float32(0.0), Float32(0.0)) + local_sum_pair_1 = (Float32(0.0), Float32(0.0)) + use_two_sum_pairs = not self.cfg.stage_kv_by_head_dim + for flat_idx in cutlass.range_constexpr(0, num_values, 2): + if cutlass.const_expr(flat_idx >= 8): + delayed_idx = flat_idx - 8 + if cutlass.const_expr(flat_idx % 4 == 0): + packed_word = _pack_float4_to_fp8_e4m3( + p_values[delayed_idx], + p_values[delayed_idx + 1], + p_values[delayed_idx + 2], + p_values[delayed_idx + 3], + ) + packed_words += (packed_word,) + if cutlass.const_expr( + use_two_sum_pairs and (delayed_idx // 2) % 2 == 1 + ): + local_sum_pair_1 = cute.arch.add_packed_f32x2( + local_sum_pair_1, + (p_values[delayed_idx], p_values[delayed_idx + 1]), + rnd="rn", + ftz=False, + ) + else: + local_sum_pair_0 = cute.arch.add_packed_f32x2( + local_sum_pair_0, + (p_values[delayed_idx], p_values[delayed_idx + 1]), + rnd="rn", + ftz=False, + ) + + p0 = cute.math.exp2(fma_values[flat_idx], fastmath=True) + if cutlass.const_expr(flat_idx + 8 < num_values): + future_idx = flat_idx + 8 + chunk_idx = future_idx // tmem_x + elem_idx = future_idx % tmem_x + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) + fma_values += (fma_pair[0], fma_pair[1]) + p1 = cute.math.exp2(fma_values[flat_idx + 1], fastmath=True) + p_values += (p0, p1) + + for delayed_idx in cutlass.range_constexpr(num_values - 8, num_values, 2): + if cutlass.const_expr(delayed_idx % 4 == 0): + packed_word = _pack_float4_to_fp8_e4m3( + p_values[delayed_idx], + p_values[delayed_idx + 1], + p_values[delayed_idx + 2], + p_values[delayed_idx + 3], + ) + packed_words += (packed_word,) + if cutlass.const_expr(use_two_sum_pairs and (delayed_idx // 2) % 2 == 1): + local_sum_pair_1 = cute.arch.add_packed_f32x2( + local_sum_pair_1, + (p_values[delayed_idx], p_values[delayed_idx + 1]), + rnd="rn", + ftz=False, + ) + else: + local_sum_pair_0 = cute.arch.add_packed_f32x2( + local_sum_pair_0, + (p_values[delayed_idx], p_values[delayed_idx + 1]), + rnd="rn", + ftz=False, + ) + + if cutlass.const_expr(use_two_sum_pairs): + local_sum_pair_0 = cute.arch.add_packed_f32x2( + local_sum_pair_0, + local_sum_pair_1, + rnd="rn", + ftz=False, + ) + + store_fragment = cutlass.Vector.from_elements(packed_words, Int32) + prims.tcgen05_st( + "32x32b", + prims.make_tmem_ptr(tmem_p_addr, cutlass.Int8), + store_fragment, + ) + cute.arch.fence_view_async_tmem_store() + return local_sum_pair_0[0] + local_sum_pair_0[1] + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def compute_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Main K-loop stage: load S from TMEM and compute unmasked row_max.""" + s_data = self._load_s_chunks(stage_info) + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def fixed_dense_k_tail_masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Exclude TMA zero-fill lanes in a partial fixed dense K/V tile.""" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + s_data = self._load_s_chunks(stage_info) + + if stage_info.loop_offset == stage_info.loop_end - Int32(1): + neg_inf = cutlass.vector.full( + [tmem_x], + self.cfg.qk_acc_dtype(-Float32.inf), + dtype=self.cfg.qk_acc_dtype, + ) + for chunk_idx in cutlass.range_constexpr(num_chunks): + valid_in_chunk = cute.math.min( + cute.math.max( + Int32(self.cfg.fixed_dense_k_tail) - Int32(chunk_idx * tmem_x), + Int32(0), + ), + Int32(tmem_x), + ) + mask = cutlass.vector.create_mask([tmem_x], [valid_in_chunk]) + s_data[chunk_idx] = cutlass.vector.where( + mask, s_data[chunk_idx], neg_inf + ) + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def packed_dense_k_masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + seqlen_k: Int32, + section: cutlass.Constexpr[FmhaStage], + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Mask scores beyond one packed request's K/V right edge.""" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + s_data = self._load_s_chunks(stage_info) + if cutlass.const_expr(section == FmhaStage.Loop): + kv_tile_idx = stage_info.loop_offset + else: + # Some schedules materialize their final score tile in TAIL. + kv_tile_idx = stage_info.loop_end + kv_base = kv_tile_idx * self.cfg.kv_tile_n + kv_end = kv_base + Int32(self.cfg.kv_tile_n) + if seqlen_k < kv_end: + neg_inf = cutlass.vector.full( + [tmem_x], + self.cfg.qk_acc_dtype(-Float32.inf), + dtype=self.cfg.qk_acc_dtype, + ) + for chunk_idx in cutlass.range_constexpr(num_chunks): + chunk_base = kv_base + Int32(chunk_idx * tmem_x) + valid_in_chunk = cute.math.min( + cute.math.max(seqlen_k - chunk_base, Int32(0)), + Int32(tmem_x), + ) + mask = cutlass.vector.create_mask([tmem_x], [valid_in_chunk]) + s_data[chunk_idx] = cutlass.vector.where( + mask, s_data[chunk_idx], neg_inf + ) + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def variable_window_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + window_start: Int32, + window_end: Int32, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Mask S using inclusive per-row VariableWindow bounds.""" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + s_data = self._load_s_chunks(stage_info) + kv_tile_base = stage_info.loop_offset * self.cfg.kv_tile_n + tile_n = self.cfg.qk_mma_tiler[1] + left_oob = cute.math.min( + cute.math.max(window_start - kv_tile_base, Int32(0)), + Int32(tile_n), + ) + right_valid = cute.math.min( + cute.math.max(window_end + Int32(1) - kv_tile_base, Int32(0)), + Int32(tile_n), + ) + for chunk_idx in cutlass.range_constexpr(num_chunks): + chunk_base = Int32(chunk_idx * tmem_x) + chunk_left = cute.math.min( + cute.math.max(left_oob - chunk_base, Int32(0)), + Int32(tmem_x), + ) + chunk_right = cute.math.min( + cute.math.max(right_valid - chunk_base, Int32(0)), + Int32(tmem_x), + ) + valid_bits = _bmsk_clamp(chunk_left, chunk_right - chunk_left) + chunk = s_data[chunk_idx] + masked_scores = [] + for quad_idx in cutlass.range_constexpr(tmem_x // 4): + quad_base = quad_idx * 4 + masked_scores.extend( + _mask_score_quad( + valid_bits >> Int32(quad_base), + chunk[quad_base], + chunk[quad_base + 1], + chunk[quad_base + 2], + chunk[quad_base + 3], + ) + ) + s_data[chunk_idx] = cutlass.Vector.from_elements( + tuple(masked_scores), self.cfg.qk_acc_dtype + ) + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def left_masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + q_offset: Int32, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Apply the bottom-right-aligned sliding-window left mask.""" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + num_softmax_warps = 4 + warp_id_in_sg = cute.arch.warp_idx() % num_softmax_warps + tmem_row_id = warp_id_in_sg * cute.arch.WARP_SIZE + row_in_tile = tmem_row_id + cute.arch.lane_idx() + + s_data = self._load_s_chunks(stage_info) + + seq_tile_coord, _, _ = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + index_q = seq_tile_coord * self.cfg.q_tile_m + row_in_tile + if cutlass.const_expr(self.cfg.has_varlen or self.cfg.has_q_offset): + window_bound_left = bottom_right_window_left_bound( + index_q, + q_offset, + self.cfg.window_size_left, + ) + kv_tile_start = bottom_right_window_tile_start( + seq_coord=seq_tile_coord, + q_tile_m=self.cfg.q_tile_m, + kv_tile_n=self.cfg.kv_tile_n, + q_offset=q_offset, + window_size_left=self.cfg.window_size_left, + ) + else: + # Exact fixed equal-length fast path: q_offset is constexpr zero. + window_bound_left = index_q - Int32(self.cfg.window_size_left) + kv_tile_start = cute.math.max( + Int32(0), + (seq_tile_coord * self.cfg.q_tile_m - self.cfg.window_size_left) + // self.cfg.kv_tile_n, + ) + kv_tile_abs = kv_tile_start + stage_info.loop_offset + + neg_inf = cutlass.vector.full( + [tmem_x], self.cfg.qk_acc_dtype(-Float32.inf), dtype=self.cfg.qk_acc_dtype + ) + all_true_mask = cutlass.vector.create_mask([tmem_x], [tmem_x]) + + for chunk_idx in cutlass.range_constexpr(num_chunks): + base_k = kv_tile_abs * self.cfg.kv_tile_n + chunk_idx * tmem_x + left_oob_end_idx = window_bound_left - base_k + left_mask_inverted = cutlass.vector.create_mask( + [tmem_x], [left_oob_end_idx] + ) + mask = left_mask_inverted ^ all_true_mask + if cutlass.const_expr(self.cfg.has_varlen or self.cfg.has_q_offset): + # Packed requests share a worst-case window span, and fixed + # bottom-right windows can begin at a non-aligned Q/K offset. + window_bound_right = index_q + q_offset + right_oob_start_idx = window_bound_right + Int32(1) - base_k + right_oob_start_idx = cute.math.min( + cute.math.max(right_oob_start_idx, Int32(0)), + Int32(tmem_x), + ) + right_mask = cutlass.vector.create_mask([tmem_x], [right_oob_start_idx]) + mask = mask & right_mask + s_data[chunk_idx] = cutlass.vector.where(mask, s_data[chunk_idx], neg_inf) + + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def loop_masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + q_offset: Int32, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Loop stage: apply causal masking for mixed Q right-offset batches.""" + s_data = self._load_s_chunks(stage_info) + s_data = self._apply_causal_mask_for_kv_tile( + stage_info, s_data, kv_tile_idx=stage_info.loop_offset, q_offset=q_offset + ) + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=p_chunk) + @cute.jit + def exp2_p( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + ) -> SoftmaxRowSumContribution: + """Apply exp2 using the runtime scale cached before the K/V loop.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), row_max, scale_softmax_log2 + ) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def right_masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + q_offset: Int32, + section: cutlass.Constexpr[FmhaStage], + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Head-paired stage: apply causal/window mask and compute row_max. + + Unlike the query-paired tail mask, this keeps Q0/Q1 on the same + sequence tile; q_half selects a Q head, not a later sequence tile. + Sliding-window tails need both bounds because the final tile can also + contain keys to the left of the visible window. + """ + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + num_softmax_warps = 4 + warp_id_in_sg = cute.arch.warp_idx() % num_softmax_warps + s_data = self._load_s_chunks(stage_info) + if cutlass.const_expr(self.cfg.is_causal): + seq_tile_coord, _, _ = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + tmem_row_id = warp_id_in_sg * cute.arch.WARP_SIZE + row_in_tile = tmem_row_id + cute.arch.lane_idx() + index_q = seq_tile_coord * self.cfg.q_tile_m + row_in_tile + if cutlass.const_expr(self.cfg.window_size_left > 0): + if cutlass.const_expr(self.cfg.has_varlen or self.cfg.has_q_offset): + kv_tile_start = bottom_right_window_tile_start( + seq_coord=seq_tile_coord, + q_tile_m=self.cfg.q_tile_m, + kv_tile_n=self.cfg.kv_tile_n, + q_offset=q_offset, + window_size_left=self.cfg.window_size_left, + ) + else: + kv_tile_start = cute.math.max( + Int32(0), + (seq_tile_coord * self.cfg.q_tile_m - self.cfg.window_size_left) + // self.cfg.kv_tile_n, + ) + if cutlass.const_expr(self.needs_window_tail_left_mask): + window_bound_left = bottom_right_window_left_bound( + index_q, + q_offset, + self.cfg.window_size_left, + ) + else: + kv_tile_start = Int32(0) + if cutlass.const_expr(section == FmhaStage.Loop): + base_k = (kv_tile_start + stage_info.loop_offset) * self.cfg.kv_tile_n + else: + # Tail uses the first tile after the loop domain. + base_k = (kv_tile_start + stage_info.loop_end) * self.cfg.kv_tile_n + for chunk_idx in cutlass.range_constexpr(num_chunks): + chunk_base_k = base_k + chunk_idx * tmem_x + window_bound_right = index_q + q_offset + right_oob_start_idx = window_bound_right + Int32(1) - chunk_base_k + right_oob_start_idx = cute.math.min( + cute.math.max(right_oob_start_idx, Int32(0)), + Int32(tmem_x), + ) + mask = cutlass.vector.create_mask([tmem_x], [right_oob_start_idx]) + if cutlass.const_expr(self.needs_window_tail_left_mask): + left_oob_end_idx = window_bound_left - chunk_base_k + left_mask_inverted = cutlass.vector.create_mask( + [tmem_x], [left_oob_end_idx] + ) + all_true_mask = cutlass.vector.create_mask([tmem_x], [tmem_x]) + left_mask = left_mask_inverted ^ all_true_mask + mask = mask & left_mask + neg_inf = cutlass.vector.full( + [tmem_x], + self.cfg.qk_acc_dtype(-Float32.inf), + dtype=self.cfg.qk_acc_dtype, + ) + s_data[chunk_idx] = cutlass.vector.where( + mask, s_data[chunk_idx], neg_inf + ) + return self._reduce_row_max(s_data, row_max) + + @cute.jit + def _apply_causal_mask_for_kv_tile( + self, + stage_info: StageInfo, + s_data: SoftmaxChunks, + kv_tile_idx: Int32, + q_offset: Int32, + ) -> SoftmaxChunks: + """Apply the query-paired right-edge causal mask to a loaded S tile. + + Query-paired maps q_half=1 to the next sequence tile, so the row index + includes q_half * q_tile_m. Head-paired causal tails use + right_masked_row_max() instead. q_offset is cached once per work tile + so varlen masking does not reload cum_seqlen_q/k in every K/V loop. + """ + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + num_softmax_warps = 4 + warp_id_in_sg = cute.arch.warp_idx() % num_softmax_warps + seq_coord, _, _ = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + kv_base = kv_tile_idx * self.cfg.kv_tile_n + q_min = ( + q_offset + + seq_coord * self.cfg.cta_tiler[0] + + self.q_half * self.cfg.q_tile_m + ) + k_max = kv_base + self.cfg.qk_mma_tiler[1] - Int32(1) + need_mask = q_min <= k_max + if need_mask: + q_idx = q_min + warp_id_in_sg * cute.arch.WARP_SIZE + cute.arch.lane_idx() + for chunk_idx in cutlass.range_constexpr(num_chunks): + k_chunk_base = kv_base + chunk_idx * tmem_x + num_valid = cute.math.min( + cute.math.max(q_idx - k_chunk_base + Int32(1), Int32(0)), + Int32(tmem_x), + ) + causal_mask = cutlass.vector.create_mask([tmem_x], [num_valid]) + neg_inf_vec = cutlass.vector.full_like( + s_data[chunk_idx], Float32(-Float32.inf) + ) + s_data[chunk_idx] = cutlass.vector.where( + causal_mask, s_data[chunk_idx], neg_inf_vec + ) + return s_data + + @cute.jit + def _apply_causal_mask( + self, + stage_info: StageInfo, + s_data: SoftmaxChunks, + q_offset: Int32, + ) -> SoftmaxChunks: + """Apply the tail-stage causal mask to a loaded S tile.""" + return self._apply_causal_mask_for_kv_tile( + stage_info, s_data, stage_info.loop_end, q_offset + ) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + q_offset: Int32, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Tail stage: load S, apply causal mask, and compute row_max.""" + s_data = self._load_s_chunks(stage_info) + if cutlass.const_expr(self.cfg.is_causal): + s_data = self._apply_causal_mask(stage_info, s_data, q_offset) + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=p_chunk) + @cute.jit + def masked_exp2_p( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + ) -> SoftmaxRowSumContribution: + """Tail stage: apply exp2 using the cached runtime softmax scale.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), row_max, scale_softmax_log2 + ) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def invalid_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Tail stage for softmax group 0: identity row_max, no S load.""" + row_max_safe = row_max + if row_max == -Float32.inf: + row_max_safe = Float32(0.0) + _ = stage_info + return row_max, row_max_safe + + @consumer_work + @cute.jit + def invalid_exp2_p( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + ) -> None: + """Tail stage for softmax group 0: no-op because MMA will not read P.""" + pass + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=row_sum) + @cute.jit + def softmax_aux_reduce( + self, + stage_info: StageInfo, + *, + old_row_max: SoftmaxScalar, + row_max: SoftmaxScalar, + row_sum: SoftmaxScalar, + p_chunk: SoftmaxRowSumContribution, + scale_softmax_log2: SoftmaxScalar, + ) -> SoftmaxScalar: + """Accumulate row_sum from vector P fragments or their scalar sum.""" + _ = stage_info + if cutlass.const_expr(self.enable_early_tile_sum): + acc_scale = cute.math.exp2( + scale_softmax_log2 * (old_row_max - row_max), + fastmath=True, + ) + return row_sum * acc_scale + p_chunk + return self._row_sum_reduction( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=old_row_max) + @cute.jit + def softmax_aux_identity( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + ) -> SoftmaxScalar: + """Auxiliary identity path (no P-chunk reduction).""" + _ = stage_info + return row_max + + @cute.jit + def _row_sum_reduction( + self, + *, + old_row_max: SoftmaxScalar, + row_max: SoftmaxScalar, + row_sum: SoftmaxScalar, + p_chunk: SoftmaxChunks, + scale_softmax_log2: SoftmaxScalar, + ) -> Float32: + """Accumulate row_sum from P chunks saved by consumer_work.""" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + scale = scale_softmax_log2 + acc_scale_ = scale * (old_row_max - row_max) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5 + scaled_sum = row_sum * acc_scale + if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): + # Use four independent float2 accumulation chains for D256. A + # single 64-pair chain serializes every FADD behind the + # preceding result and leaves no row-sum ILP after P publication. + local_sum_0 = (scaled_sum, scaled_sum) + local_sum_1 = (Float32(0.0), Float32(0.0)) + local_sum_2 = (Float32(0.0), Float32(0.0)) + local_sum_3 = (Float32(0.0), Float32(0.0)) + for chunk_idx in cutlass.range_constexpr(num_chunks): + p_chunk_vec = p_chunk[chunk_idx] + for elem_idx in cutlass.range_constexpr(0, tmem_x, 8): + local_sum_0 = cute.arch.add_packed_f32x2( + local_sum_0, + (p_chunk_vec[elem_idx], p_chunk_vec[elem_idx + 1]), + rnd="rn", + ftz=False, + ) + local_sum_1 = cute.arch.add_packed_f32x2( + local_sum_1, + (p_chunk_vec[elem_idx + 2], p_chunk_vec[elem_idx + 3]), + rnd="rn", + ftz=False, + ) + local_sum_2 = cute.arch.add_packed_f32x2( + local_sum_2, + (p_chunk_vec[elem_idx + 4], p_chunk_vec[elem_idx + 5]), + rnd="rn", + ftz=False, + ) + local_sum_3 = cute.arch.add_packed_f32x2( + local_sum_3, + (p_chunk_vec[elem_idx + 6], p_chunk_vec[elem_idx + 7]), + rnd="rn", + ftz=False, + ) + local_sum_0 = cute.arch.add_packed_f32x2( + local_sum_0, local_sum_1, rnd="rn", ftz=False + ) + local_sum_2 = cute.arch.add_packed_f32x2( + local_sum_2, local_sum_3, rnd="rn", ftz=False + ) + local_sum_0 = cute.arch.add_packed_f32x2( + local_sum_0, local_sum_2, rnd="rn", ftz=False + ) + return local_sum_0[0] + local_sum_0[1] + + local_sum = (scaled_sum, scaled_sum) + for chunk_idx in cutlass.range_constexpr(num_chunks): + p_chunk_vec = p_chunk[chunk_idx] + for idx in cutlass.range_constexpr(tmem_x // 2): + local_sum = cute.arch.add_packed_f32x2( + local_sum, + (p_chunk_vec[2 * idx], p_chunk_vec[2 * idx + 1]), + rnd="rn", + ftz=False, + ) + return local_sum[0] + local_sum[1] + + +# --------------------------------------------------------------------------- +# TmemPResource -- P-ready handoff from softmax to UMMA +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class TmemPResource(MemoryResource): + """Pipeline-only P handoff for split S/P scheduling. + + Softmax stores P into the TMEM columns owned by ``TmemSPResource`` and + commits this AsyncUmma resource. The MMA task waits on it before issuing + PV, so the next QK can use the other S/P stage without using the S acquire + as an implicit P-ready wait. + """ + + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + tmem_p_offset: Constexpr[int] = field(init=False, default=None) + tmem_p_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + tmem_p_offset: int, + **kwargs: Any, + ) -> None: + """Bind the base P TMEM offset used by the split S/P pipeline.""" + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.tmem_p_offset = tmem_p_offset + self.tmem_p_base = TaskLocalVariable( + dtype=Int32, + default=Int32(tmem_p_offset), + docs="Selected staged P TMEM column base for PV MMA.", + ) + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Return no allocation because P aliases the TmemSP allocation.""" + return [] + + @cute.jit + def create_function_variables(self, context: Optional[Any] = None) -> Int32: + """Create the staged P-base dataflow slot.""" + _ = context + return Int32(self.tmem_p_offset) + + @consumer_work(returns=("tmem_p_base",)) + @cute.jit + def p_base(self, stage_info: StageInfo) -> Int32: + """Return the staged P column base for the MMA PV producer.""" + tmem_p_base = Int32(self.tmem_p_offset) + if cutlass.const_expr(self.cfg.mma_softmax_stage > 1): + tmem_p_base = tmem_p_base + stage_info.stage_idx * self.cfg.qk_mma_tiler[1] + return tmem_p_base + + +# --------------------------------------------------------------------------- +# TmemStatsResource -- TMEM correction statistics with AsyncAsync pipeline +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class TmemStatsResource(MemoryResource): + """Correction statistics with an AsyncAsync pipeline. + + Producer: Softmax writes old_max/row_max/row_sum stats. Consumer: + Correction reads them for O rescaling. Persistent D256 keeps the payload + in a small staged SMEM ring so the stats no longer alias S/P TMEM columns; + other schedules retain the original TMEM storage. + """ + + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + tmem_vec_offset: Constexpr[int] = field(init=False, default=None) + scale_softmax_log2: cute.Tensor | None = field(init=False, default=None) + output_scale: cute.Tensor | None = field(init=False, default=None) + tmem_addr_cached: TmemAddr | None = field(init=False, default=None) + # Precomputed per-warp TMEM vec address (once, before persistent loop). + tmem_vec_addr_cached: TmemAddr | None = field(init=False, default=None) + tmem_ptr_vec_cached: TmemPtr | None = field(init=False, default=None) + + _alloc: Constexpr[Optional[TmemAllocation]] = field(init=False, default=None) + _smem_alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + vec_old_max: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + vec_new_max: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + vec_row_sum: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + vec_scale: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + tmem_vec_offset: int, + scale_softmax_log2: cute.Tensor | None = None, + output_scale: cute.Tensor | None = None, + **kwargs: Any, + ) -> None: + """Bind the correction-stat TMEM vector offset and allocation.""" + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.tmem_vec_offset = tmem_vec_offset + self.scale_softmax_log2 = scale_softmax_log2 + self.output_scale = output_scale + self._alloc = TmemAllocation(f"tmem_vec_{tmem_vec_offset}", cfg.tmem_stats_cols) + self._smem_alloc = None + if cfg.stats_via_smem: + stats_rows = len(cfg.softmax0_warp_ids) * cute.arch.WARP_SIZE + # Each row publishes two packed FP32 values. Loop records hold + # (old_max, new_max), while the final record reuses the pair for + # (row_sum, new_max). + stage_bytes = stats_rows * 2 * 4 + self._smem_alloc = SmemAllocation( + f"smem_vec_{tmem_vec_offset}", + pipeline_config.num_stages * stage_bytes, + alignment=16, + ) + self.tmem_addr_cached = Int32(0) + self.tmem_vec_addr_cached = Int32(0) + self.tmem_ptr_vec_cached = _placeholder_tmem_ptr() + self.vec_old_max = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="Previous row maximum read from TMEM stats.", + ) + self.vec_new_max = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="Current row maximum read from TMEM stats.", + ) + self.vec_row_sum = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="Softmax denominator read from TMEM stats.", + ) + self.vec_scale = TaskLocalVariable( + dtype=Float32, + default=Float32(1.0), + docs="Correction scale derived from TMEM stats.", + ) + self.scale_softmax_log2_value = TaskLocalVariable( + dtype=Float32, + # Placeholder before load_scale_softmax_log2 reads the runtime tensor. + default=Float32(0.0), + docs="Softmax scale cached from the runtime scale tensor.", + ) + self.output_scale_value = TaskLocalVariable( + dtype=Float32, + # Placeholder before load_output_scale reads the runtime tensor. + default=Float32(1.0), + docs="Output scale cached from the runtime scale tensor.", + ) + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Return the TMEM allocation for one correction-stat vector.""" + if self.cfg.stats_via_smem: + return [] + return [self._alloc] + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Return the staged SMEM stats ring when TMEM aliasing is disabled.""" + if not self.cfg.stats_via_smem: + return [] + return [self._smem_alloc] + + @cute.jit + def _init_function_state(self, stage_info: StageInfo) -> None: + """Initialize vec address fields to establish DSL type before scf.while. + + Real values computed by per-work-tile auxiliary work after setmaxnreg. + + Emits the correction stat slots (vec_old_max, vec_new_max, + vec_row_sum, vec_scale) consumed by TmemO / SmemO via consumer- + to-consumer routing; producer-side old_row_max / row_max / + row_sum slots are auto-mirrored from TmemSP by + Task.init_variables. + """ + self.tmem_vec_addr_cached = Int32(0) + self.tmem_ptr_vec_cached = prims.make_tmem_ptr(Int32(0), cutlass.Int8) + _ = stage_info + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_state(self, stage_info: StageInfo) -> None: + self._init_function_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_read_state(self, stage_info: StageInfo) -> None: + self._init_function_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns="scale_softmax_log2_value") + @cute.jit + def load_scale_softmax_log2(self, stage_info: StageInfo) -> Float32: + """Load the runtime softmax scale once before the correction loop.""" + _ = stage_info + if cutlass.const_expr(self.scale_softmax_log2 is None): + # Safe fallback for validation-only resource construction. + return Float32(0.0) + return self.scale_softmax_log2[0] + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns="output_scale_value") + @cute.jit + def load_output_scale(self, stage_info: StageInfo) -> Float32: + """Load the runtime output scale once before the correction loop.""" + _ = stage_info + if cutlass.const_expr(self.output_scale is None): + # Identity fallback for validation-only resource construction. + return Float32(1.0) + return self.output_scale[0] + + @cute.jit + def _init_work_tile_state(self, stage_info: StageInfo) -> None: + """Compute per-warp TMEM vec address and pointer each tile. + + Deferred from function-scope auxiliary work so the arithmetic runs + after setmaxnreg and does not spill across the register-budget boundary. + """ + if cutlass.const_expr(self.cfg.stats_via_smem): + return + # Softmax producer and correction consumer both use 4 warps. + num_warps = 4 + warp_id_in_wg = cute.arch.warp_idx() % num_warps + tmem_raw_addr = self.tmem_addr_cached + tmem_base_row = tmem_raw_addr >> 16 + tmem_base_col = tmem_raw_addr & Int32(0xFFFF) + row_id = tmem_base_row + warp_id_in_wg * cute.arch.WARP_SIZE + self.tmem_vec_addr_cached = (row_id << 16) | ( + tmem_base_col + self.tmem_vec_offset + ) + self.tmem_ptr_vec_cached = prims.make_tmem_ptr( + self.tmem_vec_addr_cached, cutlass.Int8 + ) + _ = stage_info + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_read_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @cute.jit + def _stage_col_offset(self, stage_info: StageInfo) -> Int32 | int: + """Return the TMEM column offset for stage-scoped stats.""" + stage_col_offset = Int32(0) + if cutlass.const_expr(self.cfg.stage_scoped_tmem_stats): + stage_col_offset = stage_info.stage_idx * self.cfg.qk_mma_tiler[1] + return stage_col_offset + + @cute.jit + def _stats_smem_ptr(self, stage_info: StageInfo) -> cute.Pointer: + """Return this warp-group thread's staged SMEM stats slot.""" + context = stage_info.context + assert context is not None and context.smem_base is not None + assert self._smem_alloc is not None + stats_rows = len(self.cfg.softmax0_warp_ids) * cute.arch.WARP_SIZE + stats_elems_per_row = 2 + stage_elems = stats_rows * stats_elems_per_row + tidx, _, _ = cute.arch.thread_idx() + row_idx = tidx % stats_rows + base_ptr = context.smem_base.data_ptr() + self._smem_alloc.offset + view = cutlass.Array( + base_ptr, + dtype=Float32, + shape=(self.pipeline_config.num_stages * stage_elems,), + addrspace=3, + ) + elem_offset = stage_info.stage_idx * stage_elems + row_idx * stats_elems_per_row + return view.subview(elem_offset).data_ptr() + + @producer_work + @cute.jit + def store_vec( + self, + stage_info: StageInfo, + *, + old_row_max: SoftmaxScalar, + row_max: SoftmaxScalar, + row_sum: SoftmaxScalar, + final_stats: cutlass.Constexpr[bool] = False, + ) -> None: + """Softmax: publish correction statistics for one row. + + The TMEM-backed topology writes four elements per row: + [0] = old_row_max (previous iteration's max) + [1] = new_row_max (current iteration's max) + [2] = row_sum (accumulated softmax denominator) + [3] = padding + + The compact SMEM-backed topology writes ``[old_max, new_max]`` during + the loop. Its final publication repurposes slot 0 for ``row_sum``. + + The Correction warp reads these to compute the rescale factor: + scale = exp2(scale_log2 * (old_max - new_max)) + and to forward row_sum to SmemO for the final normalization. + """ + if cutlass.const_expr(self.cfg.stats_via_smem): + stat0 = old_row_max + if cutlass.const_expr(final_stats): + stat0 = row_sum + vec_data = cutlass.Vector.from_elements( + (stat0, row_max), + self.cfg.qk_acc_dtype, + ) + self._stats_smem_ptr(stage_info).store(vec_data, alignment=8) + else: + vec_data = cutlass.Vector.from_elements( + (old_row_max, row_max, row_sum, Float32(0.0)), + self.cfg.qk_acc_dtype, + ) + tmem_ptr_vec = prims.make_tmem_ptr( + self.tmem_vec_addr_cached + self._stage_col_offset(stage_info), + cutlass.Int8, + ) + prims.tcgen05_st( + "32x32b", + tmem_ptr_vec, + vec_data, + ) + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def _read_vec( + self, + stage_info: StageInfo, + scale_softmax_log2: SoftmaxScalar, + final_stats: cutlass.Constexpr[bool] = False, + ) -> tuple[SoftmaxScalar, SoftmaxScalar, SoftmaxScalar, SoftmaxScalar]: + """Read correction stats from TMEM and cache in consumer_vars. + + CRITICAL: This read must happen here (immediately after the + pipeline wait) rather than being deferred to TmemO.correct, + because the TMEM stats region (cols 128-131 for TmemStats1, cols 0-3 + for TmemStats0) overlaps with the S0/S1 score regions. After the MMA warp + commits O and continues to QK→S, the UMMA write to S can + overwrite the stats data. Reading here, before the O pipeline + wait, ensures the stats are captured before any S overwrite. + + TmemO.correct retrieves the cached values from this + resource's consumer_vars via a direct reference. + """ + tmem_shape_vec = "32x32b" + # Vector layout: [old_max, new_max, row_sum, pad]. + tmem_x_vec = self.cfg.tmem_stats_cols + + if cutlass.const_expr(self.cfg.stats_via_smem): + vec_rmem = self._stats_smem_ptr(stage_info).load( + count=2, + alignment=8, + ) + else: + tmem_ptr_vec = prims.make_tmem_ptr( + self.tmem_vec_addr_cached + self._stage_col_offset(stage_info), + self.cfg.qk_acc_dtype, + ) + vec_rmem = cutlass.Array(self.cfg.qk_acc_dtype, tmem_x_vec) + vec_rmem[0:tmem_x_vec] = prims.tcgen05_ld( + tmem_shape_vec, tmem_ptr_vec, num=tmem_x_vec + ) + cute.arch.fence_view_async_tmem_load() + + vec_old_max = vec_rmem[0] + vec_new_max = vec_rmem[1] + vec_row_sum = Float32(0.0) + if cutlass.const_expr(not self.cfg.stats_via_smem): + vec_row_sum = vec_rmem[2] + scale = Float32(1.0) + if cutlass.const_expr(not (self.cfg.stats_via_smem and final_stats)): + scale_ = scale_softmax_log2 * (vec_old_max - vec_new_max) + scale = cute.math.exp2(scale_, fastmath=True) + else: + vec_row_sum = vec_rmem[0] + vec_old_max = vec_new_max + _ = stage_info + return vec_old_max, vec_new_max, vec_row_sum, scale + + @consumer_work(returns=(vec_old_max, vec_new_max, vec_row_sum, vec_scale)) + @cute.jit + def read_vec( + self, + stage_info: StageInfo, + *, + scale_softmax_log2: SoftmaxScalar, + final_stats: cutlass.Constexpr[bool] = False, + ) -> tuple[SoftmaxScalar, SoftmaxScalar, SoftmaxScalar, SoftmaxScalar]: + """Read correction stats using the cached runtime softmax scale.""" + return self._read_vec(stage_info, scale_softmax_log2, final_stats) + + +# --------------------------------------------------------------------------- +# TmemOResource -- TMEM O accumulation with UmmaAsync pipeline +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class TmemOResource(MemoryResource): + """TMEM O accumulation with a topology-derived UmmaAsync pipeline. + + Producer: MMA writes P*V -> O (double-buffered O0/O1). + Consumer: Correction rescales O in-place. + + Paired schedules use two stages so MMA can commit O0, work on O1, + commit O1, then acquire O0. SMEM-stats D256 uses one stage because it + writes one physical O accumulator. + """ + + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + tmem_o0_offset: Constexpr[int] = field(init=False, default=None) + tmem_o1_offset: Constexpr[int] = field(init=False, default=None) + tmem_addr_cached: TmemAddr | None = field(init=False, default=None) + # Precomputed TMEM raw pointer (inttoptr of tmem_addr_cached). + tmem_ptr_raw_cached: TmemPtr | None = field(init=False, default=None) + # Precomputed per-warp TMEM O address base: (row_id << 16) | tmem_base_col. + # consumer_work adds tmem_o_offset to get the final O0/O1 address. + tmem_o_addr_base_cached: TmemAddr | None = field(init=False, default=None) + # P-stage base supplied by TmemPResource for split S/P scheduling. + tmem_p_base_cached: TmemAddr | None = field(init=False, default=None) + # References to TmemStats resources for reading cached correction stats. + # consumer_work reads stats from these instead of from TMEM, because + # the stats TMEM region overlaps with S0/S1 and can be overwritten by + # MMA's QK→S before correction reads it. + tmem_vec0_resource: TmemStatsResource | None = field(init=False, default=None) + tmem_vec1_resource: TmemStatsResource | None = field(init=False, default=None) + + _alloc_o0: Constexpr[Optional[TmemAllocation]] = field(init=False, default=None) + _alloc_o1: Constexpr[Optional[TmemAllocation]] = field(init=False, default=None) + + def __init__( + self, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + tmem_o0_offset: int, + tmem_o1_offset: int, + tmem_vec0_resource: TmemStatsResource | None = None, + tmem_vec1_resource: TmemStatsResource | None = None, + **kwargs: Any, + ) -> None: + """Bind O TMEM offsets and correction-stat resources.""" + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.tmem_o0_offset = tmem_o0_offset + self.tmem_o1_offset = tmem_o1_offset + self.tmem_vec0_resource = tmem_vec0_resource + self.tmem_vec1_resource = tmem_vec1_resource + self._alloc_o0 = TmemAllocation("tmem_o0", 128) + self._alloc_o1 = TmemAllocation("tmem_o1", 128) + self.tmem_addr_cached = Int32(0) + self.tmem_ptr_raw_cached = _placeholder_tmem_ptr() + self.tmem_o_addr_base_cached = Int32(0) + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Return the TMEM allocations for the double-buffered O accumulators.""" + return [self._alloc_o0, self._alloc_o1] + + @property + def loop_offset_sensitive(self) -> bool: + """Return true because PV accumulation scale depends on loop_offset.""" + # producer_work uses loop_offset to compute scale_d. + return True + + @cute.jit + def _init_function_state(self, stage_info: StageInfo) -> None: + """Precompute MMA-warp TMEM raw pointer (once, ungated). + + Per-warp O address base for correction warps is deferred to + per-work-tile auxiliary work to avoid crossing the setmaxnreg boundary. + tmem_o_addr_base_cached initialized to Int32(0) to establish DSL type. + + Pure consumer/producer of upstream emitters — emits no + consumer vars itself. Producer-side desc_v_base slot is + auto-mirrored from SmemKV; consumer-side vec_old_max / + vec_new_max slots are auto-mirrored from TmemStats via + consumer-to-consumer routing in the captured schedule. + """ + self.tmem_ptr_raw_cached = prims.make_tmem_ptr( + self.tmem_addr_cached, cutlass.Int8 + ) + self.tmem_o_addr_base_cached = Int32(0) + self.tmem_p_base_cached = Int32(0) + _ = stage_info + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_state(self, stage_info: StageInfo) -> None: + self._init_function_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_correction_state(self, stage_info: StageInfo) -> None: + self._init_function_state(stage_info) + + @cute.jit + def _init_work_tile_state(self, stage_info: StageInfo) -> None: + """Compute per-warp TMEM O address base each tile (after setmaxnreg).""" + num_correction_warps = 4 + warp_id_in_wg = cute.arch.warp_idx() % num_correction_warps + tmem_raw_addr = self.tmem_addr_cached + tmem_base_row = tmem_raw_addr >> 16 + tmem_base_col = tmem_raw_addr & Int32(0xFFFF) + row_id = tmem_base_row + warp_id_in_wg * cute.arch.WARP_SIZE + self.tmem_o_addr_base_cached = (row_id << 16) | tmem_base_col + _ = stage_info + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_correction_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def set_p_base(self, stage_info: StageInfo, *, tmem_p_base: Int32) -> None: + """Cache the P TMEM column base selected by TmemPResource.""" + _ = stage_info + self.tmem_p_base_cached = tmem_p_base + + @producer_work + @cute.jit + def pv_mma( + self, + stage_info: StageInfo, + *, + desc_v_base: prims.Tcgen05SmemDesc, + section: cutlass.Constexpr[FmhaStage], + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + inst_idx: cutlass.Constexpr[int] = 0, + is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """PV MMA: P*V -> O (double-buffered O0/O1). + + Uses captured schedule section and call index to select O0/O1 and + scale_d statically. + Reads P from TMEM (via tmem_p offset on the corresponding SP resource) + and V descriptor from SmemV consumer vars. + + In causal mode with no Q right offset, skips P0*V→O0 MMA in the last + LOOP iteration, since Softmax0's domain is N-2 but MMA's domain is N-1. + The task domain pads partial final CTAs so this slot is always outside + peer0's causal reach. + """ + if cutlass.const_expr(section == FmhaStage.Head): + writes_o0 = True + first_o0_write = True + first_o1_write_maybe = False + elif cutlass.const_expr(section == FmhaStage.Loop): + writes_o0 = inst_idx == 1 + first_o0_write = False + first_o1_write_maybe = inst_idx == 0 + else: + writes_o0 = False + first_o0_write = False + first_o1_write_maybe = True + + # In causal mode, check if O0 MMA should skip the last LOOP iteration. + skip_o0_invalid = False + if cutlass.const_expr( + self.cfg.skip_causal_invalid_peer0 + and writes_o0 + and section == FmhaStage.Loop + ): + if not is_tail: + skip_o0_invalid = stage_info.loop_offset == (stage_info.loop_end - 1) + + if not skip_o0_invalid: + tmem_ptr_raw = self.tmem_ptr_raw_cached + + if cutlass.const_expr(self.cfg.v_dtype.width == 8): + mma_kind = prims.Tcgen05MMAKind.F8F6F4 + # E4M3 operands use the Float16 encoding handle. + ab_format = cutlass.Float16 + else: + mma_kind = prims.Tcgen05MMAKind.F16 + if cutlass.const_expr(self.cfg.v_dtype == cutlass.BFloat16): + ab_format = cutlass.BFloat16 + else: + ab_format = cutlass.Float16 + + idesc_pv = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=ab_format, + b_dtype=ab_format, + n_dim=( + self.cfg.head_dim_per_stage_kv + if self.cfg.single_qkv_instance and self.cfg.pv_mma_tiler[1] == 256 + else self.cfg.pv_mma_tiler[1] + ), + m_dim=self.cfg.pv_mma_tiler[0], + # V is row-major / MN-major. + b_major=1, + ) + + pv_n_dim = self.cfg.pv_mma_tiler[1] + num_head_dim_stages = 1 + if cutlass.const_expr( + self.cfg.single_qkv_instance and self.cfg.pv_mma_tiler[1] == 256 + ): + pv_n_dim = self.cfg.head_dim_per_stage_kv + num_head_dim_stages = self.cfg.pv_mma_tiler[1] // pv_n_dim + head_dim_stage_start = 0 + num_head_dim_stages_to_issue = num_head_dim_stages + if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): + num_head_dim_stages_to_issue = 1 + head_dim_stage_start = head_dim_stage_idx + + k_dim_per_mma = 16 + if cutlass.const_expr(self.cfg.v_dtype.width != 16): + k_dim_per_mma = 32 + num_kphases_pv = self.cfg.pv_mma_tiler[2] // k_dim_per_mma + inc_tmem_p = ( + k_dim_per_mma * self.cfg.v_dtype.width // self.cfg.qk_acc_dtype.width + ) + tma_copy_iters_per_head_dim_stage = ( + self.cfg.tma_copy_qkv_iters // num_head_dim_stages + ) + if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): + tma_copy_iters_per_head_dim_stage = self.cfg.tma_copy_kv_stage_iters + inc_bytes_v = ( + k_dim_per_mma + * (pv_n_dim // tma_copy_iters_per_head_dim_stage) + * self.cfg.v_dtype.width + // 8 + ) + kv_chunk_bytes = ( + self.cfg.tma_copy_kv_bytes // self.cfg.tma_copy_kv_stage_iters + ) + head_dim_stage_bytes_v = kv_chunk_bytes * tma_copy_iters_per_head_dim_stage + + # Select O buffer and P offset at trace time (compile-time constant) + if cutlass.const_expr(self.cfg.single_qkv_instance or writes_o0): + tmem_ptr_o = tmem_ptr_raw.subview(self.tmem_o0_offset) + tmem_p_base = self.cfg.tmem_p0_offset + else: + tmem_ptr_o = tmem_ptr_raw.subview(self.tmem_o1_offset) + tmem_p_base = self.cfg.tmem_p1_offset + if cutlass.const_expr(self.cfg.single_qkv_instance): + tmem_p_base = self.tmem_p_base_cached + + # scale_d at trace time. + # O0 (even counters): HEAD initializes O0, so all + # subsequent O0 writes (counter 2, 4, ...) always accumulate. + # O1 (odd counters): first written in LOOP. Dynamic check needed + # because with loop peeling + causal domain=1, the peeled iteration + # may be the first O1 write (loop_offset=0 → scale_d=False). + # TAIL O1: always accumulates because LOOP or the peeled iteration + # already wrote O1 before TAIL runs. + if cutlass.const_expr(self.cfg.single_qkv_instance): + if cutlass.const_expr(self.cfg.has_tmem_p_pipeline): + if cutlass.const_expr(section == FmhaStage.Loop): + scale_d = stage_info.loop_offset > stage_info.loop_start + elif cutlass.const_expr(is_tail): + scale_d = stage_info.loop_end > stage_info.loop_start + else: + scale_d = False + elif cutlass.const_expr(is_tail): + scale_d = stage_info.loop_end > 0 + elif cutlass.const_expr(section == FmhaStage.Loop): + scale_d = stage_info.loop_offset > 0 + else: + scale_d = False + elif cutlass.const_expr(first_o0_write): + # Head O0 is the first O0 write. + scale_d = False + elif cutlass.const_expr(first_o1_write_maybe and section == FmhaStage.Tail): + # TAIL O1 accumulates if LOOP already wrote O1 (domain >= 1). + # When domain=0, TAIL is the first O1 write. + scale_d = stage_info.loop_end > 0 + elif cutlass.const_expr(first_o1_write_maybe): + # Loop O1 initializes on the first iteration and accumulates later. + scale_d = stage_info.loop_offset > 0 + else: + # O0 after the head write always accumulates. + scale_d = True + # Prevent LLVM from rematerializing V descriptor inside + # each elect_sync block (same pattern as QK MMA above). + desc_v_base_ = freeze_smem_descriptor(desc_v_base) + + if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): + tmem_ptr_o_stage = tmem_ptr_o.subview(head_dim_stage_start * pv_n_dim) + scale_d_stage = scale_d + for k_idx in cutlass.range_constexpr(num_kphases_pv): + dp = tmem_ptr_raw.subview(tmem_p_base + k_idx * inc_tmem_p) + increment = (inc_bytes_v * k_idx) >> 4 + dv = desc_v_base_ + increment + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + prims.CTAGroup.CTA_1, + tmem_ptr_o_stage, + dp, + dv, + idesc_pv, + scale_d_stage, + ) + scale_d_stage = True + else: + for head_dim_stage_idx in cutlass.range_constexpr( + num_head_dim_stages_to_issue + ): + tmem_ptr_o_stage = tmem_ptr_o.subview(head_dim_stage_idx * pv_n_dim) + v_stage_increment = ( + head_dim_stage_bytes_v * head_dim_stage_idx + ) >> 4 + scale_d_stage = scale_d + for k_idx in cutlass.range_constexpr(num_kphases_pv): + dp = tmem_ptr_raw.subview(tmem_p_base + k_idx * inc_tmem_p) + increment = v_stage_increment + ((inc_bytes_v * k_idx) >> 4) + dv = desc_v_base_ + increment + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + prims.CTAGroup.CTA_1, + tmem_ptr_o_stage, + dp, + dv, + idesc_pv, + scale_d_stage, + ) + scale_d_stage = True + + @consumer_work + @cute.jit + def correct( + self, + stage_info: StageInfo, + *, + vec_old_max: SoftmaxScalar, + vec_new_max: SoftmaxScalar, + vec_scale: SoftmaxScalar, + inst_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """Correction using the scale cached by TmemStatsResource.""" + self._correct_impl( + stage_info, + vec_old_max=vec_old_max, + vec_new_max=vec_new_max, + scale_softmax_log2=Float32(0.0), + vec_scale=vec_scale, + use_cached_scale=True, + inst_idx=inst_idx, + is_tail=is_tail, + ) + + @cute.jit + def _correct_impl( + self, + stage_info: StageInfo, + *, + vec_old_max: SoftmaxScalar, + vec_new_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + vec_scale: SoftmaxScalar, + use_cached_scale: Constexpr[bool], + inst_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """Correction: read cached stats, compute scale, rescale O. + + Reads the correction stats [old_max, new_max, row_sum] from the + TmemStatsResource's consumer_vars (cached by TmemStats.read_vec + right after the pipeline wait). This avoids a TMEM race: the stats + region overlaps with S0/S1, and MMA's QK->S can overwrite it + between O commit and the time Correction reads here. + + Uses inst_idx to select O0/O1 and forwards row_sum to + SmemO.producer_vars for the tail epilog. + + With skip_correction enabled: uses vote_ballot_sync to check if + old_max == new_max across all threads. If true, scale=1.0 and + we skip the expensive TMEM load+rescale+store. + """ + # Select O offset from captured correction-call position. + if cutlass.const_expr(self.cfg.single_qkv_instance or inst_idx == 0): + tmem_o_offset = self.tmem_o0_offset + else: + tmem_o_offset = self.tmem_o1_offset + + # In causal mode with no Q right offset, skip the invalid O0 correction + # in the last LOOP iteration. The task domain pads partial final CTAs so + # this slot is always outside peer0's causal reach. + skip_o0_invalid = False + if cutlass.const_expr( + self.cfg.skip_causal_invalid_peer0 + and not self.cfg.single_qkv_instance + and inst_idx == 0 + ): + # This is O0 correction; check if this is the last LOOP iteration. + if not is_tail: + skip_o0_invalid = stage_info.loop_offset == (stage_info.loop_end - 1) + + # Check if we should skip correction (when old_max == new_max) + should_rescale = True + if cutlass.const_expr(self.cfg.enable_skip_correction): + vote_ballot_cnt = cute.arch.vote_ballot_sync(vec_old_max != vec_new_max) + should_rescale = vote_ballot_cnt != Int32(0) + + scale = Float32(1.0) + if should_rescale: + if cutlass.const_expr(use_cached_scale): + scale = vec_scale + else: + scale_ = scale_softmax_log2 * (vec_old_max - vec_new_max) + scale = cute.math.exp2(scale_, fastmath=True) + + # PTX ISA 9.7.16.6.4.4: Non-pipelined instructions, different thread. + # MMA (Thread 0) does tcgen05.mma → tcgen05.commit on O_full. + # Correction (Thread 1) does mbarrier.try_wait on O_full → tcgen05.ld. + # The fence orders the prior tcgen05.commit's completion with our tcgen05.ld. + from cutlass.experimental import primitives as _prims + + _prims.tcgen05_fence("after") + + # Only rescale if old_max != new_max AND not in invalid O0 iteration + if should_rescale and not skip_o0_invalid: + # Load O, rescale, store back + tmem_o_addr = self.tmem_o_addr_base_cached + tmem_o_offset + + tmem_shape = "32x32b" + tmem_x = 16 + + num_iters = self.cfg.cta_tiler[2] // tmem_x + for i in cutlass.range_constexpr(num_iters): + tmem_tile_addr = tmem_o_addr + i * tmem_x + tmem_ptr = cutlass.inttoptr( + tmem_tile_addr, + mem_space=6, + dtype=self.cfg.pv_acc_dtype, + ) + + # Load from TMEM as vector, scale, store back + o_vec = prims.tcgen05_ld(tmem_shape, tmem_ptr, num=tmem_x) + cute.arch.fence_view_async_tmem_load() + scale_vec = cutlass.vector.full_like(o_vec, scale) + o_scaled = o_vec * scale_vec + prims.tcgen05_st(tmem_shape, tmem_ptr, o_scaled) + + cute.arch.fence_view_async_tmem_store() + # else: skip TMEM rescale entirely when scale=1.0 + + +# --------------------------------------------------------------------------- +# SmemOResource -- SMEM O buffer with AsyncAsync pipeline +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class SmemOResource(MemoryResource): + """SMEM buffer for one O-subtile (1-stage AsyncAsync pipeline). + + Each instance (``smem_o_0``, ``smem_o_1``) owns a distinct smem + region for its subtile, so the checker can verify that stage-0 + and stage-1 accesses never conflict. + + Producer: CorrectionTask (correction_epilog writes converted O to SMEM). + Consumer: EpilogueTask (TMA stores O from SMEM to GMEM). + """ + + sO_array: cutlass.Array = field(init=False, default=None) + tmem_addr_cached: TmemAddr | None = field(init=False, default=None) + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + tmem_o_addr_base_cached: TmemAddr | None = field(init=False, default=None) + # Reference to the TmemStats resource for this stage's correction stats. + tmem_vec_resource: TmemStatsResource | None = field(init=False, default=None) + stage_idx: Constexpr[int] = field(init=False, default=0) + _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + head_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + batch_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + seq_coord_q: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + stage_idx: int = 0, + tmem_vec_resource: TmemStatsResource | None = None, + **kwargs: Any, + ) -> None: + """Bind one output subtile stage and reserve its SMEM staging buffer.""" + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.stage_idx = stage_idx + self.tmem_vec_resource = tmem_vec_resource + stage_elements = cfg.sO_stage_elements + size_bytes = stage_elements * cfg.o_dtype.width // 8 + self._alloc = SmemAllocation( + f"smem_o_{stage_idx}", size_bytes, alignment=cfg.buffer_align_bytes + ) + self.sO_array = _placeholder_smem_array(cfg.o_dtype) + self.tmem_addr_cached = Int32(0) + self.tmem_o_addr_base_cached = Int32(0) + self.head_coord = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Q/O head coordinate for the output subtile.", + ) + self.batch_coord = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Batch coordinate for the output subtile.", + ) + self.seq_coord_q = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Output row coordinate for the output subtile.", + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Return the SMEM allocation for this O staging subtile.""" + return [self._alloc] + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + """Derive sO_array from context (once, ungated). + + Per-warp TMEM O address base deferred to per-work-tile auxiliary work + to avoid crossing the setmaxnreg boundary. + tmem_o_addr_base_cached initialized to Int32(0) to establish DSL type. + + Emits per-tile output coordinates consumed downstream by + GmemO via the EpilogueTask; producer-side vec_row_sum / + vec_scale slots are auto-mirrored from TmemStats by + Task.init_variables. + """ + smem_base = stage_info.context.smem_base + stage_elements = self.cfg.sO_stage_elements + self.sO_array = cutlass.Array( + smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.o_dtype, + shape=(stage_elements,), + addrspace=3, + ) + self.tmem_o_addr_base_cached = Int32(0) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_output_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @cute.jit + def _init_work_tile_state(self, stage_info: StageInfo) -> None: + """Compute per-warp TMEM O address base each tile (after setmaxnreg).""" + num_correction_warps = 4 + warp_id_in_wg = cute.arch.warp_idx() % num_correction_warps + tmem_raw_addr = self.tmem_addr_cached + tmem_base_row = tmem_raw_addr >> 16 + tmem_base_col = tmem_raw_addr & Int32(0xFFFF) + row_id = tmem_base_row + warp_id_in_wg * cute.arch.WARP_SIZE + self.tmem_o_addr_base_cached = (row_id << 16) | tmem_base_col + _ = stage_info + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_output_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @producer_work + @cute.jit + def store_o( + self, + stage_info: StageInfo, + *, + vec_row_sum: SoftmaxScalar, + vec_scale: SoftmaxScalar, + output_scale: SoftmaxScalar, + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + ) -> None: + """Store O using the output scale cached before the loop.""" + self._store_o( + stage_info, + vec_row_sum=vec_row_sum, + vec_scale=vec_scale, + output_scale=output_scale, + head_dim_stage_idx=head_dim_stage_idx, + ) + + @cute.jit + def _store_o( + self, + stage_info: StageInfo, + *, + vec_row_sum: SoftmaxScalar, + vec_scale: SoftmaxScalar, + output_scale: SoftmaxScalar, + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + ) -> None: + """Correct and normalize O in one TMEM pass, then stage it in SMEM.""" + if cutlass.const_expr(self.stage_idx == 0): + tmem_o_offset = self.cfg.tmem_o0_offset + else: + tmem_o_offset = self.cfg.tmem_o1_offset + sO_base = self.sO_array + + # Read precomputed correction scale and row_sum from TmemStats. + # vec_scale = exp2(scale_log2 * (old_max - new_max)) was computed + # in TmemStats.consumer_work. + correction_scale = vec_scale + scale = output_scale * correction_scale / vec_row_sum + + num_correction_warps = 4 + tidx, _, _ = cute.arch.thread_idx() + tid_in_wg = tidx % (cute.arch.WARP_SIZE * num_correction_warps) + + o_head_dim = self.cfg.epi_tile[1] + tma_copy_o_iters = self.cfg.tma_copy_o_iters + if cutlass.const_expr(self.cfg.stage_o_by_head_dim): + o_head_dim = self.cfg.head_dim_per_stage_kv + tma_copy_o_iters = self.cfg.tma_copy_o_stage_iters + + tmem_offset_o = ( + self.tmem_o_addr_base_cached + + tmem_o_offset + + head_dim_stage_idx * o_head_dim + ) + + tmem_shape = "32x32b" + tmem_x = 16 + num_iters = o_head_dim // tmem_x + + smem_o_swizzle = _smem_o_swizzle(self.cfg) + + d_block_size = o_head_dim // tma_copy_o_iters + row_offset = tid_in_wg * d_block_size + + for i in cutlass.range_constexpr(num_iters): + tmem_offset_tile = tmem_offset_o + i * tmem_x + + tmem_ptr = cutlass.inttoptr( + tmem_offset_tile, + mem_space=6, + dtype=self.cfg.pv_acc_dtype, + ) + + o_rmem = prims.tcgen05_ld(tmem_shape, tmem_ptr, num=tmem_x) + cute.arch.fence_view_async_tmem_load() + + scale_vec = cutlass.vector.full_like(o_rmem, scale) + o_rmem = o_rmem * scale_vec + + o_rmem_dtype = o_rmem.to(self.cfg.o_dtype) + + col_offset = (i * tmem_x) % d_block_size + block_idx = (i * tmem_x) // d_block_size + block_offset = block_idx * self.cfg.tma_copy_o_granu_elems + smem_offset = block_offset + row_offset + col_offset + smem_ptr = (sO_base.subview(smem_offset)).data_ptr() + + if cutlass.const_expr(self.cfg.o_dtype.width == 8): + o_rmem_i8 = o_rmem_dtype.bitcast(cutlass.Int8) + smem_ptr.store_swizzled(o_rmem_i8, alignment=64, swizzle=smem_o_swizzle) + else: + smem_ptr.store_swizzled( + o_rmem_dtype, alignment=64, swizzle=smem_o_swizzle + ) + + prims.fence_proxy( + kind=prims.Proxy.ASYNC_SHARED, + space=prims.SharedSpace.shared_cta, + ) + + @consumer_work(returns=(head_coord, batch_coord, seq_coord_q)) + @cute.jit + def compute_output_coords( + self, stage_info: StageInfo + ) -> tuple[Int32, Int32, Int32]: + """Return output-tile coordinates for downstream GMEM TMA store.""" + seq_coord, head_coord, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + seq_coord_q = seq_coord * self.cfg.q_tile_m * self.cfg.work_tile_q_seq_tiles + return head_coord, batch_coord, seq_coord_q + + +# --------------------------------------------------------------------------- +# GmemOResource -- global memory O output (no pipeline) +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class GmemOResource(MemoryResource): + """TMA store for one O-subtile to global memory. + + Each instance (``gmem_o_0``, ``gmem_o_1``) has its own smem + staging region aliased with the matching ``SmemOResource`` + instance, keeping per-subtile accesses independently trackable. + No pipeline — point-access only. + + Producer: EpilogueTask stores O tiles from SMEM to GMEM via TMA. + """ + + tma_o_desc: cutlass.Pointer | None = field(init=False, default=None) + cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) + sO_array: cutlass.Array = field(init=False, default=None) + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + stage_idx: Constexpr[int] = field(init=False, default=0) + _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + + def __init__( + self, + tma_o_desc: cutlass.Pointer | None, + cum_seqlen_q: cute.Tensor | None, + cfg: FmhaConfig, + stage_idx: int = 0, + **kwargs: Any, + ) -> None: + """Bind the O TMA descriptor and reserve store-side SMEM staging.""" + super().__init__(**kwargs) + self.tma_o_desc = tma_o_desc + self.cum_seqlen_q = cum_seqlen_q + self.cfg = cfg + self.stage_idx = stage_idx + stage_elements = cfg.sO_stage_elements + size_bytes = stage_elements * cfg.o_dtype.width // 8 + self._alloc = SmemAllocation( + f"gmem_o_{stage_idx}_smem", size_bytes, alignment=cfg.buffer_align_bytes + ) + self.sO_array = _placeholder_smem_array(cfg.o_dtype) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Return the SMEM allocation used as the O TMA store source.""" + return [self._alloc] + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_state(self, stage_info: StageInfo) -> None: + """Materialize the O store staging buffer; this resource emits no slots.""" + # Pure sink — producer-side head_coord / batch_coord / seq_coord_q + # slots are auto-mirrored from upstream SmemO by Task.init_variables. + smem_base = stage_info.context.smem_base + stage_elements = self.cfg.sO_stage_elements + self.sO_array = cutlass.Array( + smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.o_dtype, + shape=(stage_elements,), + addrspace=3, + ) + + @producer_work + @cute.jit + def tma_store( + self, + stage_info: StageInfo, + *, + head_coord: Int32, + batch_coord: Int32, + seq_coord_q: Int32, + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + correction_fused: cutlass.Constexpr[bool] = False, + ) -> None: + """TMA store O from SMEM to GMEM. + + Coordinates are produced by SmemOResource.consumer_work() and routed + via schedule dataflow into this producer call. In head-paired mode, + the two output stages map to consecutive Q heads instead of consecutive + Q sequence tiles. + """ + head_coord = ( + head_coord * self.cfg.work_tile_q_heads + + self.stage_idx * self.cfg.peer_q_head_stride + ) + seq_offset_o = ( + seq_coord_q + + self.stage_idx * self.cfg.peer_q_seq_tile_stride * self.cfg.q_tile_m + ) + sO_base = self.sO_array + should_store = True + q_seq_extent = Int32(0) + if cutlass.const_expr(self.cfg.has_varlen): + if cutlass.const_expr(self.cfg.has_uniform_varlen): + cuseqlen_q = batch_coord * Int32(self.cfg.uniform_seq_len_q) + seq_end = cuseqlen_q + Int32(self.cfg.uniform_seq_len_q) + else: + cuseqlen_q = Int32(self.cum_seqlen_q[batch_coord]) + seq_end = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) + seq_offset_o = cuseqlen_q + seq_offset_o + q_seq_extent = seq_end - seq_offset_o + should_store = seq_offset_o < seq_end + + tma_copy_o_iters = self.cfg.tma_copy_o_iters + if cutlass.const_expr(self.cfg.stage_o_by_head_dim): + tma_copy_o_iters = self.cfg.tma_copy_o_stage_iters + + is_store_warp = True + if cutlass.const_expr(correction_fused): + # elect_sync elects one lane *per warp*. A correction-fused call + # runs on four warps, so only the first correction warp may enter + # the TMA issue/commit/wait body. + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + is_store_warp = warp_idx == self.cfg.correction_warp_ids[0] + if is_store_warp: + if should_store: + if prims.elect_sync(): + for i in cutlass.range_constexpr(tma_copy_o_iters): + d_offset = ( + head_dim_stage_idx * self.cfg.head_dim_per_stage_kv + + i * self.cfg.tma_copy_o_granu_inner + ) + o_coords = (d_offset, head_coord, seq_offset_o, batch_coord) + if cutlass.const_expr(self.cfg.has_varlen): + o_coords = (d_offset, head_coord, seq_offset_o) + o_coords = transform_ragged_coords( + o_coords, + ragged_dim_idx=2, + ragged_box_size=self.cfg.epi_tile[0], + ragged_extent=q_seq_extent, + ) + prims.cp_async_bulk_tensor_global_shared_cta( + self.tma_o_desc, + sO_base.subview(i * self.cfg.tma_copy_o_granu_elems), + o_coords, + ) + # should_store is CTA-uniform because it depends only on batch and + # Q tile coordinates. Keep commit paired with an actual store. + if should_store: + prims.cp_async_bulk_commit_group() + if cutlass.const_expr(self.cfg.gmem_o_store_wait_after_write): + prims.cp_async_bulk_wait_group(0, read=True) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py new file mode 100644 index 000000000000..ece2f1fe1346 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py @@ -0,0 +1,2724 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task definitions for the TS FMHA kernel. + +Tasks own ordering, not data movement bodies. Each schedule below sequences +resource waits, acquires, work calls, commits, and releases for one warp role. +The resource methods contain the actual TMA, MMA, softmax, correction, and +epilogue work. + +Schedule phase terms follow TS schedule-builder naming. HEAD is the one-time +schedule before the repeated K/V tile loop, LOOP is the repeated K/V tile body, +and TAIL is the one-time cleanup and drain after LOOP exits. +""" + +from collections.abc import Callable, Generator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + +from ..stage import FmhaStage +from cutlass.experimental.task_scheduling.schedule_builder import ( + domain_loop, + schedule, + work_tile_loop, +) +from cutlass.experimental.task_scheduling.resources import MemoryResource, WorkQueue +from cutlass.experimental.task_scheduling.task import Task + +from .fmha_resources import ( + FmhaConfig, + GmemOResource, + GmemQKVResource, + S0S1SequenceResource, + SmemKVResource, + SmemOResource, + SmemPageOffsetsKvResource, + SmemQResource, + TmemOResource, + TmemPResource, + TmemSPResource, + TmemStatsResource, + TmemStatsDoneResource, +) + + +@dataclass(kw_only=True) +class PackedContextWorkQueue(WorkQueue): + """Persistent queue that skips Q tiles outside a live packed request.""" + + cfg: cutlass.Constexpr[FmhaConfig] = field(init=False, default=None) + cum_seqlen_q: Any = field(init=False, default=None) + + def __init__( + self, + cfg: FmhaConfig, + cum_seqlen_q: cute.Tensor, + **kwargs: Any, + ) -> None: + """Attach the live packed-Q metadata used by the skip predicate.""" + super().__init__(**kwargs) + self.cfg = cfg + self.cum_seqlen_q = cum_seqlen_q + + @cute.jit + def skip_work_tile_if(self, work_tile: Any) -> cutlass.Boolean: + """Skip a scheduler tile whose first Q row is outside its request.""" + seq_idx, _, batch_idx = self.cfg.work_tile_coord_indices + seq_coord = Int32(work_tile.tile_idx[seq_idx]) + if cutlass.const_expr(self.cfg.uses_causal_reversed_head_batch_seq_tile_order): + seq_coord = Int32(self.cfg.num_seq_tiles) - seq_coord - Int32(1) + batch_coord = Int32(work_tile.tile_idx[batch_idx]) + q_begin = Int32(self.cum_seqlen_q[batch_coord]) + q_end = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) + seqlen_q = q_end - q_begin + return seq_coord * Int32(self.cfg.cta_tiler[0]) >= seqlen_q + + +def _persistent_tail(work_queue: WorkQueue) -> None: + """Advance and release the persistent work tile after one task body.""" + work_queue.wait() + work_queue.get_and_advance_work_tile() + work_queue.release() + + +def _src_resources( + *resources: MemoryResource, + work_queue: WorkQueue | None, +) -> list[MemoryResource]: + """Build a task source-resource list, including WorkQueue when present.""" + src = list(resources) + if work_queue is not None: + src.append(work_queue) + return src + + +def _schedule_with_work_queue( + schedule: Callable[..., object], + *resources: MemoryResource, + work_queue: WorkQueue | None, +) -> object: + """Invoke a captured schedule with the optional WorkQueue argument.""" + if work_queue is None: + return schedule(*resources) + return schedule(*resources, work_queue) + + +def _packed_context_skip_predicate( + work_queue: WorkQueue | None, +) -> Callable[..., object] | None: + """Select the live-Q skip predicate before schedule capture creates proxies.""" + if isinstance(work_queue, PackedContextWorkQueue): + return PackedContextWorkQueue.skip_work_tile_if + return None + + +@contextmanager +def _work_tile_schedule_loop( + work_queue: WorkQueue | None, + *, + skip_if: Callable[..., object] | None = None, +) -> Generator[object | None, None, None]: + """Wrap a task body once per persistent work tile, or once for static schedules.""" + if skip_if is not None: + assert work_queue is not None + with work_tile_loop( + work_queue, + skip_if=skip_if, + ) as work_tiles: + with work_tiles.skippable(): + yield work_tiles + # Every fetched tile, including a skipped one, must advance and + # release the queue exactly once so persistent workers converge. + _persistent_tail(work_queue) + elif work_queue is not None: + with work_tile_loop(work_queue) as work_tile: + yield work_tile + _persistent_tail(work_queue) + else: + yield None + + +def _captured_loop_bounds( + task_class: type[Task], + task_kwargs: dict[str, object], +) -> tuple[object, object, object]: + """Infer ``(start, end, step)`` loop bounds for a captured schedule. + + Dense schedules pass a static ``domain``; causal schedules pass + ``num_kv_tiles`` and use the task class's ``get_domain`` as a dynamic end. + """ + loop_start = task_kwargs.pop("domain_start", 0) + loop_step = task_kwargs.pop("step", 1) + loop_end = task_kwargs.pop("domain", None) + if loop_end is None: + if "num_kv_tiles" not in task_kwargs: + raise ValueError( + "create_*_task requires a 'domain' or 'num_kv_tiles' kwarg to " + "determine the loop end." + ) + loop_end = task_class.get_domain + return loop_start, loop_end, loop_step + + +def create_load_task( + gmem_qkv: GmemQKVResource, + smem_q: SmemQResource, + smem_kv: SmemKVResource, + work_queue: WorkQueue | None, + task_class: type[Task] = Task, + smem_page_offsets_kv: SmemPageOffsetsKvResource | None = None, + smem_page_offsets_v: SmemPageOffsetsKvResource | None = None, + **task_kwargs: Any, +) -> Task: + """Create the one-warp TMA load task. + + When ``smem_page_offsets_kv`` is provided, each K/V TMA load consumes page + IDs prefetched by the auxiliary warp through the ordinary asynchronous + page-offset pipeline. + """ + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + skip_work_tile_if = _packed_context_skip_predicate(work_queue) + src = _src_resources(gmem_qkv, work_queue=work_queue) + dst = [smem_q, smem_kv] + if smem_page_offsets_kv is not None: + src.append(smem_page_offsets_kv) + if smem_page_offsets_v is not None: + src.append(smem_page_offsets_v) + if smem_q.cfg.single_qkv_instance and smem_q.cfg.has_tmem_p_pipeline: + num_head_dim_stages_k = smem_kv.cfg.num_head_dim_stages_k + num_head_dim_stages_v = smem_kv.cfg.num_head_dim_stages_v + + if smem_page_offsets_v is not None: + if smem_page_offsets_kv is None: + raise ValueError("a V page window requires a matching K page window") + pages_per_tile = smem_kv.cfg.kv_tile_n // smem_kv.cfg.num_tokens_per_page + page_window_period = smem_kv.cfg.page_table_window_entries // pages_per_tile + if ( + not isinstance(loop_start, int) + or not isinstance(loop_end, int) + or not isinstance(loop_step, int) + or loop_start != 0 + or loop_step != 1 + or loop_end < page_window_period + or loop_end % page_window_period != 0 + ): + raise ValueError( + "reused page windows require a compile-time K/V domain " + "divisible by the topology-derived page-window period" + ) + + def load_reused_page_windows_schedule_body( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + spok: SmemPageOffsetsKvResource, + spov: SmemPageOffsetsKvResource, + wq: WorkQueue | None, + ) -> None: + """Load staged K/V while retaining each page-ID window.""" + sq.init_load_state() + skv.init_load_state() + spok.init_read_state() + cached_v_page_ids = spov.init_cached_read_state() + + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + ( + _seq_coord, + head_coord, + kv_head_coord, + _head_coord_kv, + batch_coord, + seq_coord_q, + cuseqlen_q, + cuseqlen_k, + seqlen_q, + seqlen_k, + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) = gqkv.compute_coords() + sq.acquire() + sq.tma_load( + seq_coord_q=seq_coord_q, + head_coord=head_coord, + batch_coord=batch_coord, + cuseqlen_q=cuseqlen_q, + seqlen_q=seqlen_q, + inst_idx=0, + ) + sq.commit() + + def load_k_tile(*, tile_offset: int) -> None: + for head_dim_stage_idx in range(num_head_dim_stages_k): + skv.try_acquire() + skv.acquire() + skv.k_load_stage( + stage_id=head_dim_stage_idx, + tile_offset=tile_offset, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + + def cache_v_tile(*, tile_offset: int) -> None: + nonlocal cached_v_page_ids + cached_v_page_ids = spov.cache_tile_page_ids( + cached_page_ids=cached_v_page_ids, + kv_tile_start=kv_tile_start, + tile_offset=tile_offset, + ) + + def load_v_tile( + *, tile_offset: int, reuse_cached_page_ids: bool = False + ) -> None: + for head_dim_stage_idx in range(num_head_dim_stages_v): + skv.try_acquire() + skv.acquire() + if reuse_cached_page_ids: + skv.v_load_stage_cached( + cached_v_page_ids=cached_v_page_ids, + stage_id=head_dim_stage_idx, + tile_offset=tile_offset, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + else: + skv.v_load_stage( + stage_id=head_dim_stage_idx, + tile_offset=tile_offset, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + + # Window zero: K stays one tile ahead of V. Cache the final + # V IDs before releasing the window because its last V + # tile is delayed across the boundary. + spok.wait() + spok.read_offsets() + load_k_tile(tile_offset=0) + load_k_tile(tile_offset=1) + spov.wait() + load_v_tile(tile_offset=0) + for tile_delta in range(2, page_window_period - 1): + load_k_tile(tile_offset=tile_delta) + load_v_tile(tile_offset=tile_delta - 1) + load_k_tile(tile_offset=page_window_period - 1) + spok.release() + load_v_tile(tile_offset=page_window_period - 2) + cache_v_tile(tile_offset=page_window_period - 1) + spov.release() + + # Each structural iteration consumes one complete K/V page + # window. Only register page IDs cross the loop boundary. + with domain_loop( + page_window_period, + loop_end, + page_window_period, + ): + spok.wait() + spok.read_offsets() + load_k_tile(tile_offset=0) + load_v_tile(tile_offset=-1, reuse_cached_page_ids=True) + load_k_tile(tile_offset=1) + spov.wait() + load_v_tile(tile_offset=0) + for tile_delta in range(2, page_window_period - 1): + load_k_tile(tile_offset=tile_delta) + load_v_tile(tile_offset=tile_delta - 1) + load_k_tile(tile_offset=page_window_period - 1) + spok.release() + load_v_tile(tile_offset=page_window_period - 2) + cache_v_tile(tile_offset=page_window_period - 1) + spov.release() + + load_v_tile( + tile_offset=page_window_period - 1, + reuse_cached_page_ids=True, + ) + + @schedule + def load_reused_page_windows_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + spok: SmemPageOffsetsKvResource, + spov: SmemPageOffsetsKvResource, + wq: WorkQueue | None = None, + ) -> None: + load_reused_page_windows_schedule_body(gqkv, sq, skv, spok, spov, wq) + + captured_schedule = _schedule_with_work_queue( + load_reused_page_windows_schedule, + gmem_qkv, + smem_q, + smem_kv, + smem_page_offsets_kv, + smem_page_offsets_v, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=smem_kv.cfg.load_warp_id, + num_warps=1, + schedule=captured_schedule, + num_registers=smem_kv.cfg.num_regs_other, + name="LoadTask", + **task_kwargs, + ) + + def load_schedule_body( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + spo: SmemPageOffsetsKvResource | None, + wq: WorkQueue | None, + ) -> None: + sq.init_load_state() + skv.init_load_state() + if spo is not None: + spo.init_read_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + coords = gqkv.compute_coords() + ( + _seq_coord, + head_coord, + kv_head_coord, + _head_coord_kv, + batch_coord, + seq_coord_q, + cuseqlen_q, + cuseqlen_k, + seqlen_q, + seqlen_k, + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) = coords + sq.acquire() + sq.tma_load( + seq_coord_q=seq_coord_q, + head_coord=head_coord, + batch_coord=batch_coord, + cuseqlen_q=cuseqlen_q, + seqlen_q=seqlen_q, + inst_idx=0, + ) + sq.commit() + + for head_dim_stage_idx in range(num_head_dim_stages_k): + skv.try_acquire() + if spo is not None and head_dim_stage_idx == 0: + spo.wait() + spo.read_offsets() + skv.acquire() + skv.k_load_stage( + stage_id=head_dim_stage_idx, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + if spo is not None: + spo.release() + + with domain_loop(loop_start + 1, loop_end, loop_step): + for head_dim_stage_idx in range(num_head_dim_stages_k): + skv.try_acquire() + if spo is not None and head_dim_stage_idx == 0: + spo.wait() + spo.read_offsets() + skv.acquire() + skv.k_load_stage( + stage_id=head_dim_stage_idx, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + if spo is not None: + spo.release() + + for head_dim_stage_idx in range(num_head_dim_stages_v): + skv.try_acquire() + if spo is not None and head_dim_stage_idx == 0: + spo.wait() + spo.read_offsets() + skv.acquire() + skv.v_load_stage( + stage_id=head_dim_stage_idx, + previous=True, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + if spo is not None: + spo.release() + + for head_dim_stage_idx in range(num_head_dim_stages_v): + skv.try_acquire() + if spo is not None and head_dim_stage_idx == 0: + spo.wait() + spo.read_offsets() + skv.acquire() + skv.v_load_stage( + stage_id=head_dim_stage_idx, + previous=False, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + if spo is not None: + spo.release() + + @schedule + def load_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + wq: WorkQueue | None = None, + ) -> None: + load_schedule_body(gqkv, sq, skv, None, wq) + + @schedule + def load_page_offsets_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + spo: SmemPageOffsetsKvResource, + wq: WorkQueue | None = None, + ) -> None: + load_schedule_body(gqkv, sq, skv, spo, wq) + + if smem_page_offsets_kv is None: + captured_schedule = _schedule_with_work_queue( + load_schedule, gmem_qkv, smem_q, smem_kv, work_queue=work_queue + ) + else: + captured_schedule = _schedule_with_work_queue( + load_page_offsets_schedule, + gmem_qkv, + smem_q, + smem_kv, + smem_page_offsets_kv, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=smem_kv.cfg.load_warp_id, + num_warps=1, + schedule=captured_schedule, + num_registers=smem_kv.cfg.num_regs_other, + name="LoadTask", + **task_kwargs, + ) + + if smem_q.cfg.single_qkv_instance: + raise ValueError("single-instance context requires the staged TMEM-P topology") + if smem_page_offsets_kv is not None or smem_page_offsets_v is not None: + raise ValueError("paired context resolves paged K/V IDs directly") + + def load_schedule_body( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + wq: WorkQueue | None, + ) -> None: + """Load paired Q instances and their directly addressed K/V tiles.""" + sq.init_load_state() + skv.init_load_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): # noqa: SIM117 + # The first K-loop iteration also loads Q0/Q1. Later iterations + # only stream the next K/V tiles through the SmemKV pipeline. + with domain_loop(loop_start, loop_end, loop_step) as d: + with d.first_iter(): + ( + _seq_coord, + head_coord, + kv_head_coord, + _head_coord_kv, + batch_coord, + seq_coord_q, + cuseqlen_q, + cuseqlen_k, + seqlen_q, + seqlen_k, + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) = gqkv.compute_coords() + # Load Q0 for the first Q tile in this work tile. + sq.acquire() + sq.tma_load( + seq_coord_q=seq_coord_q, + head_coord=head_coord, + batch_coord=batch_coord, + cuseqlen_q=cuseqlen_q, + seqlen_q=seqlen_q, + inst_idx=0, + ) + sq.commit() + # Throttle TMA before reserving a KV stage. + skv.try_acquire() + # Load Ki, with K0 handled by the first iteration. + skv.acquire() + skv.k_load( + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + with d.first_iter(): + # Load Q1 for the second Q tile in this work tile. + sq.acquire() + sq.tma_load( + seq_coord_q=seq_coord_q, + head_coord=head_coord, + batch_coord=batch_coord, + cuseqlen_q=cuseqlen_q, + seqlen_q=seqlen_q, + inst_idx=1, + ) + sq.commit() + # Throttle TMA before reserving a KV stage. + skv.try_acquire() + # Load Vi, with V0 handled by the first iteration. + skv.acquire() + skv.v_load( + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + seqlen_k=seqlen_k, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + skv.commit() + + @schedule + def load_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + wq: WorkQueue | None = None, + ) -> None: + """Contiguous-KV captured schedule.""" + # Mypy retains the earlier branch's five-argument closure signature. + load_schedule_body(gqkv, sq, skv, wq) # type: ignore[call-arg] + + captured_schedule = _schedule_with_work_queue( + load_schedule, gmem_qkv, smem_q, smem_kv, work_queue=work_queue + ) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=smem_kv.cfg.load_warp_id, + num_warps=1, + schedule=captured_schedule, + num_registers=smem_kv.cfg.num_regs_other, + name="LoadTask", + **task_kwargs, + ) + + +def create_mma_task( + gmem_qkv: GmemQKVResource, + smem_q: SmemQResource, + smem_kv: SmemKVResource, + tmem_sp0: TmemSPResource, + tmem_sp1: TmemSPResource | None, + tmem_p0: TmemPResource | None, + tmem_o: TmemOResource, + tmem_vec_done_0: TmemStatsDoneResource, + tmem_vec_done_1: TmemStatsDoneResource | None, + work_queue: WorkQueue | None, + task_class: type[Task] = Task, + **task_kwargs: Any, +) -> Task: + """Create the one-warp MMA compute task.""" + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + skip_work_tile_if = _packed_context_skip_predicate(work_queue) + src = _src_resources(gmem_qkv, smem_q, smem_kv, work_queue=work_queue) + + if ( + smem_q.cfg.single_qkv_instance + and smem_q.cfg.has_tmem_p_pipeline + and tmem_p0 is not None + ): + split_src = _src_resources( + gmem_qkv, smem_q, smem_kv, tmem_p0, work_queue=work_queue + ) + num_head_dim_stages_k = smem_kv.cfg.num_head_dim_stages_k + num_head_dim_stages_v = smem_kv.cfg.num_head_dim_stages_v + loop_carried_head_dim_stages = 2 + if ( + num_head_dim_stages_k != loop_carried_head_dim_stages + or num_head_dim_stages_v != loop_carried_head_dim_stages + ): + raise ValueError("loop-carried split S/P scheduling expects two K/V stages") + + @schedule + def mma_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + tp0: TmemPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + sq.init_descriptor_state() + skv.init_descriptor_state() + sp0.init_mma_state() + to.init_mma_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + sp0.init_mma_work_tile_state() + to.init_mma_work_tile_state() + v_seqlen_k = Int32(0) + v_kv_tile_start = Int32(0) + if cutlass.const_expr(smem_q.cfg.use_paged_kv): + ( + _seq_coord, + _head_coord, + _kv_head_coord, + _head_coord_kv, + _batch_coord, + _seq_coord_q, + _cuseqlen_q, + _cuseqlen_k, + _seqlen_q, + v_seqlen_k, + v_kv_tile_start, + _kv_request_begin, + _kv_page_idx_ub, + ) = gqkv.compute_coords() + + sq.wait() + desc_q0_base = sq.q0_desc(inst_idx=0) + if not smem_q.cfg.stats_via_smem: + vd0.acquire() + sp0.acquire() + for head_dim_stage_idx in range(num_head_dim_stages_k): + skv.wait() + desc_k_base = skv.k_desc() + sp0.qk_mma( + desc_q_base=desc_q0_base, + desc_k_base=desc_k_base, + section=FmhaStage.Head, + head_dim_stage_idx=head_dim_stage_idx, + ) + skv.release() + sp0.commit() + if not smem_q.cfg.stats_via_smem: + vd0.commit() + sp0.acquire() + + # Loop offset i is local to the steady state. LoadTask has + # already advanced K by one tile, so these waits consume K(i+1) + # for QK and V(i) for PV. + with domain_loop(loop_start, loop_end, loop_step): + if not smem_q.cfg.stats_via_smem: + vd0.acquire() + skv.wait() + desc_k_base = skv.k_desc() + sp0.qk_mma( + desc_q_base=desc_q0_base, + desc_k_base=desc_k_base, + section=FmhaStage.Loop, + head_dim_stage_idx=0, + ) + skv.release() + + skv.wait() + desc_k_base = skv.k_desc() + sp0.qk_mma( + desc_q_base=desc_q0_base, + desc_k_base=desc_k_base, + section=FmhaStage.Loop, + head_dim_stage_idx=1, + ) + skv.release() + sp0.commit() + if not smem_q.cfg.stats_via_smem: + vd0.commit() + sp0.acquire() + + to.acquire() + tp0.wait() + tmem_p_base = tp0.p_base() + to.set_p_base(tmem_p_base=tmem_p_base) + + skv.wait() + if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): + desc_v_base = skv.v_desc_paged( + section=FmhaStage.Loop, + seqlen_k=v_seqlen_k, + kv_tile_start=v_kv_tile_start, + ) + else: + desc_v_base = skv.v_desc() + to.pv_mma( + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + head_dim_stage_idx=0, + ) + skv.release() + + skv.wait() + if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): + desc_v_base = skv.v_desc_paged( + section=FmhaStage.Loop, + seqlen_k=v_seqlen_k, + kv_tile_start=v_kv_tile_start, + ) + else: + desc_v_base = skv.v_desc() + to.pv_mma( + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + head_dim_stage_idx=1, + ) + skv.release() + to.commit() + tp0.release() + + sq.release() + to.acquire() + tp0.wait() + tmem_p_base = tp0.p_base() + to.set_p_base(tmem_p_base=tmem_p_base) + for head_dim_stage_idx in range(num_head_dim_stages_v): + skv.wait() + if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): + desc_v_base = skv.v_desc_paged( + section=FmhaStage.Tail, + seqlen_k=v_seqlen_k, + kv_tile_start=v_kv_tile_start, + ) + else: + desc_v_base = skv.v_desc() + to.pv_mma( + desc_v_base=desc_v_base, + section=FmhaStage.Tail, + head_dim_stage_idx=head_dim_stage_idx, + is_tail=True, + ) + skv.release() + to.commit() + tp0.release() + if not smem_q.cfg.stats_via_smem: + vd0.acquire() + sp0.commit() + if not smem_q.cfg.stats_via_smem: + vd0.commit() + tp0.wait() + tp0.release() + + captured_schedule = _schedule_with_work_queue( + mma_schedule, + gmem_qkv, + smem_q, + smem_kv, + tmem_sp0, + tmem_p0, + tmem_o, + tmem_vec_done_0, + work_queue=work_queue, + ) + return task_class( + src_resources=split_src, + dst_resources=[tmem_sp0, tmem_o] + + ([] if smem_q.cfg.stats_via_smem else [tmem_vec_done_0]), + warp_idx=smem_q.cfg.mma_warp_id, + num_warps=1, + schedule=captured_schedule, + name="MmaTask", + num_registers=smem_q.cfg.num_regs_other, + **task_kwargs, + ) + + if smem_q.cfg.single_qkv_instance: + num_head_dim_stages_k = smem_kv.cfg.num_head_dim_stages_k + num_head_dim_stages_v = smem_kv.cfg.num_head_dim_stages_v + + @schedule + def mma_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + desc_q0_base, _desc_q1_base = sq.create_function_variables() + desc_k_base, desc_v_base = skv.create_function_variables() + sp0.create_function_variables() + to.create_function_variables() + vd0.create_function_variables() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + v_seqlen_k = Int32(0) + v_kv_tile_start = Int32(0) + if cutlass.const_expr(smem_q.cfg.use_paged_kv): + ( + _seq_coord, + _head_coord, + _kv_head_coord, + _head_coord_kv, + _batch_coord, + _seq_coord_q, + _cuseqlen_q, + _cuseqlen_k, + _seqlen_q, + v_seqlen_k, + v_kv_tile_start, + _kv_request_begin, + _kv_page_idx_ub, + ) = gqkv.compute_coords() + if wq is not None: + sp0.create_work_tile_variables() + to.create_work_tile_variables() + + with domain_loop(loop_start, loop_end, loop_step) as d: + with d.first_iter(): + sq.wait() + desc_q0_base = sq.q0_desc(inst_idx=0) + vd0.acquire() + sp0.acquire() + for head_dim_stage_idx in range(num_head_dim_stages_k): + skv.wait() + desc_k_base = skv.k_desc() + sp0.qk_mma( + desc_q_base=desc_q0_base, + desc_k_base=desc_k_base, + section=FmhaStage.Loop, + head_dim_stage_idx=head_dim_stage_idx, + ) + skv.release() + sp0.commit() + with d.first_iter(): + vd0.commit() + to.acquire() + sp0.acquire() + sp0.p_read() + for head_dim_stage_idx in range(num_head_dim_stages_v): + skv.wait() + if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): + desc_v_base = skv.v_desc_paged( + section=FmhaStage.Loop, + seqlen_k=v_seqlen_k, + kv_tile_start=v_kv_tile_start, + ) + else: + desc_v_base = skv.v_desc() + to.pv_mma( + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + head_dim_stage_idx=head_dim_stage_idx, + ) + skv.release() + to.commit() + + sq.release() + sp0.commit() + + captured_schedule = _schedule_with_work_queue( + mma_schedule, + gmem_qkv, + smem_q, + smem_kv, + tmem_sp0, + tmem_o, + tmem_vec_done_0, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=[tmem_sp0, tmem_o, tmem_vec_done_0], + warp_idx=smem_q.cfg.mma_warp_id, + num_warps=1, + schedule=captured_schedule, + name="MmaTask", + num_registers=smem_q.cfg.num_regs_other, + **task_kwargs, + ) + + if tmem_sp1 is None or tmem_vec_done_1 is None: + raise ValueError("paired MMA scheduling requires peer-1 resources") + + @schedule + def mma_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Captured schedule for interleaved QK and PV MMA work.""" + sq.init_descriptor_state() + skv.init_descriptor_state() + sp0.init_mma_state() + sp1.init_mma_state() + to.init_mma_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + v_seqlen_k = Int32(0) + v_kv_tile_start = Int32(0) + if cutlass.const_expr(smem_q.cfg.use_paged_kv): + ( + _seq_coord, + _head_coord, + _kv_head_coord, + _head_coord_kv, + _batch_coord, + _seq_coord_q, + _cuseqlen_q, + _cuseqlen_k, + _seqlen_q, + v_seqlen_k, + v_kv_tile_start, + _kv_request_begin, + _kv_page_idx_ub, + ) = gqkv.compute_coords() + # HEAD: consume Q0, K0, Q1, and V0. TmemStatsDone starts empty, so + # the first acquire succeeds without priming. On later work tiles, + # correction has released the previous stats slot. + # + # Consume Q0, K0, then QK(Q0,K0)→S0. + sq.wait() + desc_q0_base = sq.q0_desc(inst_idx=0) + skv.wait() + desc_k_base = skv.k_desc() + if not smem_q.cfg.stats_via_smem: + vd0.acquire() + sp0.acquire() + sp0.qk_mma( + desc_q_base=desc_q0_base, + desc_k_base=desc_k_base, + section=FmhaStage.Head, + ) + sp0.commit() + if not smem_q.cfg.stats_via_smem: + vd0.commit() + # Consume Q1, then QK(Q1,K0)→S1. + sq.wait() + desc_q1_base = sq.q1_desc(inst_idx=1) + if not smem_q.cfg.stats_via_smem: + vd1.acquire() + sp1.acquire() + sp1.qk_mma( + desc_q_base=desc_q1_base, + desc_k_base=desc_k_base, + section=FmhaStage.Head, + ) + sp1.commit() + if not smem_q.cfg.stats_via_smem: + vd1.commit() + # Q0/Q1 stay live because UMMA reads Q throughout the K-loop. + # Release K0 (done with QK→S0 and QK→S1), then consume V0. + skv.release() + skv.wait() + if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): + desc_v_base = skv.v_desc_paged( + section=FmhaStage.Head, + seqlen_k=v_seqlen_k, + kv_tile_start=v_kv_tile_start, + ) + else: + desc_v_base = skv.v_desc() + # Acquire O first (off critical path), then acquire SP0 and run PV→O0. + to.acquire() + sp0.acquire() + sp0.p_read() + to.pv_mma(desc_v_base=desc_v_base, section=FmhaStage.Head) + to.commit() + + # LOOP: interleave QK and PV work while preserving the previous V + # tile until its PV MMA has consumed it: + # QK0(deferred commit) -> PV1(V_prev, release V_prev) -> + # QK1(commit) -> release Ki+1 -> wait Vi+1 -> PV0(no commit) + with domain_loop(loop_start, loop_end, loop_step): + skv.wait() + desc_k_base = skv.k_desc() + # QK0: QK(Q0,Ki+1) → S0 (no acquire; handle held from PV0). + sp0.qk_mma( + desc_q_base=desc_q0_base, + desc_k_base=desc_k_base, + section=FmhaStage.Loop, + ) + sp0.commit() + # PV1(V_prev): P1 * V_prev → O1. + to.acquire() + sp1.acquire() + sp1.p_read() + to.pv_mma(desc_v_base=desc_v_base, section=FmhaStage.Loop) + to.commit() + # Release V_prev after PV1 UMMA consumed SMEM data. + skv.release() + # QK1: QK(Q1,Ki+1) → S1 (no acquire; handle held from PV1). + sp1.qk_mma( + desc_q_base=desc_q1_base, + desc_k_base=desc_k_base, + section=FmhaStage.Loop, + ) + sp1.commit() + # Release Ki+1, then wait Vi+1. + skv.release() + skv.wait() + if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): + desc_v_base = skv.v_desc_paged( + section=FmhaStage.Loop, + tile_offset=1, + seqlen_k=v_seqlen_k, + kv_tile_start=v_kv_tile_start, + ) + else: + desc_v_base = skv.v_desc() + # PV0: P0 * Vi+1 → O0. + to.acquire() + sp0.acquire() + sp0.p_read() + to.pv_mma( + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + inst_idx=1, + ) + to.commit() + + # TAIL: release Qs, close the deferred SP state, and run the final + # PV→O1 MMA. + sq.release() + sq.release() + sp0.commit() + to.acquire() + sp1.acquire() + sp1.p_read() + to.pv_mma( + desc_v_base=desc_v_base, + section=FmhaStage.Tail, + is_tail=True, + ) + to.commit() + skv.release() + sp1.commit() + + captured_schedule = _schedule_with_work_queue( + mma_schedule, + gmem_qkv, + smem_q, + smem_kv, + tmem_sp0, + tmem_sp1, + tmem_o, + tmem_vec_done_0, + tmem_vec_done_1, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=[tmem_sp0, tmem_sp1, tmem_o] + + ([] if smem_q.cfg.stats_via_smem else [tmem_vec_done_0, tmem_vec_done_1]), + warp_idx=12, + num_warps=1, + schedule=captured_schedule, + name="MmaTask", + num_registers=smem_q.cfg.num_regs_other, + **task_kwargs, + ) + + +def create_softmax_task( + index: int, + tmem_sp: TmemSPResource, + tmem_vec: TmemStatsResource, + tmem_p: TmemPResource | None, + s0s1_seq: S0S1SequenceResource | None, + work_queue: WorkQueue | None, + task_class: type[Task] = Task, + **task_kwargs: Any, +) -> Task: + """Create a four-warp Softmax task. + + index=0: warps 0-3 (Softmax0Task) — S0-S1 producer (acquire/commit) + index=1: warps 4-7 (Softmax1Task) — S0-S1 consumer (wait/release) + + Args: + task_class: Task subclass used to instantiate the softmax schedule. + """ + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + skip_work_tile_if = _packed_context_skip_predicate(work_queue) + + # A missing S0-S1 sequence resource selects the single-QKV-instance path. + # For D>128, the TMEM P pipeline gives S and P independent readiness + # handoffs. MMA can issue next-tile QK on one stage while previous-tile PV + # uses the other. + if s0s1_seq is None: + if tmem_p is not None and tmem_sp.cfg.has_tmem_p_pipeline: + src = _src_resources(tmem_sp, work_queue=work_queue) + dst = [tmem_vec, tmem_p] + + @schedule + def softmax_schedule( + sp: TmemSPResource, + vec: TmemStatsResource, + tp: TmemPResource, + wq: WorkQueue | None = None, + ) -> None: + p_chunk = sp.init_softmax_state() + scale_softmax_log2 = sp.load_scale_softmax_log2() + vec.init_store_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + old_row_max, row_max, row_sum, q_offset = ( + sp.init_softmax_work_tile_state() + ) + vec.init_store_work_tile_state() + if tmem_sp.uses_varlen_q_offset_cache: + q_offset = sp.cache_q_offset() + if tmem_sp.uses_packed_dense_k_mask: + seqlen_k = sp.cache_seqlen_k() + window_start = Int32(0) + window_end = Int32(0) + if tmem_sp.uses_variable_window: + window_start, window_end = sp.cache_variable_window_bounds() + vec.acquire() + with domain_loop(loop_start, loop_end, loop_step): + # Softmax(i): wait for QK(Q,Ki) -> S(i). + sp.wait() + if tmem_sp.uses_variable_window: + old_row_max, row_max = sp.variable_window_row_max( + row_max=row_max, + window_start=window_start, + window_end=window_end, + ) + elif tmem_sp.uses_left_window_loop_mask: + old_row_max, row_max = sp.left_masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + elif tmem_sp.uses_varlen_loop_right_mask: + old_row_max, row_max = sp.right_masked_row_max( + row_max=row_max, + q_offset=q_offset, + section=FmhaStage.Loop, + ) + elif tmem_sp.uses_query_paired_q_offset_loop_mask: + old_row_max, row_max = sp.loop_masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + elif tmem_sp.uses_fixed_dense_k_tail_mask: + old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max( + row_max=row_max, + ) + elif tmem_sp.uses_packed_dense_k_mask: + old_row_max, row_max = sp.packed_dense_k_masked_row_max( + row_max=row_max, + seqlen_k=seqlen_k, + section=FmhaStage.Loop, + ) + else: + old_row_max, row_max = sp.compute_row_max(row_max=row_max) + # Stats(i): S(i) -> row max/sum for correction. + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + # P(i): acquire the matching P-ready handoff stage. + tp.acquire() + # P(i): exp2(S(i)) -> P(i) in the same TMEM stage. + p_chunk = sp.exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + # P(i) ready: P(i) -> PV(Pi,Vi). + tp.commit() + # S/P(i): release softmax ownership for the next QK stage. + sp.release() + # Aux(i): finish the row-sum reduction after releasing SP. + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + vec.acquire() + + if tmem_sp.uses_head_paired_causal_tail_mask: + # Tail S: consume and mask the final head-paired score tile. + sp.wait() + old_row_max, row_max = sp.right_masked_row_max( + row_max=row_max, + q_offset=q_offset, + section=FmhaStage.Tail, + ) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + # Tail P: publish the final probability tile to PV. + tp.acquire() + p_chunk = sp.exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + tp.commit() + sp.release() + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + vec.acquire() + # Drain the final SP/P-ready slots and publish identity stats. + sp.wait() + sp.release() + tp.acquire() + tp.commit() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + elif tmem_sp.uses_query_paired_causal_tail_mask: + # Tail S: consume and mask the final query-paired score tile. + sp.wait() + old_row_max, row_max = sp.masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + # Tail P: publish the final probability tile to PV. + tp.acquire() + p_chunk = sp.masked_exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + tp.commit() + sp.release() + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + if tmem_sp.uses_query_paired_invalid_tail: + # Invalid peer tail: consume its padded SP slot without PV. + sp.wait() + old_row_max, row_max = sp.invalid_row_max(row_max=row_max) + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + sp.invalid_exp2_p(row_max=row_max) + sp.release() + # Drain the final SP/P-ready slots and publish identity stats. + sp.wait() + sp.release() + tp.acquire() + tp.commit() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + elif tmem_sp.cfg.is_causal: + # Tail S: consume and causally mask the final score tile. + sp.wait() + old_row_max, row_max = sp.masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + # Tail P: publish the final probability tile to PV. + tp.acquire() + p_chunk = sp.masked_exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + tp.commit() + sp.release() + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + # Drain the final SP/P-ready slots and publish identity stats. + sp.wait() + sp.release() + tp.acquire() + tp.commit() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + else: + # Non-causal cleanup: drain SP and the matching P-ready slot. + sp.wait() + sp.release() + tp.acquire() + tp.commit() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + if tmem_sp.cfg.stats_via_smem: + # Balance the two-stage stats cursor before the next + # captured persistent work tile. The context task + # runtime carries pipeline state across work tiles, + # while each tile's static call layout begins at the + # same stage; the empty record keeps both in phase. + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + + captured_schedule = _schedule_with_work_queue( + softmax_schedule, tmem_sp, tmem_vec, tmem_p, work_queue=work_queue + ) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=index * 4, + num_warps=4, + schedule=captured_schedule, + num_registers=tmem_sp.cfg.num_regs_softmax, + name=f"Softmax{index}Task", + **task_kwargs, + ) + + # Non-split single-instance fallback: softmax writes P into the current + # SP stage and releases that same resource for MMA to consume directly. + src = _src_resources(tmem_sp, work_queue=work_queue) + dst = [tmem_vec] + + @schedule + def softmax_schedule( + sp: TmemSPResource, + vec: TmemStatsResource, + wq: WorkQueue | None = None, + ) -> None: + old_row_max, row_max, row_sum, p_chunk, q_offset = ( + sp.create_function_variables() + ) + vec.create_function_variables() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + if wq is not None: + old_row_max, row_max, row_sum, q_offset = ( + sp.create_work_tile_variables( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + q_offset=q_offset, + ) + ) + vec.create_work_tile_variables() + if tmem_sp.uses_varlen_q_offset_cache: + q_offset = sp.cache_q_offset() + if tmem_sp.uses_packed_dense_k_mask: + seqlen_k = sp.cache_seqlen_k() + window_start = Int32(0) + window_end = Int32(0) + if tmem_sp.uses_variable_window: + window_start, window_end = sp.cache_variable_window_bounds() + vec.acquire() + with domain_loop(loop_start, loop_end, loop_step): + sp.wait() + if tmem_sp.uses_variable_window: + old_row_max, row_max = sp.variable_window_row_max( + row_max=row_max, + window_start=window_start, + window_end=window_end, + ) + elif tmem_sp.uses_left_window_loop_mask: + old_row_max, row_max = sp.left_masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + elif tmem_sp.uses_varlen_loop_right_mask: + old_row_max, row_max = sp.right_masked_row_max( + row_max=row_max, + q_offset=q_offset, + section=FmhaStage.Loop, + ) + elif tmem_sp.uses_query_paired_q_offset_loop_mask: + old_row_max, row_max = sp.loop_masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + elif tmem_sp.uses_fixed_dense_k_tail_mask: + old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max( + row_max=row_max + ) + elif tmem_sp.uses_packed_dense_k_mask: + old_row_max, row_max = sp.packed_dense_k_masked_row_max( + row_max=row_max, + seqlen_k=seqlen_k, + section=FmhaStage.Loop, + ) + else: + old_row_max, row_max = sp.row_max(row_max) + vec.store_vec( + old_row_max, + row_max, + row_sum, + ) + vec.commit() + p_chunk = sp.exp2_p(row_max) + sp.release() + row_sum = sp.softmax_post_release_reduce( + old_row_max, row_max, row_sum, p_chunk + ) + vec.acquire() + + if tmem_sp.uses_head_paired_causal_tail_mask: + sp.wait() + old_row_max, row_max = sp.right_masked_row_max( + row_max=row_max, + q_offset=q_offset, + section=FmhaStage.Tail, + ) + vec.store_vec( + old_row_max, + row_max, + row_sum, + ) + vec.commit() + p_chunk = sp.exp2_p(row_max) + sp.release() + row_sum = sp.softmax_post_release_reduce( + old_row_max, row_max, row_sum, p_chunk + ) + vec.acquire() + sp.wait() + sp.release() + old_row_max = sp.softmax_post_release_identity(row_max) + vec.store_vec( + old_row_max, + row_max, + row_sum, + final_stats=True, + ) + vec.commit() + elif tmem_sp.uses_query_paired_causal_tail_mask: + sp.wait() + old_row_max, row_max = sp.masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + vec.store_vec(old_row_max, row_max, row_sum) + vec.commit() + p_chunk = sp.masked_exp2_p( + row_max=row_max, + ) + sp.release() + row_sum = sp.softmax_post_release_reduce( + old_row_max, row_max, row_sum, p_chunk + ) + if tmem_sp.uses_query_paired_invalid_tail: + sp.wait() + old_row_max, row_max = sp.invalid_row_max(row_max) + vec.acquire() + vec.store_vec(old_row_max, row_max, row_sum) + vec.commit() + sp.invalid_exp2_p(row_max=row_max) + sp.release() + sp.wait() + sp.release() + old_row_max = sp.softmax_post_release_identity(row_max) + vec.acquire() + vec.store_vec( + old_row_max, + row_max, + row_sum, + final_stats=True, + ) + vec.commit() + elif tmem_sp.cfg.is_causal: + sp.wait() + old_row_max, row_max = sp.masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + vec.store_vec(old_row_max, row_max, row_sum) + vec.commit() + p_chunk = sp.masked_exp2_p( + row_max=row_max, + ) + sp.release() + row_sum = sp.softmax_post_release_reduce( + old_row_max, row_max, row_sum, p_chunk + ) + sp.wait() + sp.release() + old_row_max = sp.softmax_post_release_identity(row_max) + vec.acquire() + vec.store_vec( + old_row_max, + row_max, + row_sum, + final_stats=True, + ) + vec.commit() + else: + sp.wait() + sp.release() + old_row_max = sp.softmax_post_release_identity(row_max) + vec.store_vec(old_row_max, row_max, row_sum) + vec.commit() + + captured_schedule = _schedule_with_work_queue( + softmax_schedule, tmem_sp, tmem_vec, work_queue=work_queue + ) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=index * 4, + num_warps=4, + schedule=captured_schedule, + num_registers=tmem_sp.cfg.num_regs_softmax, + name=f"Softmax{index}Task", + **task_kwargs, + ) + + # Paired QKV instances use separate SP resources. S0S1SequenceResource + # orders their P stores, so this path does not need the TMEM P handoff. + if s0s1_seq is not None and index == 1: + src = _src_resources(tmem_sp, s0s1_seq, work_queue=work_queue) + else: + src = _src_resources(tmem_sp, work_queue=work_queue) + dst = [tmem_vec] + if s0s1_seq is not None and index == 0: + dst.append(s0s1_seq) + + @schedule + def softmax_schedule( + sp: TmemSPResource, + vec: TmemStatsResource, + seq: S0S1SequenceResource, + wq: WorkQueue | None = None, + ) -> None: + """Captured schedule for one softmax warp group.""" + if tmem_sp.enable_early_tile_sum: + # The contribution is produced and consumed inside each iteration; + # do not carry even the scalar tile sum through the persistent loop. + sp.init_softmax_state_early() + else: + p_chunk = sp.init_softmax_state() + scale_softmax_log2 = sp.load_scale_softmax_log2() + vec.init_store_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + # Recompute per-tile SP/Vec TMEM state. + old_row_max, row_max, row_sum, q_offset = sp.init_softmax_work_tile_state() + vec.init_store_work_tile_state() + if tmem_sp.uses_varlen_q_offset_cache: + q_offset = sp.cache_q_offset() + if tmem_sp.uses_packed_dense_k_mask: + seqlen_k = sp.cache_seqlen_k() + window_start = Int32(0) + window_end = Int32(0) + if tmem_sp.uses_variable_window: + window_start, window_end = sp.cache_variable_window_bounds() + # Reserve a stats slot before the first softmax result is published. + vec.acquire() + with domain_loop(loop_start, loop_end, loop_step): + sp.wait() + # Compute row max and publish vec. + if tmem_sp.uses_variable_window: + old_row_max, row_max = sp.variable_window_row_max( + row_max=row_max, + window_start=window_start, + window_end=window_end, + ) + elif tmem_sp.uses_left_window_loop_mask: + old_row_max, row_max = sp.left_masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + elif tmem_sp.uses_varlen_loop_right_mask: + old_row_max, row_max = sp.right_masked_row_max( + row_max=row_max, + q_offset=q_offset, + section=FmhaStage.Loop, + ) + elif tmem_sp.uses_query_paired_q_offset_loop_mask: + old_row_max, row_max = sp.loop_masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + elif tmem_sp.uses_fixed_dense_k_tail_mask: + old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max( + row_max=row_max, + ) + elif tmem_sp.uses_packed_dense_k_mask: + old_row_max, row_max = sp.packed_dense_k_masked_row_max( + row_max=row_max, + seqlen_k=seqlen_k, + section=FmhaStage.Loop, + ) + else: + old_row_max, row_max = sp.compute_row_max(row_max=row_max) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + if s0s1_seq is None: + pass + elif index == 0: + # Softmax0 is the S0-S1 producer: acquire/commit sequence. + seq.acquire() + else: + # Softmax1 is the S0-S1 consumer: wait/release sequence. + seq.wait() + # Apply softmax and write P. + p_chunk = sp.exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + if s0s1_seq is None: + pass + elif index == 0: + seq.commit() + else: + seq.release() + sp.release() + # Reduction. + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + # Acquire vec for next iter. + vec.acquire() + + if tmem_sp.uses_head_paired_causal_tail_mask: + # Head-paired maps Q0/Q1 to adjacent Hq slices at the same S + # tile. Its tail mask uses right_masked_row_max(), which does + # not add the query-paired q_half * q_tile_m sequence advance. + sp.wait() + old_row_max, row_max = sp.right_masked_row_max( + row_max=row_max, + q_offset=q_offset, + section=FmhaStage.Tail, + ) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + if s0s1_seq is None: + pass + elif index == 0: + seq.acquire() + else: + seq.wait() + p_chunk = sp.exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + if s0s1_seq is None: + pass + elif index == 0: + seq.commit() + else: + seq.release() + sp.release() + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + vec.acquire() + sp.wait() + sp.release() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + elif tmem_sp.uses_query_paired_causal_tail_mask: + # Query-paired maps Q1 to the next S tile. Its generic causal + # tail uses masked_row_max(), which includes q_half * q_tile_m + # so each peer tile is masked at the right sequence boundary. + sp.wait() + old_row_max, row_max = sp.masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + if s0s1_seq is not None: + seq.acquire() + p_chunk = sp.masked_exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + if s0s1_seq is not None: + seq.commit() + sp.release() + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + if tmem_sp.uses_query_paired_invalid_tail: + sp.wait() + old_row_max, row_max = sp.invalid_row_max(row_max=row_max) + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + if s0s1_seq is not None: + seq.acquire() + sp.invalid_exp2_p(row_max=row_max) + if s0s1_seq is not None: + seq.commit() + sp.release() + sp.wait() + sp.release() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + elif tmem_sp.cfg.is_causal: + # Causal softmax1 TAIL handles masked rows and cleanup. + sp.wait() + old_row_max, row_max = sp.masked_row_max( + row_max=row_max, + q_offset=q_offset, + ) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + ) + vec.commit() + if s0s1_seq is not None: + seq.wait() + p_chunk = sp.masked_exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + if s0s1_seq is not None: + seq.release() + sp.release() + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + sp.wait() + sp.release() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + else: + # Non-causal TAIL commits the reserved stats slot and lets MMA + # complete its cleanup path. + sp.wait() + sp.release() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + + captured_schedule = _schedule_with_work_queue( + softmax_schedule, tmem_sp, tmem_vec, s0s1_seq, work_queue=work_queue + ) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=index * 4, + num_warps=4, + schedule=captured_schedule, + num_registers=tmem_sp.cfg.num_regs_softmax, + name=f"Softmax{index}Task", + **task_kwargs, + ) + + +def create_correction_task( + tmem_vec0: TmemStatsResource, + tmem_vec1: TmemStatsResource | None, + tmem_o: TmemOResource, + smem_o_0: SmemOResource, + smem_o_1: SmemOResource | None, + gmem_o_0: GmemOResource, + gmem_o_1: GmemOResource | None, + tmem_vec_done_0: TmemStatsDoneResource, + tmem_vec_done_1: TmemStatsDoneResource | None, + work_queue: WorkQueue | None, + task_class: type[Task] = Task, + **task_kwargs: Any, +) -> Task: + """Create the four-warp Correction task (warps 8-11).""" + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + skip_work_tile_if = _packed_context_skip_predicate(work_queue) + + def _create_single_instance_task() -> Task: + fuse_epilogue = smem_o_0.cfg.fuse_epilogue_into_correction + num_o_head_dim_stages = smem_o_0.cfg.num_o_head_dim_stages + + src = _src_resources( + tmem_vec0, + tmem_o, + *([] if tmem_vec0.cfg.stats_via_smem else [tmem_vec_done_0]), + *([smem_o_0] if fuse_epilogue else []), + work_queue=work_queue, + ) + + @schedule + def correction_schedule( + v0: TmemStatsResource, + to: TmemOResource, + so0: SmemOResource, + go0: GmemOResource, + vd0: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + v0.init_read_state() + scale_softmax_log2 = v0.load_scale_softmax_log2() + output_scale = v0.load_output_scale() + to.init_correction_state() + so0.init_store_state() + if fuse_epilogue: + go0.init_store_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + v0.init_read_work_tile_state() + to.init_correction_work_tile_state() + so0.init_store_work_tile_state() + + v0.wait() + v0.release() + if not tmem_vec0.cfg.stats_via_smem: + vd0.wait() + vd0.release() + with domain_loop(loop_start, loop_end, loop_step): + v0.wait() + vec_old_max, vec_new_max, _, vec_scale = v0.read_vec( + scale_softmax_log2=scale_softmax_log2, + ) + if not tmem_vec0.cfg.stats_via_smem: + vd0.wait() + vd0.release() + to.wait() + to.correct( + vec_old_max=vec_old_max, + vec_new_max=vec_new_max, + vec_scale=vec_scale, + inst_idx=0, + ) + v0.release() + to.release() + + v0.wait() + _, _, vec_row_sum, vec_scale = v0.read_vec( + scale_softmax_log2=scale_softmax_log2, + final_stats=True, + ) + if not tmem_vec0.cfg.stats_via_smem: + vd0.wait() + vd0.release() + v0.release() + to.wait() + for head_dim_stage_idx in range(num_o_head_dim_stages): + so0.acquire() + so0.store_o( + vec_row_sum=vec_row_sum, + vec_scale=vec_scale, + output_scale=output_scale, + head_dim_stage_idx=head_dim_stage_idx, + ) + so0.commit() + if fuse_epilogue: + # The same four-warp group consumes the completed SMEM + # stage. Only its first warp issues TMA; all four wait + # and release the pipeline stage together. + so0.wait() + head_coord, batch_coord, seq_coord_q = ( + so0.compute_output_coords() + ) + go0.tma_store( + head_coord=head_coord, + batch_coord=batch_coord, + seq_coord_q=seq_coord_q, + head_dim_stage_idx=head_dim_stage_idx, + correction_fused=True, + ) + so0.release() + to.release() + if tmem_vec0.cfg.stats_via_smem: + # Consume the cursor-balancing record emitted by Softmax. + v0.wait() + v0.release() + + captured_schedule = _schedule_with_work_queue( + correction_schedule, + tmem_vec0, + tmem_o, + smem_o_0, + gmem_o_0, + tmem_vec_done_0, + work_queue=work_queue, + ) + dst = [smem_o_0] + if fuse_epilogue: + dst.append(gmem_o_0) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=smem_o_0.cfg.correction_warp_ids[0], + num_warps=4, + schedule=captured_schedule, + num_registers=smem_o_0.cfg.num_regs_correction, + name="CorrectionTask", + **task_kwargs, + ) + + def _create_paired_task() -> Task: + if tmem_vec1 is None or smem_o_1 is None or tmem_vec_done_1 is None: + raise ValueError("paired correction scheduling requires peer-1 resources") + + src = _src_resources( + tmem_vec0, + tmem_vec1, + tmem_o, + *( + [] + if tmem_vec0.cfg.stats_via_smem + else [tmem_vec_done_0, tmem_vec_done_1] + ), + work_queue=work_queue, + ) + + @schedule + def correction_schedule( + v0: TmemStatsResource, + v1: TmemStatsResource, + to: TmemOResource, + so0: SmemOResource, + so1: SmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Captured schedule for O rescale and SMEM staging.""" + v0.init_read_state() + v1.init_read_state() + scale_softmax_log2_v0 = v0.load_scale_softmax_log2() + scale_softmax_log2_v1 = v1.load_scale_softmax_log2() + output_scale0 = v0.load_output_scale() + output_scale1 = v1.load_output_scale() + to.init_correction_state() + so0.init_store_state() + so1.init_store_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + # Per-tile TMEM/SMEM cached addresses are computed here. + v0.init_read_work_tile_state() + v1.init_read_work_tile_state() + to.init_correction_work_tile_state() + so0.init_store_work_tile_state() + so1.init_store_work_tile_state() + # The empty stats pipeline needs no priming. Discard its first + # slot and retain TmemStats1 for the first loop cross-release. + v0.wait() + v0.release() + v1.wait() + # The correction loop consumes vec/O pairs in alternating order so + # each half can unblock the other half's next producer. + with domain_loop(loop_start, loop_end, loop_step): + # Part 1: consume TmemStats0 + O0, release TmemStats1. + v0.wait() + vec_old_max, vec_new_max, _, vec_scale = v0.read_vec( + scale_softmax_log2=scale_softmax_log2_v0, + ) + to.wait() + to.correct( + vec_old_max=vec_old_max, + vec_new_max=vec_new_max, + vec_scale=vec_scale, + inst_idx=0, + ) + v1.release() + to.release() + # Part 2: consume TmemStats1 + O1, release TmemStats0. + v1.wait() + vec_old_max, vec_new_max, _, vec_scale = v1.read_vec( + scale_softmax_log2=scale_softmax_log2_v1, + ) + to.wait() + to.correct( + vec_old_max=vec_old_max, + vec_new_max=vec_new_max, + vec_scale=vec_scale, + inst_idx=1, + ) + v0.release() + to.release() + # TAIL: consume remaining stats, release tmem-stats-done gates, and + # stage corrected O0/O1 into SMEM for the epilogue task. + v1.release() + v0.wait() + _, _, vec_row_sum, vec_scale = v0.read_vec( + scale_softmax_log2=scale_softmax_log2_v0, + final_stats=True, + ) + if not tmem_vec0.cfg.stats_via_smem: + vd0.wait() + vd0.release() + v0.release() + to.wait() + so0.acquire() + so0.store_o( + vec_row_sum=vec_row_sum, + vec_scale=vec_scale, + output_scale=output_scale0, + ) + so0.commit() + to.release() + v1.wait() + _, _, vec_row_sum, vec_scale = v1.read_vec( + scale_softmax_log2=scale_softmax_log2_v1, + final_stats=True, + ) + if not tmem_vec0.cfg.stats_via_smem: + vd1.wait() + vd1.release() + v1.release() + to.wait() + so1.acquire() + so1.store_o( + vec_row_sum=vec_row_sum, + vec_scale=vec_scale, + output_scale=output_scale1, + ) + so1.commit() + to.release() + + captured_schedule = _schedule_with_work_queue( + correction_schedule, + tmem_vec0, + tmem_vec1, + tmem_o, + smem_o_0, + smem_o_1, + tmem_vec_done_0, + tmem_vec_done_1, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=[smem_o_0, smem_o_1], + warp_idx=8, + num_warps=4, + schedule=captured_schedule, + num_registers=smem_o_0.cfg.num_regs_correction, + name="CorrectionTask", + **task_kwargs, + ) + + if smem_o_0.cfg.single_qkv_instance: + return _create_single_instance_task() + return _create_paired_task() + + +def create_epilogue_task( + smem_o_0: SmemOResource, + smem_o_1: SmemOResource | None, + gmem_o_0: GmemOResource, + gmem_o_1: GmemOResource | None, + work_queue: WorkQueue | None, + task_class: type[Task] = Task, + **task_kwargs: Any, +) -> Task: + """Create the one-warp Epilogue store task (warp 14).""" + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + skip_work_tile_if = _packed_context_skip_predicate(work_queue) + + def _create_single_instance_task() -> Task: + src = _src_resources(smem_o_0, work_queue=work_queue) + num_o_head_dim_stages = smem_o_0.cfg.num_o_head_dim_stages + + @schedule + def epilogue_schedule( + so0: SmemOResource, + go0: GmemOResource, + wq: WorkQueue | None = None, + ) -> None: + so0.init_output_state() + go0.init_store_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + so0.init_output_work_tile_state() + with domain_loop(loop_start, loop_end, loop_step): + pass + for head_dim_stage_idx in range(num_o_head_dim_stages): + so0.wait() + head_coord, batch_coord, seq_coord_q = so0.compute_output_coords() + go0.tma_store( + head_coord=head_coord, + batch_coord=batch_coord, + seq_coord_q=seq_coord_q, + head_dim_stage_idx=head_dim_stage_idx, + ) + so0.release() + + captured_schedule = _schedule_with_work_queue( + epilogue_schedule, + smem_o_0, + gmem_o_0, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=[gmem_o_0], + warp_idx=gmem_o_0.cfg.epilogue_warp_id, + num_warps=1, + schedule=captured_schedule, + num_registers=gmem_o_0.cfg.num_regs_other, + name="EpilogueTask", + **task_kwargs, + ) + + def _create_paired_task() -> Task: + if smem_o_1 is None or gmem_o_1 is None: + raise ValueError("paired epilogue scheduling requires peer-1 resources") + + src = _src_resources(smem_o_0, smem_o_1, work_queue=work_queue) + + @schedule + def epilogue_schedule( + so0: SmemOResource, + so1: SmemOResource, + go0: GmemOResource, + go1: GmemOResource, + wq: WorkQueue | None = None, + ) -> None: + """Captured schedule for GMEM O stores.""" + so0.init_output_state() + so1.init_output_state() + go0.init_store_state() + go1.init_store_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + # Per-tile SMEM O address base is computed in work vars. + so0.init_output_work_tile_state() + so1.init_output_work_tile_state() + with domain_loop(loop_start, loop_end, loop_step): + pass + # Store the first corrected O tile through gmem_o_0. + so0.wait() + head_coord, batch_coord, seq_coord_q = so0.compute_output_coords() + go0.tma_store( + head_coord=head_coord, + batch_coord=batch_coord, + seq_coord_q=seq_coord_q, + ) + so0.release() + # Store the second corrected O tile through gmem_o_1. + so1.wait() + head_coord, batch_coord, seq_coord_q = so1.compute_output_coords() + go1.tma_store( + head_coord=head_coord, + batch_coord=batch_coord, + seq_coord_q=seq_coord_q, + ) + so1.release() + + captured_schedule = _schedule_with_work_queue( + epilogue_schedule, + smem_o_0, + smem_o_1, + gmem_o_0, + gmem_o_1, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=[gmem_o_0, gmem_o_1], + warp_idx=14, + num_warps=1, + schedule=captured_schedule, + num_registers=gmem_o_0.cfg.num_regs_other, + name="EpilogueTask", + **task_kwargs, + ) + + if smem_o_0.cfg.single_qkv_instance: + return _create_single_instance_task() + return _create_paired_task() + + +def create_padding_task( + work_queue: WorkQueue | None, + warp_idx: int = 15, + num_warps: int = 1, + num_registers: int = 32, + name: str = "PaddingTask", + task_class: type[Task] = Task, + **task_kwargs: Any, +) -> Task: + """Create the one-warp padding task (warp 15 in the D128 schedule). + + Required in ALL modes (persistent and non-persistent) because + ``setmaxnreg.sync`` requires every warp in the warp group to + participate. In D128, warps 12-15 form warp group 3; without the padding + task its final warp never calls ``setmaxregister``, deadlocking the group. + + In persistent mode the task also consumes work_queue tiles so that + the auxiliary warp participates in the persistent outer loop. + + In CLC dynamic mode, the padding task is replaced by a scheduler task + (see ``create_scheduler_task``). + """ + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + skip_work_tile_if = _packed_context_skip_predicate(work_queue) + src = _src_resources(work_queue=work_queue) + + @schedule + def padding_schedule(wq: WorkQueue | None = None) -> None: + """Captured schedule for warp-group register participation.""" + with ( + _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if), + domain_loop(loop_start, loop_end, loop_step), + ): + pass + + captured_schedule = _schedule_with_work_queue( + padding_schedule, work_queue=work_queue + ) + return task_class( + src_resources=src, + dst_resources=[], + warp_idx=warp_idx, + num_warps=num_warps, + schedule=captured_schedule, + num_registers=num_registers, + name=name, + **task_kwargs, + ) + + +def _prefetch_page_offsets_for_work_tile( + spo: SmemPageOffsetsKvResource, + *, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + loop_start: int, + loop_end: int, + loop_step: int, + staged_single_instance: bool, +) -> None: + """Produce page-ID stages in the same logical order as the load task.""" + if staged_single_instance: + # D>128 overlaps QK(i) with PV(i-1): K0, then K_i/V_{i-1}, then V_last. + # One page-ID stage is shared by all head-dimension slices of a logical + # K or V tile, so this producer fires once per tile rather than per slice. + spo.acquire() + spo.load_k( + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spo.commit() + with domain_loop(loop_start + 1, loop_end, loop_step): + spo.acquire() + spo.load_k( + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spo.commit() + spo.acquire() + spo.load_v( + previous=True, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spo.commit() + spo.acquire() + spo.load_v( + previous=False, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spo.commit() + return + + with domain_loop(loop_start, loop_end, loop_step): + spo.acquire() + spo.load_k( + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spo.commit() + spo.acquire() + spo.load_v( + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spo.commit() + + +def _prefetch_reused_page_windows_for_work_tile( + spok: SmemPageOffsetsKvResource, + spov: SmemPageOffsetsKvResource, + *, + kv_tile_start: Int32, + kv_request_begin: Int32, + kv_page_idx_ub: Int32, + loop_start: object, + loop_end: object, + loop_step: object, + page_window_period: int, +) -> None: + """Publish independent K/V page windows at their structural cadence.""" + if ( + not isinstance(loop_start, int) + or not isinstance(loop_end, int) + or not isinstance(loop_step, int) + or loop_start != 0 + or loop_step != 1 + or loop_end < page_window_period + or loop_end % page_window_period != 0 + ): + raise ValueError( + "reused page windows require a compile-time K/V domain " + "divisible by the topology-derived page-window period" + ) + + spok.acquire() + spok.load_k( + tile_offset=0, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spok.commit() + spov.acquire() + spov.load_v( + tile_offset=0, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spov.commit() + + with domain_loop(page_window_period, loop_end, page_window_period): + spok.acquire() + spok.load_k( + tile_offset=0, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spok.commit() + spov.acquire() + spov.load_v( + tile_offset=0, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + spov.commit() + + +def create_page_offsets_task( + gmem_qkv: GmemQKVResource, + smem_page_offsets_kv: SmemPageOffsetsKvResource, + work_queue: WorkQueue | None, + task_class: type[Task] = Task, + num_registers: int = 32, + smem_page_offsets_v: SmemPageOffsetsKvResource | None = None, + **task_kwargs: Any, +) -> Task: + """Create the one-warp paged-KV page-offsets prefetch task. + + Replaces ``create_padding_task`` when ``cfg.use_paged_kv`` is True. The + configuration's empty/scheduler warp prefetches page-table entries into + SMEM so the load warp can read cached page IDs when issuing paged TMA + copies. This also preserves that warp's ``setmaxnreg.sync`` participation. + + Paired D128 CLC does not instantiate this task: its load warp reads page + IDs directly, leaving warp 15 exclusively responsible for CLC. Staged + D256 can use this task with CLC because page offsets run on the empty warp + while the freed epilogue warp owns scheduling. No task therefore combines + dynamic work-queue and page-offset production through the public DSL API. + """ + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + skip_work_tile_if = _packed_context_skip_predicate(work_queue) + src = _src_resources(gmem_qkv, work_queue=work_queue) + dst = [smem_page_offsets_kv] + if smem_page_offsets_v is not None: + dst.append(smem_page_offsets_v) + staged_single_instance = ( + smem_page_offsets_kv.cfg.single_qkv_instance + and smem_page_offsets_kv.cfg.has_tmem_p_pipeline + ) + page_window_period = smem_page_offsets_kv.cfg.page_table_window_entries // ( + smem_page_offsets_kv.cfg.kv_tile_n + // smem_page_offsets_kv.cfg.num_tokens_per_page + ) + + def page_offsets_schedule_body( + gqkv: GmemQKVResource, + spo: SmemPageOffsetsKvResource, + spov: SmemPageOffsetsKvResource | None, + wq: WorkQueue | None = None, + ) -> None: + """Captured schedule for K/V page-table prefetch.""" + spo.init_load_state() + if spov is not None: + spov.init_load_state() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): + ( + kv_tile_start, + kv_request_begin, + kv_page_idx_ub, + ) = gqkv.compute_page_coords() + if spov is None: + _prefetch_page_offsets_for_work_tile( + spo, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + loop_start=loop_start, + loop_end=loop_end, + loop_step=loop_step, + staged_single_instance=staged_single_instance, + ) + else: + _prefetch_reused_page_windows_for_work_tile( + spo, + spov, + kv_tile_start=kv_tile_start, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + loop_start=loop_start, + loop_end=loop_end, + loop_step=loop_step, + page_window_period=page_window_period, + ) + + @schedule + def page_offsets_schedule( + gqkv: GmemQKVResource, + spo: SmemPageOffsetsKvResource, + wq: WorkQueue | None = None, + ) -> None: + page_offsets_schedule_body(gqkv, spo, None, wq) + + @schedule + def reused_page_windows_schedule( + gqkv: GmemQKVResource, + spok: SmemPageOffsetsKvResource, + spov: SmemPageOffsetsKvResource, + wq: WorkQueue | None = None, + ) -> None: + page_offsets_schedule_body(gqkv, spok, spov, wq) + + if smem_page_offsets_v is None: + captured_schedule = _schedule_with_work_queue( + page_offsets_schedule, + gmem_qkv, + smem_page_offsets_kv, + work_queue=work_queue, + ) + else: + captured_schedule = _schedule_with_work_queue( + reused_page_windows_schedule, + gmem_qkv, + smem_page_offsets_kv, + smem_page_offsets_v, + work_queue=work_queue, + ) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=smem_page_offsets_kv.cfg.empty_warp_id, + num_warps=1, + schedule=captured_schedule, + num_registers=num_registers, + name="PageTableTask", + **task_kwargs, + ) + + +def create_scheduler_task( + work_queue: WorkQueue, + warp_idx: int = 15, + num_registers: int = 32, + task_class: type[Task] = Task, + **task_kwargs: Any, +) -> Task: + """Create the one-warp CLC scheduler task (warp 15 in D128). + + Replaces the padding task in CLC dynamic persistent mode. + Issues CLC tile-fetch queries (producer side) and participates in + the persistent outer loop. Still satisfies the ``setmaxnreg.sync`` + requirement for the final warp group. + """ + loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) + + @schedule + def scheduler_schedule(wq: WorkQueue) -> None: + """Captured schedule for CLC work-tile fetches.""" + with _work_tile_schedule_loop(wq): + with domain_loop(loop_start, loop_end, loop_step): + pass + # Producer side: issue CLC tile-fetch query. + wq.acquire() + wq.fetch_work_tile() + wq.commit() + + captured_schedule = scheduler_schedule(work_queue) + return task_class( + src_resources=[work_queue], + dst_resources=[work_queue], + warp_idx=warp_idx, + num_warps=1, + schedule=captured_schedule, + num_registers=num_registers, + name="SchedulerTask", + **task_kwargs, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py new file mode 100644 index 000000000000..239e381b96de --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Low-level helper intrinsics for FMHA context TS resources.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 + + +@cute.jit +def variable_window_cta_min_start( + cta_starts: cute.Tensor, + *, + batch_coord: Int32, + seq_coord: Int32, + q_stride: int | Int32, + tile_size_q: cutlass.Constexpr[int], +) -> Int32: + """Load the plan-time minimum variable-window start for one Q CTA.""" + num_seq_tiles = cute.ceil_div(q_stride, tile_size_q) + return Int32(cta_starts[batch_coord * num_seq_tiles + seq_coord]) + + +def bottom_right_window_left_bound( + query_idx: int | Int32, + q_offset: int | Int32, + window_size_left: int, +) -> int | Int32: + """Return the inclusive left bound for bottom-right causal attention.""" + return query_idx + q_offset - window_size_left + + +def bottom_right_window_tile_start( + *, + seq_coord: int | Int32, + q_tile_m: int | Int32, + kv_tile_n: int | Int32, + q_offset: int | Int32, + window_size_left: int, +) -> int | Int32: + """Return the first K/V tile intersecting a bottom-right left window.""" + raw_start = ( + bottom_right_window_left_bound( + seq_coord * q_tile_m, + q_offset, + window_size_left, + ) + // kv_tile_n + ) + if isinstance(raw_start, int): + return max(0, raw_start) + return cute.math.max(Int32(0), raw_start) + + +def bottom_right_window_max_tiles( + *, + q_tile_m: int, + kv_tile_n: int, + window_size_left: int, +) -> int: + """Return the offset-independent maximum K/V span for one Q tile. + + The visible interval before sequence-boundary clipping is inclusive and + has ``window_size_left + q_tile_m`` tokens. Its alignment against a K/V + tile can require one additional tile at each boundary, so the maximum + intersecting span is ``ceil((length + kv_tile_n - 1) / kv_tile_n)``. + Packed-ragged scheduling uses this bound because each request can have a + different bottom-right Q/K offset while a task must keep one loop domain. + """ + if q_tile_m <= 0 or kv_tile_n <= 0 or window_size_left < 0: + raise ValueError("tile sizes must be positive and window size non-negative") + numerator = window_size_left + q_tile_m + kv_tile_n - 1 + return (numerator + kv_tile_n - 1) // kv_tile_n + + +@cute.jit +def freeze_smem_descriptor(desc): + """Copy a shared-memory descriptor through a register to prevent rematerialization.""" + return cute.arch.inline_ptx( + "mov.b64 {$w0}, {$r0};", + write_only_types=[Int64], + read_only_args=[desc], + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py new file mode 100644 index 000000000000..195d7b50c9ee --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Paged-KV helpers for the task-scheduled FMHA *context* kernel. + +Context has no multi-CTA-KV split, folds any sliding-window prefix into +``kv_tile_start``, and uses one ``kv_tile_start + loop_offset`` expression for +the current K/V tile index. Page IDs come from fixed row-strided block tables. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + +from cutlass.experimental.task_scheduling.resources import StageInfo + +from .fmha_resources import FmhaConfig + + +@cute.jit +def _load_runtime_seq_len_kv( + seq_lens_kv: cute.Pointer | None, + max_seq_len_kv: Int32 | int, + batch_coord: Int32, +) -> Int32: + # Context resolves (head, batch, seq) directly from the work tile, so the + # batch coord is always known by the caller — no persistent-vs-launch + # fallback indirection required. + if cutlass.const_expr(seq_lens_kv is None): + return Int32(max_seq_len_kv) + return Int32(seq_lens_kv[batch_coord]) + + +@cute.jit +def _runtime_last_valid_page_idx(cfg: FmhaConfig, seq_len_kv: Int32) -> Int32: + num_pages = cute.ceil_div(seq_len_kv, cfg.num_tokens_per_page) + return cute.math.max(num_pages - Int32(1), Int32(0)) + + +@cute.jit +def _load_block_table_row_bounds( + block_table_row_stride: Int32, + cfg: FmhaConfig, + seq_len_kv: Int32, + batch_coord: Int32, +) -> tuple[Int32, Int32]: + """Return one fixed-table row base and runtime-valid inclusive page bound.""" + + row_begin = batch_coord * block_table_row_stride + return row_begin, _runtime_last_valid_page_idx(cfg, seq_len_kv) + + +@cute.jit +def _resolve_kv_tile_idx_context( + stage_info: StageInfo, + kv_tile_start: Int32, + tile_offset: cutlass.Constexpr[int] = 0, +) -> Int32: + # Context's current K/V tile index is kv_tile_start + loop_offset, with an + # optional compile-time shift for the staged D>128 schedule's previous V. + # No multi-CTA-KV transform; sliding-window prefix skipping is already + # folded into kv_tile_start by GmemQKVResource. + return kv_tile_start + Int32(stage_info.loop_offset) + Int32(tile_offset) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/__init__.py new file mode 100644 index 000000000000..6bf2639138df --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FMHA TS decode resources, schedules, kernel, and runner.""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py new file mode 100644 index 000000000000..de53ccb9c7de --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py @@ -0,0 +1,580 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Validate paged request metadata and canonical BSR for one-shot planning. + +A BSR Q-block row is one ``block_indptr`` row keyed by +``(batch, kv_head, q_block)``. The maximum row width gives the one-shot API the +same semantic ``max_blocks_per_row`` bound that reusable plans receive from +their caller. Token-mask contents belong to the run-time prepare kernel and +are not read here. + +Paged inspection first validates live sequence lengths and page rows, then +four warps validate four BSR Q-block rows per CTA. Both publish one validation +status plus the maximum row width in one Int64 summary; no route payload is +constructed. +""" + +import functools +from collections.abc import Callable + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings import driver as cuda_drv + +from .fmha_decode_resources.helpers_common import _warp_broadcast_i32 + +_WARPS_PER_CTA = 4 +_WARP_SIZE = 32 +_THREADS_PER_CTA = _WARPS_PER_CTA * _WARP_SIZE +_COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" + +# ``summary`` is a zero-initialized Int64[2] shared with the host wrapper. +_SUMMARY_ERROR_CODE = 0 +_SUMMARY_MAX_BSR_BLOCK_COUNT = 1 +_SUMMARY_FIELDS = 2 + +_BSR_ERROR_NONE = 0 +_BSR_ERROR_NOT_STRICTLY_INCREASING = 1 +_BSR_ERROR_INDEX_OUT_OF_RANGE = 2 +_BSR_ERROR_INVALID_INDPTR = 3 +_ERROR_INVALID_SEQ_LEN = 4 +_ERROR_INVALID_PAGE_INDPTR = 5 +_ERROR_INSUFFICIENT_PAGE_CAPACITY = 6 +_ERROR_INVALID_PHYSICAL_PAGE_ID = 7 + + +@cute.jit +def _validate_bsr_row_lane( + block_indices: cute.Tensor, + bsr_row_begin: cutlass.Int32, + bsr_row_end: cutlass.Int32, + lane_idx: cutlass.Int32, + num_kv_blocks: cutlass.Int32, +) -> cutlass.Int32: + """Validate one lane stripe of a canonical ordered BSR row.""" + + error_code = cutlass.Int32(_BSR_ERROR_NONE) + selected_kv_block_count = cutlass.Int64(bsr_row_end) - cutlass.Int64(bsr_row_begin) + bsr_entry_offset = cutlass.Int64(lane_idx) + while bsr_entry_offset < selected_kv_block_count: + entry_position = cutlass.Int64(bsr_row_begin) + bsr_entry_offset + block_id = cutlass.Int32(block_indices[entry_position]) + in_range = block_id >= 0 and block_id < num_kv_blocks + if not in_range: + error_code = cutlass.Int32(_BSR_ERROR_INDEX_OUT_OF_RANGE) + else: + if entry_position > cutlass.Int64(bsr_row_begin): + previous_block_id = cutlass.Int32( + block_indices[entry_position - cutlass.Int64(1)] + ) + if ( + block_id <= previous_block_id + and error_code < _BSR_ERROR_NOT_STRICTLY_INCREASING + ): + error_code = cutlass.Int32(_BSR_ERROR_NOT_STRICTLY_INCREASING) + bsr_entry_offset += cutlass.Int64(_WARP_SIZE) + return error_code + + +@cute.jit +def _inspect_bsr_row( + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + batch_idx: cutlass.Int32, + kv_head_idx: cutlass.Int32, + q_block_row_idx: cutlass.Int32, + lane_idx: cutlass.Int32, + bsr_row_is_valid: cutlass.Boolean, + num_kv_blocks: cutlass.Int32, +) -> tuple[cutlass.Int32, cutlass.Int32]: + """Validate one BSR row against a static or live Int32 upper bound.""" + + bsr_row_begin = cutlass.Int32(0) + bsr_row_end = cutlass.Int32(0) + row_range_is_valid = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0) and bsr_row_is_valid: + bsr_row_begin = cutlass.Int32( + block_indptr[batch_idx, kv_head_idx, q_block_row_idx] + ) + bsr_row_end = cutlass.Int32( + block_indptr[batch_idx, kv_head_idx, q_block_row_idx + 1] + ) + row_range_is_valid = cutlass.Int32( + bsr_row_begin >= cutlass.Int32(0) + and bsr_row_begin <= bsr_row_end + and bsr_row_end <= cutlass.Int32(cute.size(block_indices)) + ) + bsr_row_begin = _warp_broadcast_i32(bsr_row_begin, 0) + bsr_row_end = _warp_broadcast_i32(bsr_row_end, 0) + row_range_is_valid = _warp_broadcast_i32(row_range_is_valid, 0) + + error_code = cutlass.Int32(_BSR_ERROR_NONE) + if bsr_row_is_valid: + if row_range_is_valid == cutlass.Int32(0): + error_code = cutlass.Int32(_BSR_ERROR_INVALID_INDPTR) + else: + error_code = _validate_bsr_row_lane( + block_indices, + bsr_row_begin, + bsr_row_end, + lane_idx, + num_kv_blocks, + ) + error_code = cutlass.Int32(cute.arch.warp_redux_sync(error_code, "max")) + return error_code, bsr_row_end - bsr_row_begin + + +@cute.jit +def _inspect_live_paged_bsr_row( + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + seq_lens_kv: cute.Tensor, + batch_idx: cutlass.Int32, + kv_head_idx: cutlass.Int32, + q_block_row_idx: cutlass.Int32, + lane_idx: cutlass.Int32, + bsr_row_is_valid: cutlass.Boolean, + kv_block_size: cutlass.Constexpr[int], +) -> tuple[cutlass.Int32, cutlass.Int32]: + """Inspect one live BSR row against the current request sequence length.""" + + error_code = cutlass.Int32(_BSR_ERROR_NONE) + selected_kv_block_count = cutlass.Int32(0) + num_live_kv_blocks = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0) and bsr_row_is_valid: + live_seq_len_kv = cutlass.Int32(seq_lens_kv[batch_idx]) + num_live_kv_blocks = (live_seq_len_kv - cutlass.Int32(1)) // cutlass.Int32( + kv_block_size + ) + cutlass.Int32(1) + num_live_kv_blocks = _warp_broadcast_i32(num_live_kv_blocks, 0) + error_code, selected_kv_block_count = _inspect_bsr_row( + block_indptr, + block_indices, + batch_idx, + kv_head_idx, + q_block_row_idx, + lane_idx, + bsr_row_is_valid, + num_live_kv_blocks, + ) + return error_code, selected_kv_block_count + + +class _InspectBlockSparseBsr: + """Validate canonical BSR against static or live request metadata.""" + + def __init__( + self, + *, + batch_size: int, + num_kv_heads: int, + seq_len_q: int, + seq_len_kv: int | None, + q_block_size: int, + kv_block_size: int, + ) -> None: + self.num_kv_heads = num_kv_heads + self.num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + self.num_kv_blocks = ( + 0 + if seq_len_kv is None + else (seq_len_kv + kv_block_size - 1) // kv_block_size + ) + self.kv_block_size = kv_block_size + self.total_bsr_row_count = batch_size * num_kv_heads * self.num_q_block_rows + + @cute.jit + def _launch( + self, + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + seq_lens_kv: cute.Tensor | None, + summary: cute.Tensor, + stream: cuda_drv.CUstream, + ) -> None: + self.kernel( + block_indptr, + block_indices, + seq_lens_kv, + summary, + ).launch( + grid=[ + (self.total_bsr_row_count + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, + 1, + 1, + ], + block=[_THREADS_PER_CTA, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + seq_lens_kv: cute.Tensor | None, + summary: cute.Tensor, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = thread_idx // _WARP_SIZE + lane_idx = thread_idx % _WARP_SIZE + + linear_bsr_row_idx = block_idx * _WARPS_PER_CTA + warp_idx + bsr_row_is_valid = linear_bsr_row_idx < self.total_bsr_row_count + safe_linear_bsr_row_idx = ( + linear_bsr_row_idx if bsr_row_is_valid else cutlass.Int32(0) + ) + q_block_row_idx = safe_linear_bsr_row_idx % self.num_q_block_rows + linear_batch_head_idx = safe_linear_bsr_row_idx // self.num_q_block_rows + kv_head_idx = linear_batch_head_idx % self.num_kv_heads + batch_idx = linear_batch_head_idx // self.num_kv_heads + + if cutlass.const_expr(seq_lens_kv is None): + error_code, selected_kv_block_count = _inspect_bsr_row( + block_indptr, + block_indices, + batch_idx, + kv_head_idx, + q_block_row_idx, + lane_idx, + bsr_row_is_valid, + cutlass.Int32(self.num_kv_blocks), + ) + else: + error_code, selected_kv_block_count = _inspect_live_paged_bsr_row( + block_indptr, + block_indices, + seq_lens_kv, + batch_idx, + kv_head_idx, + q_block_row_idx, + lane_idx, + bsr_row_is_valid, + self.kv_block_size, + ) + + if lane_idx == cutlass.Int32(0) and bsr_row_is_valid: + if error_code != cutlass.Int32(_BSR_ERROR_NONE): + cute.arch.atomic_max( + summary.iterator + _SUMMARY_ERROR_CODE, + cutlass.Int64(error_code), + sem="relaxed", + scope="gpu", + ) + else: + cute.arch.atomic_max( + summary.iterator + _SUMMARY_MAX_BSR_BLOCK_COUNT, + cutlass.Int64(selected_kv_block_count), + sem="relaxed", + scope="gpu", + ) + + @cute.jit + def __call__( + self, + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + seq_lens_kv: cute.Tensor | None, + summary: cute.Tensor, + stream: cuda_drv.CUstream, + ) -> None: + self._launch(block_indptr, block_indices, seq_lens_kv, summary, stream) + + +class _InspectPagedKvMetadata: + """Validate live lengths and page rows with one warp per request.""" + + def __init__( + self, + *, + batch_size: int, + minimum_seq_len_kv: int, + max_seq_len_kv: int, + page_size: int, + ) -> None: + self.batch_size = batch_size + self.minimum_seq_len_kv = minimum_seq_len_kv + self.max_seq_len_kv = max_seq_len_kv + self.page_size = page_size + + @cute.jit + def __call__( + self, + paged_kv_indptr: cute.Tensor, + paged_kv_indices: cute.Tensor, + seq_lens_kv: cute.Tensor, + num_physical_kv_pages: cutlass.Int64, + summary: cute.Tensor, + stream: cuda_drv.CUstream, + ) -> None: + self.kernel( + paged_kv_indptr, + paged_kv_indices, + seq_lens_kv, + num_physical_kv_pages, + summary, + ).launch( + grid=[ + (self.batch_size + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, + 1, + 1, + ], + block=[_THREADS_PER_CTA, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + paged_kv_indptr: cute.Tensor, + paged_kv_indices: cute.Tensor, + seq_lens_kv: cute.Tensor, + num_physical_kv_pages: cutlass.Int64, + summary: cute.Tensor, + ) -> None: + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = thread_idx // _WARP_SIZE + lane_idx = thread_idx % _WARP_SIZE + batch_idx = block_idx * _WARPS_PER_CTA + warp_idx + request_is_valid = batch_idx < self.batch_size + + request_begin = cutlass.Int32(0) + request_end = cutlass.Int32(0) + request_range_is_valid = cutlass.Int32(0) + error_code = cutlass.Int32(_BSR_ERROR_NONE) + if lane_idx == 0 and request_is_valid: + seq_len_kv = cutlass.Int32(seq_lens_kv[batch_idx]) + seq_len_is_valid = cutlass.Boolean( + seq_len_kv >= cutlass.Int32(self.minimum_seq_len_kv) + and seq_len_kv <= cutlass.Int32(self.max_seq_len_kv) + ) + if not seq_len_is_valid: + error_code = cutlass.Int32(_ERROR_INVALID_SEQ_LEN) + + request_begin = cutlass.Int32(paged_kv_indptr[batch_idx]) + request_end = cutlass.Int32(paged_kv_indptr[batch_idx + 1]) + num_page_indices = cutlass.Int32(cute.size(paged_kv_indices)) + request_range_is_valid = cutlass.Int32( + paged_kv_indptr[cutlass.Int32(0)] == cutlass.Int32(0) + and request_begin >= cutlass.Int32(0) + and request_begin <= request_end + and request_end <= num_page_indices + ) + if request_range_is_valid == cutlass.Int32(0): + error_code = cutlass.Int32(_ERROR_INVALID_PAGE_INDPTR) + elif seq_len_is_valid: + required_pages = (seq_len_kv - cutlass.Int32(1)) // cutlass.Int32( + self.page_size + ) + cutlass.Int32(1) + if request_end - request_begin < required_pages: + error_code = cutlass.Int32(_ERROR_INSUFFICIENT_PAGE_CAPACITY) + + request_begin = _warp_broadcast_i32(request_begin, 0) + request_end = _warp_broadcast_i32(request_end, 0) + request_range_is_valid = _warp_broadcast_i32(request_range_is_valid, 0) + if request_is_valid and request_range_is_valid != cutlass.Int32(0): + page_offset = cutlass.Int64(lane_idx) + request_page_count = cutlass.Int64(request_end) - cutlass.Int64( + request_begin + ) + while page_offset < request_page_count: + page_position = cutlass.Int64(request_begin) + page_offset + physical_page_id = cutlass.Int32(paged_kv_indices[page_position]) + if ( + physical_page_id < cutlass.Int32(0) + or cutlass.Int64(physical_page_id) >= num_physical_kv_pages + ): + error_code = cutlass.Int32(_ERROR_INVALID_PHYSICAL_PAGE_ID) + page_offset += cutlass.Int64(_WARP_SIZE) + + error_code = cutlass.Int32(cute.arch.warp_redux_sync(error_code, "max")) + if ( + lane_idx == cutlass.Int32(0) + and request_is_valid + and error_code != cutlass.Int32(_BSR_ERROR_NONE) + ): + cute.arch.atomic_max( + summary.iterator + _SUMMARY_ERROR_CODE, + cutlass.Int64(error_code), + sem="relaxed", + scope="gpu", + ) + + +class _InspectPagedBlockSparseMetadata: + """Launch request inspection before live-BSR inspection into one summary.""" + + def __init__( + self, + *, + inspect_requests: Callable[..., None], + inspect_bsr: Callable[..., None], + ) -> None: + self.inspect_requests = inspect_requests + self.inspect_bsr = inspect_bsr + + @cute.jit + def __call__( + self, + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + paged_kv_indptr: cute.Tensor, + paged_kv_indices: cute.Tensor, + seq_lens_kv: cute.Tensor, + num_physical_kv_pages: cutlass.Int64, + summary: cute.Tensor, + stream: cuda_drv.CUstream, + ) -> None: + self.inspect_requests( + paged_kv_indptr, + paged_kv_indices, + seq_lens_kv, + num_physical_kv_pages, + summary, + stream, + ) + self.inspect_bsr( + block_indptr, + block_indices, + seq_lens_kv, + summary, + stream, + ) + + +def _fake_compact( + dtype: type, + shape: tuple[object, ...], + *, + alignment: int, +) -> cute.Tensor: + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(reversed(range(len(shape)))), + assumed_align=alignment, + ) + + +@functools.cache +def compile_block_sparse_inspection( + *, + device_index: int, + batch_size: int, + num_kv_heads: int, + seq_len_q: int, + seq_len_kv: int, + q_block_size: int, + kv_block_size: int, +) -> Callable[..., None]: + """Compile one geometry specialization while keeping ``indices[nnz]`` dynamic. + + Tensor ranks, dtypes, compact strides, and all attention geometry are part + of the specialization. Only the flat ``block_indices`` extent is symbolic; + tensor contents may of course vary between calls to the cached function. + """ + + num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + logical_nnz = cute.sym_int() + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = _InspectBlockSparseBsr( + batch_size=batch_size, + num_kv_heads=num_kv_heads, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + q_block_size=q_block_size, + kv_block_size=kv_block_size, + ) + with torch.cuda.device(device_index): + return cute.compile( + kernel, + _fake_compact( + cutlass.Int32, + (batch_size, num_kv_heads, num_q_block_rows + 1), + alignment=4, + ), + _fake_compact(cutlass.Int32, (logical_nnz,), alignment=4), + None, + _fake_compact(cutlass.Int64, (_SUMMARY_FIELDS,), alignment=8), + stream, + options=_COMPILE_OPTIONS, + ) + + +@functools.cache +def compile_paged_block_sparse_metadata_inspection( + *, + device_index: int, + batch_size: int, + num_kv_heads: int, + seq_len_q: int, + minimum_seq_len_kv: int, + max_seq_len_kv: int, + q_block_size: int, + kv_block_size: int, + page_size: int, +) -> Callable[..., None]: + """Compile one paged metadata entry that launches request then live-BSR.""" + + num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + logical_page_capacity = cute.sym_int() + logical_nnz = cute.sym_int() + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + inspect_requests = _InspectPagedKvMetadata( + batch_size=batch_size, + minimum_seq_len_kv=minimum_seq_len_kv, + max_seq_len_kv=max_seq_len_kv, + page_size=page_size, + ) + inspect_bsr = _InspectBlockSparseBsr( + batch_size=batch_size, + num_kv_heads=num_kv_heads, + seq_len_q=seq_len_q, + seq_len_kv=None, + q_block_size=q_block_size, + kv_block_size=kv_block_size, + ) + inspect_paged_block_sparse_metadata = _InspectPagedBlockSparseMetadata( + inspect_requests=inspect_requests, + inspect_bsr=inspect_bsr, + ) + + with torch.cuda.device(device_index): + return cute.compile( + inspect_paged_block_sparse_metadata, + _fake_compact( + cutlass.Int32, + (batch_size, num_kv_heads, num_q_block_rows + 1), + alignment=4, + ), + _fake_compact(cutlass.Int32, (logical_nnz,), alignment=4), + _fake_compact(cutlass.Int32, (batch_size + 1,), alignment=4), + _fake_compact(cutlass.Int32, (logical_page_capacity,), alignment=4), + _fake_compact(cutlass.Int32, (batch_size,), alignment=4), + cutlass.Int64(1), + _fake_compact(cutlass.Int64, (_SUMMARY_FIELDS,), alignment=8), + stream, + options=_COMPILE_OPTIONS, + ) + + +__all__ = [ + "compile_block_sparse_inspection", + "compile_paged_block_sparse_metadata_inspection", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py new file mode 100644 index 000000000000..6a61d0be86cd --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py @@ -0,0 +1,733 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prepare live canonical BSR rows for the PrimTS FMHA route consumer. + +The kernel converts caller-owned semantic KV blocks into fixed-stride route +metadata on every run. Route origins are logical KV-token coordinates; +the paged specialization also resolves each origin to a physical page ID for +the attention load path. One warp handles one BSR row and iterates only that +row's live routes, while four warps share a CTA. + +``row_route_offsets`` is a separate plan-owned immutable Int32 tensor. +``route_workspace`` contains only mutable row counts and route metadata +described by ``_BlockSparseRouteLayout``. Payload outside each live row +count is intentionally stale. ``max_blocks_per_row`` is the plan-declared +semantic BSR-block limit, which remains distinct from packed-route capacity. +""" + +import math +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda_drv +from cutlass.cute.testing import assert_ as runtime_assert +from cutlass.experimental import primitives as prims + +from ..._block_sparse.prepared import ( + _PREPARED_ROUTE_IS_FULL_FLAG, + _BlockSparseRouteLayout, +) +from .block_sparse_inspect import _validate_bsr_row_lane +from .fmha_decode_resources.helpers_common import _warp_broadcast_i32 + + +_WARPS_PER_CTA = 4 +_WARP_SIZE = 32 +_THREADS_PER_CTA = _WARPS_PER_CTA * _WARP_SIZE + + +@dataclass(frozen=True) +class _PreparedRouteConfig: + """Compile-time geometry shared by contiguous and paged route packing.""" + + num_kv_heads: int + num_q_block_rows: int + num_kv_blocks: int + num_rows: int + seq_len_kv: int + kv_block_size: int + atom_size: int + logical_origins_per_route: int + token_words_per_route: int + atom_valid_mask_word_offset: int + route_flags_word_offset: int + token_words_word_offset: int + has_token_bits: bool + route_metadata_stride_words: int + route_metadata_base_word_offset: int + + @staticmethod + def create( + *, + layout: _BlockSparseRouteLayout, + num_kv_heads: int, + seq_len_q: int, + seq_len_kv: int, + q_block_size: int, + kv_block_size: int, + ) -> "_PreparedRouteConfig": + """Build shared prepare geometry without adding a storage-mode flag.""" + + num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + return _PreparedRouteConfig( + num_kv_heads=num_kv_heads, + num_q_block_rows=num_q_block_rows, + num_kv_blocks=(seq_len_kv + kv_block_size - 1) // kv_block_size, + num_rows=layout.num_rows, + seq_len_kv=seq_len_kv, + kv_block_size=kv_block_size, + atom_size=layout.atom_size, + logical_origins_per_route=layout.logical_origins_per_route, + token_words_per_route=layout.token_words_per_route, + atom_valid_mask_word_offset=layout.atom_valid_mask_word_offset, + route_flags_word_offset=layout.route_flags_word_offset, + token_words_word_offset=( + layout.token_words_word_offset + if layout.token_words_word_offset is not None + else 0 + ), + has_token_bits=layout.has_token_bits, + route_metadata_stride_words=layout.route_metadata_stride_words, + route_metadata_base_word_offset=layout.route_metadata_base_word_offset, + ) + + +def _positive_i32_ceil_div( + value: cutlass.Int32, + divisor: cutlass.Constexpr[int], +) -> cutlass.Int32: + """Ceil-divide a positive Int32 without overflowing its upper bound.""" + + return (value - cutlass.Int32(1)) // cutlass.Int32(divisor) + cutlass.Int32(1) + + +@cute.jit +def _retained_atom_count( + block_indices: cute.Tensor, + row_begin: cutlass.Int32, + row_end: cutlass.Int32, + kv_block_size: cutlass.Constexpr[int], + atom_size: cutlass.Constexpr[int], + seq_len_kv: cutlass.Int32, +) -> cutlass.Int32: + """Count selected atoms whose logical origin precedes ``seq_len_kv``.""" + + row_nnz = row_end - row_begin + retained_atoms = cutlass.Int32(0) + if row_nnz > cutlass.Int32(0): + atoms_per_block = kv_block_size // atom_size + retained_atoms = (row_nnz - cutlass.Int32(1)) * cutlass.Int32(atoms_per_block) + last_block_idx = cutlass.Int32(block_indices[row_end - cutlass.Int32(1)]) + last_block_origin = last_block_idx * cutlass.Int32(kv_block_size) + remaining_tokens = cutlass.Int32(seq_len_kv) - last_block_origin + runtime_assert( + remaining_tokens > cutlass.Int32(0), + "block_indices row exceeds the live KV block range", + ) + retained_last_atoms = (remaining_tokens - cutlass.Int32(1)) // cutlass.Int32( + atom_size + ) + cutlass.Int32(1) + if retained_last_atoms > cutlass.Int32(atoms_per_block): + retained_last_atoms = cutlass.Int32(atoms_per_block) + retained_atoms = retained_atoms + retained_last_atoms + return retained_atoms + + +@cute.jit +def _resolve_route_logical_atom_origin( + block_indices: cute.Tensor, + row_begin: cutlass.Int32, + row_end: cutlass.Int32, + route_idx: cutlass.Int32, + atom_in_route: cutlass.Int32, + kv_block_size: cutlass.Constexpr[int], + atom_size: cutlass.Constexpr[int], + logical_origins_per_route: cutlass.Constexpr[int], + seq_len_kv: cutlass.Int32, +) -> tuple[cutlass.Int32, cutlass.Boolean]: + """Resolve one route atom to its logical KV-token origin.""" + + atoms_per_block = kv_block_size // atom_size + flat_atom_idx = route_idx * cutlass.Int32(logical_origins_per_route) + atom_in_route + bsr_entry_offset = flat_atom_idx // cutlass.Int32(atoms_per_block) + atom_in_block = flat_atom_idx % cutlass.Int32(atoms_per_block) + valid = cutlass.Boolean(bsr_entry_offset < row_end - row_begin) + logical_origin = cutlass.Int32(-1) + if valid: + block_idx = cutlass.Int32(block_indices[row_begin + bsr_entry_offset]) + block_origin = block_idx * cutlass.Int32(kv_block_size) + atom_offset = atom_in_block * cutlass.Int32(atom_size) + valid = cutlass.Boolean(atom_offset < cutlass.Int32(seq_len_kv) - block_origin) + if valid: + logical_origin = block_origin + atom_offset + return logical_origin, valid + + +@cute.jit +def _load_coarse_token_word( + block_indices: cute.Tensor, + kv_valid_bits: cute.Tensor, + row_begin: cutlass.Int32, + row_end: cutlass.Int32, + route_idx: cutlass.Int32, + logical_word_idx: cutlass.Int32, + batch_idx: cutlass.Int32, + kv_block_size: cutlass.Constexpr[int], + atom_size: cutlass.Constexpr[int], + logical_origins_per_route: cutlass.Constexpr[int], + seq_len_kv: cutlass.Int32, +) -> cutlass.Uint32: + """Load one logical K32 word from a coarse atom larger than K32.""" + + logical_word = cutlass.Uint32(0) + words_per_atom = atom_size // 32 + atom_in_route = logical_word_idx // cutlass.Int32(words_per_atom) + word_in_atom = logical_word_idx % cutlass.Int32(words_per_atom) + logical_origin, valid = _resolve_route_logical_atom_origin( + block_indices, + row_begin, + row_end, + route_idx, + atom_in_route, + kv_block_size, + atom_size, + logical_origins_per_route, + seq_len_kv, + ) + logical_word_origin = logical_origin + word_in_atom * cutlass.Int32(32) + if valid and logical_word_origin < cutlass.Int32(seq_len_kv): + valid_bits_word_idx = logical_word_origin >> cutlass.Int32(5) + logical_word = cutlass.Uint32(kv_valid_bits[batch_idx, valid_bits_word_idx]) + remaining_tokens = cutlass.Int32(seq_len_kv) - logical_word_origin + if remaining_tokens < cutlass.Int32(32): + logical_word = logical_word & ( + (cutlass.Uint32(1) << remaining_tokens) - cutlass.Uint32(1) + ) + return logical_word + + +@cute.jit +def _load_atom_token_chunk( + kv_valid_bits: cute.Tensor, + batch_idx: cutlass.Int32, + logical_origin: cutlass.Int32, + origin_is_valid: cutlass.Boolean, + atom_size: cutlass.Constexpr[int], + seq_len_kv: cutlass.Int32, +) -> cutlass.Uint32: + """Load the <=K32 mask chunk owned by one resolved-origin lane.""" + + token_chunk = cutlass.Uint32(0) + if origin_is_valid: + valid_bits_word_idx = logical_origin >> cutlass.Int32(5) + source_word = cutlass.Uint32(kv_valid_bits[batch_idx, valid_bits_word_idx]) + token_chunk = source_word >> (logical_origin & cutlass.Int32(31)) + token_chunk = token_chunk & cutlass.Uint32((1 << atom_size) - 1) + remaining_tokens = cutlass.Int32(seq_len_kv) - logical_origin + if remaining_tokens < cutlass.Int32(atom_size): + token_chunk = token_chunk & ( + (cutlass.Uint32(1) << remaining_tokens) - cutlass.Uint32(1) + ) + return token_chunk + + +@cute.jit +def _resolve_prepared_bsr_row( + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + linear_row_idx: cutlass.Int32, + lane_idx: cutlass.Int32, + row_is_valid: cutlass.Boolean, + cfg: cutlass.Constexpr[_PreparedRouteConfig], +) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Resolve one trusted canonical runtime BSR row.""" + + row_begin = cutlass.Int32(0) + row_end = cutlass.Int32(0) + batch_idx = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0) and row_is_valid: + q_block_row_idx = linear_row_idx % cfg.num_q_block_rows + linear_batch_head_idx = linear_row_idx // cfg.num_q_block_rows + kv_head_idx = linear_batch_head_idx % cfg.num_kv_heads + batch_idx = linear_batch_head_idx // cfg.num_kv_heads + row_begin = cutlass.Int32(block_indptr[batch_idx, kv_head_idx, q_block_row_idx]) + row_end = cutlass.Int32( + block_indptr[batch_idx, kv_head_idx, q_block_row_idx + 1] + ) + num_indices = cutlass.Int32(cute.size(block_indices)) + runtime_assert( + row_begin >= cutlass.Int32(0) + and row_begin <= row_end + and row_end <= num_indices, + "block_indptr row must be bounded and monotone", + ) + row_begin = _warp_broadcast_i32(row_begin, 0) + row_end = _warp_broadcast_i32(row_end, 0) + batch_idx = _warp_broadcast_i32(batch_idx, 0) + + row_error_code = cutlass.Int32(0) + if row_is_valid: + row_error_code = _validate_bsr_row_lane( + block_indices, + row_begin, + row_end, + lane_idx, + cfg.num_kv_blocks, + ) + row_error_code = cutlass.Int32(cute.arch.warp_redux_sync(row_error_code, "max")) + if lane_idx == cutlass.Int32(0) and row_is_valid: + runtime_assert( + row_error_code == cutlass.Int32(0), + "block_indices row must be canonical and in range", + ) + return row_begin, row_end, batch_idx + + +@cute.jit +def _publish_prepared_route_count( + block_indices: cute.Tensor, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + row_begin: cutlass.Int32, + row_end: cutlass.Int32, + linear_row_idx: cutlass.Int32, + lane_idx: cutlass.Int32, + row_is_valid: cutlass.Boolean, + max_blocks_per_row: cutlass.Int32, + seq_len_kv: cutlass.Int32, + cfg: cutlass.Constexpr[_PreparedRouteConfig], +) -> tuple[cutlass.Int32, cutlass.Int32]: + """Assert semantic capacity, publish the header, and return its live span.""" + + row_route_begin = cutlass.Int32(0) + required_route_count = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0) and row_is_valid: + row_route_begin = cutlass.Int32(row_route_offsets[linear_row_idx]) + selected_block_count = row_end - row_begin + runtime_assert( + selected_block_count <= max_blocks_per_row, + "selected BSR blocks exceed planned semantic capacity", + ) + retained_atom_count = _retained_atom_count( + block_indices, + row_begin, + row_end, + cfg.kv_block_size, + cfg.atom_size, + seq_len_kv, + ) + required_route_count = ( + retained_atom_count + cutlass.Int32(cfg.logical_origins_per_route - 1) + ) // cutlass.Int32(cfg.logical_origins_per_route) + route_workspace[linear_row_idx] = required_route_count + row_route_begin = _warp_broadcast_i32(row_route_begin, 0) + required_route_count = _warp_broadcast_i32(required_route_count, 0) + return required_route_count, row_route_begin + + +@cute.jit +def _store_prepared_route_validity( + block_indices: cute.Tensor, + kv_valid_bits: cute.Tensor, + route_workspace: cute.Tensor, + row_begin: cutlass.Int32, + row_end: cutlass.Int32, + route_idx: cutlass.Int32, + batch_idx: cutlass.Int32, + lane_idx: cutlass.Int32, + logical_origin: cutlass.Int32, + logical_origin_is_valid: cutlass.Boolean, + stored_atom_is_full: cutlass.Boolean, + route_metadata_word_index: cutlass.Int32, + seq_len_kv: cutlass.Int32, + cfg: cutlass.Constexpr[_PreparedRouteConfig], +) -> None: + """Store storage-independent atom, token, and route validity metadata.""" + + stored_atom_valid_mask = cutlass.Int32( + cute.arch.vote_ballot_sync(logical_origin_is_valid) + ) + structural_route_is_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) or stored_atom_is_full + ) + route_is_full = structural_route_is_full + if cutlass.const_expr(cfg.has_token_bits): + token_word = cutlass.Uint32(0) + if cutlass.const_expr(cfg.atom_size <= 32): + token_chunk = _load_atom_token_chunk( + kv_valid_bits, + batch_idx, + logical_origin, + logical_origin_is_valid, + cfg.atom_size, + seq_len_kv, + ) + atoms_per_word = 32 // cfg.atom_size + if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): + atom_in_word = lane_idx % cutlass.Int32(atoms_per_word) + token_word = token_chunk << ( + atom_in_word * cutlass.Int32(cfg.atom_size) + ) + active_origin_lanes = (1 << cfg.logical_origins_per_route) - 1 + for shuffle_step in cutlass.range_constexpr( + int(math.log2(atoms_per_word)) + ): + peer_word = cutlass.Uint32( + prims.shfl_sync( + thread_mask=active_origin_lanes, + val=token_word, + offset=1 << shuffle_step, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + token_word = token_word | peer_word + if atom_in_word == cutlass.Int32(0): + logical_word_idx = lane_idx // cutlass.Int32(atoms_per_word) + route_workspace[ + route_metadata_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + logical_word_idx + ] = cutlass.Int32(token_word) + full_atom_mask = cutlass.Uint32((1 << cfg.atom_size) - 1) + token_route_is_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) + or token_chunk == full_atom_mask + ) + else: + if lane_idx < cutlass.Int32(cfg.token_words_per_route): + token_word = _load_coarse_token_word( + block_indices, + kv_valid_bits, + row_begin, + row_end, + route_idx, + lane_idx, + batch_idx, + cfg.kv_block_size, + cfg.atom_size, + cfg.logical_origins_per_route, + seq_len_kv, + ) + route_workspace[ + route_metadata_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + lane_idx + ] = cutlass.Int32(token_word) + token_route_is_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.token_words_per_route) + or token_word == cutlass.Uint32(0xFFFFFFFF) + ) + route_is_full = cutlass.Boolean( + structural_route_is_full and token_route_is_full + ) + + if lane_idx == cutlass.Int32(0): + route_workspace[ + route_metadata_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) + ] = stored_atom_valid_mask + route_workspace[ + route_metadata_word_index + cutlass.Int32(cfg.route_flags_word_offset) + ] = ( + cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + if route_is_full + else cutlass.Int32(0) + ) + + +@cute.jit +def _paged_request_page_range_is_valid( + request_begin: cutlass.Int32, + request_end: cutlass.Int32, + num_indices: cutlass.Int32, + required_pages: cutlass.Int32, +) -> cutlass.Boolean: + """Validate one request's page-table range before any index load.""" + + return cutlass.Boolean( + request_begin >= cutlass.Int32(0) + and request_begin <= request_end + and request_end <= num_indices + and request_end - request_begin >= required_pages + ) + + +@cute.jit +def _resolve_paged_route_atom_page_id( + paged_kv_indices: cute.Tensor, + request_begin: cutlass.Int32, + logical_origin: cutlass.Int32, + logical_origin_is_valid: cutlass.Boolean, + lane_idx: cutlass.Int32, + page_size: cutlass.Constexpr[int], + num_physical_kv_pages: cutlass.Int64, +) -> cutlass.Int32: + """Resolve one trusted selected logical atom to its raw physical page ID.""" + + physical_page_id = cutlass.Int32(-1) + page_id_is_valid = cutlass.Boolean(True) + if logical_origin_is_valid: + logical_page_idx = logical_origin // cutlass.Int32(page_size) + candidate_page_id = cutlass.Int32( + paged_kv_indices[request_begin + logical_page_idx] + ) + physical_page_id = candidate_page_id + page_id_is_valid = cutlass.Boolean( + candidate_page_id >= cutlass.Int32(0) + and cutlass.Int64(candidate_page_id) < num_physical_kv_pages + ) + page_ids_are_valid = cute.arch.vote_all_sync(page_id_is_valid) + if lane_idx == cutlass.Int32(0): + runtime_assert( + page_ids_are_valid, + "paged_kv_indices contains an out-of-range physical page ID", + ) + return physical_page_id + + +class _PrepareBlockSparseRoutes: + """Prepare contiguous or paged sparse routes for one static geometry.""" + + def __init__( + self, + *, + batch_size: int, + num_kv_heads: int, + seq_len_q: int, + seq_len_kv: int, + q_block_size: int, + kv_block_size: int, + kv_route_size: int, + has_token_bits: bool, + page_size: int | None = None, + mask_type: str, + ) -> None: + if mask_type not in ("dense", "causal"): + raise ValueError(f"unsupported mask_type: {mask_type}") + num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + num_rows = batch_size * num_kv_heads * num_q_block_rows + layout = _BlockSparseRouteLayout.create( + kv_route_size=kv_route_size, + kv_block_size=kv_block_size, + page_size=page_size, + has_token_bits=has_token_bits, + route_metadata_capacity=0, + num_rows=num_rows, + ) + self.cfg = _PreparedRouteConfig.create( + layout=layout, + num_kv_heads=num_kv_heads, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + q_block_size=q_block_size, + kv_block_size=kv_block_size, + ) + self.route_layout = layout + self.page_size = page_size if page_size is not None else 1 + self.minimum_seq_len_kv = seq_len_q if mask_type == "causal" else 1 + self.physical_page_ids_word_offset = ( + layout.physical_page_ids_word_offset if layout.is_paged else 0 + ) + self.route_metadata_base_word_offset = layout.route_metadata_base_word_offset + + @cute.jit + def __call__( + self, + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + kv_valid_bits: cute.Tensor, + seq_lens_kv: cute.Tensor | None, + paged_kv_indptr: cute.Tensor | None, + paged_kv_indices: cute.Tensor | None, + num_physical_kv_pages: cutlass.Int64, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + max_blocks_per_row: cutlass.Int32, + stream: cuda_drv.CUstream, + ) -> None: + """Launch four independent row preparers per CTA.""" + + self.kernel( + block_indptr, + block_indices, + kv_valid_bits, + seq_lens_kv, + paged_kv_indptr, + paged_kv_indices, + num_physical_kv_pages, + row_route_offsets, + route_workspace, + max_blocks_per_row, + ).launch( + grid=[ + (self.cfg.num_rows + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, + 1, + 1, + ], + block=[_THREADS_PER_CTA, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + kv_valid_bits: cute.Tensor, + seq_lens_kv: cute.Tensor | None, + paged_kv_indptr: cute.Tensor | None, + paged_kv_indices: cute.Tensor | None, + num_physical_kv_pages: cutlass.Int64, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + max_blocks_per_row: cutlass.Int32, + ) -> None: + """Pack logical routes and, when paged, translate physical locators.""" + + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = thread_idx // _WARP_SIZE + lane_idx = thread_idx % _WARP_SIZE + linear_row_idx = block_idx * _WARPS_PER_CTA + warp_idx + row_is_valid = linear_row_idx < self.cfg.num_rows + + row_begin, row_end, batch_idx = _resolve_prepared_bsr_row( + block_indptr, + block_indices, + linear_row_idx, + lane_idx, + row_is_valid, + self.cfg, + ) + + request_begin = cutlass.Int32(0) + live_seq_len_kv = cutlass.Int32(self.cfg.seq_len_kv) + if cutlass.const_expr(self.route_layout.is_paged): + raw_seq_len_kv = cutlass.Int32(self.cfg.seq_len_kv) + if lane_idx == cutlass.Int32(0) and row_is_valid: + raw_seq_len_kv = cutlass.Int32(seq_lens_kv[batch_idx]) + runtime_assert( + raw_seq_len_kv >= cutlass.Int32(self.minimum_seq_len_kv) + and raw_seq_len_kv <= cutlass.Int32(self.cfg.seq_len_kv), + "seq_lens_kv is outside the planned live-length range", + ) + raw_seq_len_kv = _warp_broadcast_i32(raw_seq_len_kv, 0) + live_seq_len_kv = raw_seq_len_kv + + if lane_idx == cutlass.Int32(0) and row_is_valid: + required_pages = _positive_i32_ceil_div( + live_seq_len_kv, + self.page_size, + ) + request_begin = cutlass.Int32(paged_kv_indptr[batch_idx]) + request_end = cutlass.Int32( + paged_kv_indptr[batch_idx + cutlass.Int32(1)] + ) + metadata_starts_at_zero = cutlass.Boolean( + paged_kv_indptr[cutlass.Int32(0)] == cutlass.Int32(0) + ) + runtime_assert( + metadata_starts_at_zero + and _paged_request_page_range_is_valid( + request_begin, + request_end, + cutlass.Int32(cute.size(paged_kv_indices)), + required_pages, + ), + "paged_kv_indptr row lacks the required live page capacity", + ) + request_begin = _warp_broadcast_i32(request_begin, 0) + + route_count, row_route_begin = _publish_prepared_route_count( + block_indices, + row_route_offsets, + route_workspace, + row_begin, + row_end, + linear_row_idx, + lane_idx, + row_is_valid, + max_blocks_per_row, + live_seq_len_kv, + self.cfg, + ) + + route_idx = cutlass.Int32(0) + while route_idx < route_count: + route_ordinal = row_route_begin + route_idx + route_metadata_word_index = cutlass.Int32( + self.cfg.route_metadata_base_word_offset + ) + route_ordinal * cutlass.Int32(self.cfg.route_metadata_stride_words) + logical_origin = cutlass.Int32(-1) + logical_origin_is_valid = cutlass.Boolean(False) + physical_page_id = cutlass.Int32(-1) + atom_is_full = cutlass.Boolean(False) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + ( + logical_origin, + logical_origin_is_valid, + ) = _resolve_route_logical_atom_origin( + block_indices, + row_begin, + row_end, + route_idx, + lane_idx, + self.cfg.kv_block_size, + self.cfg.atom_size, + self.cfg.logical_origins_per_route, + live_seq_len_kv, + ) + if cutlass.const_expr(self.route_layout.is_paged): + physical_page_id = _resolve_paged_route_atom_page_id( + paged_kv_indices, + request_begin, + logical_origin, + logical_origin_is_valid, + lane_idx, + self.page_size, + num_physical_kv_pages, + ) + if logical_origin_is_valid: + atom_is_full = cutlass.Boolean( + logical_origin + <= live_seq_len_kv - cutlass.Int32(self.cfg.atom_size) + ) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + route_workspace[route_metadata_word_index + lane_idx] = logical_origin + if cutlass.const_expr(self.route_layout.is_paged): + route_workspace[ + route_metadata_word_index + + cutlass.Int32(self.physical_page_ids_word_offset) + + lane_idx + ] = physical_page_id + + _store_prepared_route_validity( + block_indices, + kv_valid_bits, + route_workspace, + row_begin, + row_end, + route_idx, + batch_idx, + lane_idx, + logical_origin, + logical_origin_is_valid, + atom_is_full, + route_metadata_word_index, + live_seq_len_kv, + self.cfg, + ) + route_idx = route_idx + cutlass.Int32(1) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py new file mode 100644 index 000000000000..c67e4c83bf10 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py @@ -0,0 +1,3985 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Configuration for the FMHA decode TS kernel. + +The class FmhaDecodeConfig encapsulates static configuration parameters +that should be set before kernel compilation. +""" + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, fields, replace + +import cutlass.utils as utils +from cutlass import BFloat16, Float16, Float32, Float8E4M3FN + +from ..._block_sparse.common import ( + _block_sparse_kv_atom_size, + _select_block_sparse_q_tile_size, + _validate_sparse_kv_block_size, +) +from ...split_kv_mode_policy import select_split_kv_modes +from .fmha_decode_constants import ( + AUTO_LAUNCH_TILE_SIZE_KV, + BITS_PER_BYTE, + BYTES_PER_KIB, + FALLBACK_SM_COUNT_B200, + FP32_BYTES, + FP8_OUTPUT_ELEMENTS_PER_REG_GROUP, + FP8_P_PACKED_REGS_PER_Q_REPEAT, + FP8_VALUES_PER_REG, + FP16_OUTPUT_ELEMENTS_PER_REG_GROUP, + FP16_P_PACKED_REGS_PER_Q_REPEAT, + FP16_VALUES_PER_REG, + KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES, + KV_TILE_256_SHARED_FIFO_STAGES, + MAX_CLUSTER_DIM_X, + MAX_CLUSTER_PARTIAL_SMEM_BYTES, + MAX_KV_STAGE_SMEM_KIB, + MAX_WARP_GROUPS, + MIN_LOOP_ITERS_PER_SPLIT, + PARALLEL_REDUCTION_BYTES_PER_SLICE, + PARALLEL_REDUCTION_THREADS_PER_CTA, + PARTIAL_O_ELEMENT_BYTES, + PARTIAL_STATS_VALUES_PER_ROW, + Q_REPETITION_GROUP_HEADS, + Q_ROW_ALIGNMENT_BYTES, + REDUCTION_BYTES_PER_SLICE, + REDUCTION_THREADS_PER_CTA, + SPLIT_KV_MIN_TOKENS_PER_CTA, + TMEM_COLUMNS_PER_ROW, + TMEM_ROW_STRIDE, + TOTAL_SMEM_BUDGET_KIB, + WARP_THREADS, +) + +ConfigValue = int | float | bool | str | type | None +_GroupedKeepsProfileKey = tuple[type, type, type, int, int, int, int] + +# Public APIs use the strings ``dense`` and ``causal``. Keep the value carried +# through FmhaDecodeConfig as a small integer so mask selection remains a +# compile-time predicate in CuTe DSL kernels. +DENSE = 0 +CAUSAL = 1 +MASK_TYPES = ("dense", "causal") + +# Every correction lane reduces one packed 16-byte partial-O vector. The +# default four-warp correction group therefore owns a 2-KiB reducer slice. +SPLIT_REDUCTION_VECTOR_BYTES_PER_THREAD = 16 + +_GROUPED_KEEPS_MAIN_PROFILE: _GroupedKeepsProfileKey = ( + Float16, + Float16, + Float16, + 128, + 0, + 2, + 2, +) +# Block-sparse feature compatibility is validated by +# ``validate_block_sparse_profile``. This set only selects the Keeps MMA +# resource recipes qualified for that already-validated launch domain. +_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES = { + _GROUPED_KEEPS_MAIN_PROFILE, + (BFloat16, BFloat16, BFloat16, 128, 0, 2, 2), +} +_GROUPED_KEEPS_STATIC_ONLY_PROFILES = { + (Float8E4M3FN, Float8E4M3FN, Float16, 128, 0, 2, 2), + (BFloat16, BFloat16, BFloat16, 64, 0, 2, 2), + (Float16, Float16, Float16, 256, 128, 1, 1), +} + +_KV_TILE_256_PHYSICAL_DEFAULTS: Mapping[str, ConfigValue] = { + "tmem_s_cols": 128, + "tmem_stats_cols": 32, + "tmem_p_cols": 64, + "tmem_o_cols": 128, + "mma_tile_m_bmm1": 64, + "mma_tile_n_bmm1": 256, + "mma_tile_m_bmm2": 64, + "mma_tile_n_bmm2": 256, + "q_stages": 1, + "kv_stages": KV_TILE_256_SHARED_FIFO_STAGES, + "head_dim_per_stage_kv": 0, + "num_insts_kv": 2, + "o_stages": 2, +} + +# KV256 currently uses one validated 16-warp role assignment. Keep these +# defaults next to the physical profile; explicit alternatives remain intact +# until the common profile validator determines whether they are supported. +_KV_TILE_256_TASK_TOPOLOGY_DEFAULTS: Mapping[str, int] = { + "softmax0_warp_idx": 0, + "softmax0_num_warps": 4, + "softmax1_warp_idx": 4, + "softmax1_num_warps": 4, + "correction_warp_idx": 8, + "correction_num_warps": 4, + "mma_warp_idx": 12, + "mma_num_warps": 1, + "load_warp_idx": 13, + "load_num_warps": 1, + "page_offsets_warp_idx": 14, + "page_offsets_num_warps": 1, + "scheduler_warp_idx": 13, + "scheduler_num_warps": 1, +} + +_KV_TILE_256_TUNABLE_FIELDS = frozenset(("kv_stages",)) + +# Public cost-model collection uses the FP8 proxy for every source dtype. The +# original fixed-Q1 ratio-32 requests exercise a partial grouped-Q tile; the +# shape-aware path also admits complete fixed multi-Q tiles at any legal head +# ratio. These are profile families rather than shape exceptions: batch size, +# KV length, tile choice, and legal GMEM split fanout remain unrestricted by +# this declaration. +_GROUPED_KEEPS_PAGED_FP8_PROFILES = { + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 64, 0, 2, 2), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 128, 0, 2, 2), + (Float8E4M3FN, Float8E4M3FN, Float16, 256, 128, 1, 1), +} + + +def normalize_mask_type( + mask_type: str | int | None, + *, + sliding_window_causal: bool = False, +) -> int: + """Return the constexpr mask id for a public mask selection. + + ``None`` uses the decode defaults: ordinary decode is dense, while + requesting a causal sliding window implies a causal right bound. An + explicit dense mask together with a causal sliding window is contradictory + and rejected instead of silently changing the requested mask. + """ + + if mask_type is None: + return CAUSAL if sliding_window_causal else DENSE + if isinstance(mask_type, bool): + raise ValueError("mask_type must be 'dense' or 'causal', not a boolean") + if isinstance(mask_type, int): + if mask_type not in (DENSE, CAUSAL): + raise ValueError("internal mask_type must be DENSE or CAUSAL") + normalized = mask_type + elif isinstance(mask_type, str): + value = mask_type.lower() + if value not in MASK_TYPES: + raise ValueError( + f"mask_type must be one of {MASK_TYPES}, got {mask_type!r}" + ) + normalized = CAUSAL if value == "causal" else DENSE + else: + raise ValueError( + f"mask_type must be one of {MASK_TYPES}, got {type(mask_type).__name__}" + ) + if sliding_window_causal and normalized != CAUSAL: + raise ValueError( + "sliding_window_causal requires mask_type='causal'; omit mask_type " + "to select causal implicitly or pass mask_type='causal' explicitly" + ) + return normalized + + +def mask_type_name(mask_type: int) -> str: + """Return the public string for a normalized constexpr mask id.""" + + if mask_type == DENSE: + return "dense" + if mask_type == CAUSAL: + return "causal" + raise ValueError(f"internal mask_type must be DENSE or CAUSAL, got {mask_type!r}") + + +def _q_tokens_per_cta( + rows_per_cta: int, heads_q_per_kv: int, groups_tokens_heads_q: bool +) -> int: + """Return complete Q tokens packed into one CTA.""" + if groups_tokens_heads_q: + return rows_per_cta // heads_q_per_kv + return 1 + + +def _q_tma_rows_per_cta( + rows_per_cta: int, heads_q_per_kv: int, groups_tokens_heads_q: bool +) -> int: + """Return rows accounted by Q TMA for every CTA.""" + if groups_tokens_heads_q: + return ( + _q_tokens_per_cta(rows_per_cta, heads_q_per_kv, groups_tokens_heads_q) + * heads_q_per_kv + ) + return rows_per_cta + + +@dataclass(frozen=True) +class QTileGeometry: + """Host-side ownership and padding contract for one decode Q CTA.""" + + rows_per_cta: int + heads_q_per_kv: int + groups_tokens_heads_q: bool + + @property + def tokens_per_cta(self) -> int: + """Return complete Q tokens packed into one CTA.""" + return _q_tokens_per_cta( + self.rows_per_cta, + self.heads_q_per_kv, + self.groups_tokens_heads_q, + ) + + @property + def head_ctas_per_token(self) -> int: + """Return the number of head-band CTAs assigned to one Q token.""" + if self.groups_tokens_heads_q: + return 1 + return (self.heads_q_per_kv + self.rows_per_cta - 1) // self.rows_per_cta + + @property + def tma_rows_per_cta(self) -> int: + """Return rows accounted by Q TMA for every CTA.""" + return _q_tma_rows_per_cta( + self.rows_per_cta, + self.heads_q_per_kv, + self.groups_tokens_heads_q, + ) + + def num_q_ctas(self, seq_len_q: int) -> int: + """Return the number of Q CTAs for one ``(batch, KV head)`` pair.""" + if seq_len_q < 0: + raise ValueError("seq_len_q must be non-negative") + if self.groups_tokens_heads_q: + return (seq_len_q + self.tokens_per_cta - 1) // self.tokens_per_cta + return seq_len_q * self.head_ctas_per_token + + +@dataclass(frozen=True) +class GroupedQMmaCandidate: + """One grouped-Q MMA geometry candidate.""" + + variant: str + tile_size_q: int + q_tokens_per_cta: int + q_tiles: int + + +@dataclass(frozen=True) +class GroupedQLaunchCandidate: + """One production-valid grouped-Q MMA and KV-split launch recipe.""" + + mma: GroupedQMmaCandidate + split_kv_mode: str + splits_kv: int + base_ctas: int + launched_ctas: int + seq_len_per_cta_kv: int + waves: int + modeled_time: float + + +_GROUPED_Q_MMA_TILES = ( + ("swaps_mma_ab", 8), + ("swaps_mma_ab", 16), + ("swaps_mma_ab", 32), + ("keeps_mma_ab", 64), + ("keeps_mma_ab", 128), +) + +# Empirical GQA-generation factors derived from matched reference measurements. +# The selector reuses the measured relative costs but resolves its own legal +# profiles, Q geometry, split fanout, and cluster promotion instead of copying +# a final shape policy. +_GROUPED_Q_MAINLOOP_COST = { + 8: 1.0, + 16: 1.2, + 32: 1.48, + 64: 1.68, + 128: 2.2, +} +_GROUPED_Q_REDUCTION_COST = { + 8: 1.0, + 16: 1.03, + 32: 1.08, + 64: 1.2, + 128: 1.32, +} +_GROUPED_Q_REDUCTION_SEQ_LEN_FACTOR = 128.0 + + +def enumerate_grouped_q_mma_candidates( + *, heads_q_per_kv: int, seq_len_q: int +) -> tuple[GroupedQMmaCandidate, ...]: + """Return grouped-Q candidates without applying launch policy. + + A CTA always owns an integral number of complete Q-head groups. TileQ may + leave structural padding rows and the final CTA may own fewer tokens than + its capacity; the common Q geometry and row masks already represent both + cases. This helper deliberately does not inspect dtypes, reduction modes, + scheduler state, or mutate ``FmhaDecodeConfig``. + """ + if heads_q_per_kv <= 0: + raise ValueError("heads_q_per_kv must be positive") + if seq_len_q <= 0: + raise ValueError("seq_len_q must be positive") + + candidates = [] + for variant, tile_size_q in _GROUPED_Q_MMA_TILES: + if tile_size_q < heads_q_per_kv: + continue + q_tokens_per_cta = tile_size_q // heads_q_per_kv + candidates.append( + GroupedQMmaCandidate( + variant=variant, + tile_size_q=tile_size_q, + q_tokens_per_cta=q_tokens_per_cta, + q_tiles=(seq_len_q + q_tokens_per_cta - 1) // q_tokens_per_cta, + ) + ) + return tuple(candidates) + + +def make_grouped_q_launch_candidate( + candidate: GroupedQMmaCandidate, + *, + splits_kv: int, + seq_len_kv: int, + tile_size_kv: int, + num_insts_kv: int, + batch_size: int, + num_heads_kv: int, + service_capacity: int, +) -> GroupedQLaunchCandidate: + """Build one empirical grouped-Q launch-cost record. + + The model follows the measured GQA-generation cost structure: + + ``(mainloop_factor[TileQ] * seq_len_per_cta_kv +`` + `` reduction_factor[TileQ] * 128 * splits_kv) * waves``. + + This is a small deterministic default heuristic. It deliberately uses the + same factors for every dtype; the production profile validator still + determines which actual-dtype recipes are legal. + """ + if splits_kv <= 0: + raise ValueError("splits_kv must be positive") + if seq_len_kv <= 0: + raise ValueError("seq_len_kv must be positive") + if tile_size_kv <= 0: + raise ValueError("tile_size_kv must be positive") + if num_insts_kv <= 0: + raise ValueError("num_insts_kv must be positive") + if batch_size <= 0: + raise ValueError("batch_size must be positive") + if num_heads_kv <= 0: + raise ValueError("num_heads_kv must be positive") + if service_capacity <= 0: + raise ValueError("service_capacity must be positive") + kv_tiles = (seq_len_kv + tile_size_kv - 1) // tile_size_kv + seq_len_per_cta_kv = ((kv_tiles + splits_kv - 1) // splits_kv) * tile_size_kv + base_ctas = batch_size * num_heads_kv * candidate.q_tiles + launched_ctas = base_ctas * splits_kv + waves = (launched_ctas + service_capacity - 1) // service_capacity + modeled_time = ( + _GROUPED_Q_MAINLOOP_COST[candidate.tile_size_q] * seq_len_per_cta_kv + + _GROUPED_Q_REDUCTION_COST[candidate.tile_size_q] + * _GROUPED_Q_REDUCTION_SEQ_LEN_FACTOR + * splits_kv + ) * waves + return GroupedQLaunchCandidate( + mma=candidate, + split_kv_mode=("gmem_reduction" if splits_kv > 1 else "disabled"), + splits_kv=splits_kv, + base_ctas=base_ctas, + launched_ctas=launched_ctas, + seq_len_per_cta_kv=seq_len_per_cta_kv, + waves=waves, + modeled_time=modeled_time, + ) + + +def select_grouped_q_launch_candidate( + candidates: Sequence[GroupedQLaunchCandidate], +) -> GroupedQLaunchCandidate: + """Select the lowest-cost legal grouped-Q launch recipe. + + The score compares actual Q-grid waves, so a wider tile can win only when + it removes enough repeated KV mainloop work to repay its higher per-tile + cost. This is important for partial final tiles: TileQ128 and TileQ64 can + have different Q-grid multiplicities even when both are legal. + """ + if not candidates: + raise ValueError("at least one grouped-Q launch candidate is required") + + return min( + candidates, + key=lambda recipe: ( + recipe.modeled_time, + recipe.waves, + recipe.splits_kv, + -recipe.mma.tile_size_q, + ), + ) + + +def select_grouped_q_direct_wave_candidate( + candidates: Sequence[GroupedQLaunchCandidate], +) -> GroupedQLaunchCandidate: + """Select a direct recipe using the same mainloop-aware launch score.""" + direct = tuple(recipe for recipe in candidates if recipe.splits_kv == 1) + if not direct: + raise ValueError("at least one direct grouped-Q candidate is required") + return select_grouped_q_launch_candidate(direct) + + +def make_q_tile_geometry( + *, + rows_per_cta: int, + heads_q_per_kv: int, + groups_tokens_heads_q: bool, +) -> QTileGeometry: + """Build the common grouped or one-token Q-tile geometry contract.""" + if rows_per_cta <= 0: + raise ValueError("rows_per_cta must be positive") + if heads_q_per_kv <= 0: + raise ValueError("heads_q_per_kv must be positive") + + if groups_tokens_heads_q: + if rows_per_cta < heads_q_per_kv: + raise ValueError("grouped Q tiles require rows_per_cta >= heads_q_per_kv") + return QTileGeometry( + rows_per_cta=rows_per_cta, + heads_q_per_kv=heads_q_per_kv, + groups_tokens_heads_q=True, + ) + + return QTileGeometry( + rows_per_cta=rows_per_cta, + heads_q_per_kv=heads_q_per_kv, + groups_tokens_heads_q=False, + ) + + +@dataclass +class FmhaDecodeConfig: + """ + Kernel-wide configuration for fmha_decode. + """ + + # ------------------------------------------------------------------ + # Problem shape + # ------------------------------------------------------------------ + # Per-head embedding dimension D. Supported profiles use 64, 128, or 256. + headdim: int = 128 + # Number of KV tiles in the launch (= ceil(seq_len_kv / tile_size_kv)). + # Populated by the launcher once seq_len_kv is known. + total_kv_tiles: int = 0 + # Q rows placed on the small MMA N dimension. Shape-aware selection derives + # this tile from the Q-heads-per-KV-head ratio and grouping policy. + tile_size_q: int = 8 + # Q-layout metadata. Raw shape-less configs retain neutral values; + # shape-aware selection fills them and defaults to grouped Q. + max_seq_len_q: int = 1 + # Select the packed Q/O ABI. Q and O are laid out as + # [sum_q_tokens, num_heads_q, head_dim] and indexed by cu_seqlens_q. + use_variable_seqlens_q: bool = False + heads_q_per_kv: int = 0 + groups_tokens_heads_q: bool = False + # K/V tokens per tile along the K-sequence dimension; also the MMA "M" + # dimension for BMM1 under SwapsMmaAb. + tile_size_kv: int = 128 + # Select a block-sparse launch. A positive semantic Q block is legal when + # the selected physical Q tile stays within one route row. KV blocks remain + # restricted to 8/16/32 or positive multiples of 64 and are assembled into + # a profile-selected fixed KV128 or KV256 route. + use_block_sparse: bool = False + q_block_size: int = 0 + kv_block_size: int = 0 + # Optional batch-wide physical-token validity metadata shared by every head + # and sparse row in that batch item. Keeps profiles derive their route-full + # fast path from these words at run time; it is intentionally not a separate + # plan/config specialization. + use_kv_valid_bits: bool = False + # Let independent warps issue the two fine-route K/V instruction streams. + # The wrapper resolves this from immutable route capacity; the kernel does + # not inspect live BSR morphology to select its task topology. + use_parallel_sparse_kv_loads: bool = False + # Number of K/V instances that the loop processes per step. Two instances + # (K0/V0 and K1/V1) let two parallel SoftmaxTask groups consume alternating + # tiles, improving SM utilization when tile_size_q is tiny. + num_insts_kv: int = 2 # K/V instances per loop step + + # ------------------------------------------------------------------ + # Data types + # ------------------------------------------------------------------ + # Q element type. One of Float16 / BFloat16 / Float8E4M3FN. Must equal + # kv_dtype — mixed Q/KV element types are not supported yet; enforced by + # the guard in make_decode_config. + q_dtype: type = Float16 + # K and V element type. One of Float16 / BFloat16 / Float8E4M3FN. + kv_dtype: type = Float16 + # Output O element type. One of Float16 / BFloat16 / Float8E4M3FN. + out_dtype: type = Float16 + # Accumulator type (BMM accumulators and softmax stats), always Float32 + # in the currently supported recipes. + acc_dtype: type = Float32 + + # ------------------------------------------------------------------ + # Software pipeline depths + # ------------------------------------------------------------------ + # Q TMA pipeline depth. Q is reloaded only across persistent work tiles, + # so a shallow ring is sufficient. + q_stages: int = 2 + # K/V TMA pipeline depth. Deeper KV staging hides the long + # GMEM→SMEM latency in the BMM1↔BMM2 chain. + kv_stages: int = 4 + + # ------------------------------------------------------------------ + # TMEM column counts + # ------------------------------------------------------------------ + # Columns in TMEM for one BMM1 S = Q·Kᵀ accumulator (per softmax instance). + # Equal to tile_size_q because SwapsMmaAb puts Q heads on the N axis. + tmem_s_cols: int = 8 # per instance (tileSizeQ) + # Columns in TMEM for the per-row softmax statistics (running max/sum and + # scratch) owned by one softmax instance. + tmem_stats_cols: int = 32 # per instance (softmax local stats) + # P operand width for the staged-D256 Keeps TMEM overlay; zero for Swaps. + tmem_p_cols: int = 0 + # Columns in TMEM for one O = Σ P·V accumulator (per O stage). Mirrors + # tile_size_q for the same reason as tmem_s_cols. + tmem_o_cols: int = 8 # per instance (tileSizeQ) + # Number of O accumulator stages held simultaneously in TMEM so MMA can + # start the next BMM2 while CorrectionTask rescales the previous one. + o_stages: int = 2 # O0 and O1 in TMEM + + # ------------------------------------------------------------------ + # MMA atom shapes under SwapsMmaAb + # ------------------------------------------------------------------ + # BMM1 computes S = K · Qᵀ: K provides + # the M axis (tile_size_kv), Q provides the N axis (tile_size_q). + mma_tile_m_bmm1: int = 128 + mma_tile_n_bmm1: int = 8 + # BMM2 computes O = V · Pᵀ: V provides M (headdim_v == headdim), P + # provides N (tile_size_q), and the K-reduction axis is tile_size_kv. + mma_tile_m_bmm2: int = 128 + mma_tile_n_bmm2: int = 8 + + # ------------------------------------------------------------------ + # Warp specialization layout (4 warp groups × 4 warps = 16 warps total) + # NOTE: please update `_active_warp_roles` after new roles are added. + # ------------------------------------------------------------------ + # Softmax0Task: WG0 (warps 0–3) handles even K/V instances (K0/V0). + softmax0_warp_idx: int = 0 # WG0: warps 0-3 + softmax0_num_warps: int = 4 + # Softmax1Task: WG1 (warps 4–7) handles odd K/V instances (K1/V1). + softmax1_warp_idx: int = 4 # WG1: warps 4-7 + softmax1_num_warps: int = 4 + # CorrectionTask: WG2 (warps 8–11) rescales O across instances and writes + # the epilogue to GMEM. + correction_warp_idx: int = 8 # WG2: warps 8-11 + correction_num_warps: int = 4 + # MmaTask: a single warp in WG3 issues tcgen05 MMA instructions for BMM1/BMM2. + mma_warp_idx: int = 12 # WG3: warp 12 + mma_num_warps: int = 1 + # LoadTask: a single warp in WG3 issues TMA loads for Q/K/V. + load_warp_idx: int = 13 # WG3: warp 13 + load_num_warps: int = 1 + # PageTableTask (paged-KV only): warp 14 prefetches logical→physical + # page IDs that LoadTask consumes when issuing the TMA copies. + page_offsets_warp_idx: int = 14 # WG3: warp 14 for paged-KV page table prefetch + page_offsets_num_warps: int = 1 + # SMEM pipeline depth for the prefetched page-offset table. + page_offsets_stages: int = 6 + # SchedulerTask: under persistent scheduling, warp 13 runs the CLC tile + # scheduler instead of doing TMA loads. + scheduler_warp_idx: int = 13 # Persistent: warp 13 + scheduler_num_warps: int = 1 + # ClcLoadTask: under persistent scheduling, warp 15 issues the CLC + # response loads. + clc_load_warp_idx: int = 15 + # PaddingTask placement is derived after selecting the active task roles. + # Each active warp group is compacted first, then its unused tail warps are + # assigned to the corresponding padding task. Persistent layouts retain + # all four warp groups even when the last group contains only padding. + wg0_padding_warp_idx: int = 4 + wg0_padding_num_warps: int = 0 + wg1_padding_warp_idx: int = 8 + wg1_padding_num_warps: int = 0 + wg2_padding_warp_idx: int = 12 + wg2_padding_num_warps: int = 0 + wg3_padding_warp_idx: int = 16 + wg3_padding_num_warps: int = 0 + + # ------------------------------------------------------------------ + # Task-local register allocation + # ------------------------------------------------------------------ + # The full TileQ128 Keeps graph has the largest live softmax/correction + # fragments and needs registers moved from its descriptor-only task group. + # A long KV256 graph instead moves a smaller share from Softmax to its + # heavier correction tail. Short KV256 loops avoid the fixed hand-off cost. + @property + def uses_task_register_reallocation(self) -> bool: + return self.use_keeps_mma_ab and ( + self.tile_size_q == 128 + or ( + self.tile_size_q == 64 + and self.tile_size_kv == 256 + and self.total_kv_tiles >= KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES + ) + ) + + @property + def softmax_task_num_registers(self) -> int | None: + if not self.uses_task_register_reallocation: + return None + return 176 if self.tile_size_kv == 256 else 184 + + @property + def correction_task_num_registers(self) -> int | None: + if not self.uses_task_register_reallocation: + return None + return 104 if self.tile_size_kv == 256 else 88 + + @property + def mma_load_task_num_registers(self) -> int | None: + return 56 if self.uses_task_register_reallocation else None + + # ------------------------------------------------------------------ + # SMEM allocation alignment + # ------------------------------------------------------------------ + # Alignment (bytes) for SMEM tensor allocations; 1024 is required for + # the swizzled TMA descriptors used here. + stensor_align: int = 1024 + + @property + def smem_q_tile_bytes(self) -> int: + """SMEM bytes for one Q tile (tile_size_q × headdim elements).""" + return self.tile_size_q * self.headdim * self.q_dtype_bytes + + @property + def q_tokens_per_cta(self) -> int: + """Complete Q tokens packed into one grouped CTA.""" + return _q_tokens_per_cta( + self.tile_size_q, + self.heads_q_per_kv, + self.groups_tokens_heads_q, + ) + + @property + def num_q_ctas(self) -> int: + """Return Q-CTA groups per ``(batch, KV head)`` launch tile. + + Shape-aware decode configs always carry a positive head ratio. Keep + raw shape-less configs conservative so this scheduling predicate can + never specialize geometry that has not been resolved yet. + """ + if self.heads_q_per_kv <= 0 or self.max_seq_len_q <= 0: + return max(self.max_seq_len_q, 1) + q_geometry = make_q_tile_geometry( + rows_per_cta=self.tile_size_q, + heads_q_per_kv=self.heads_q_per_kv, + groups_tokens_heads_q=self.groups_tokens_heads_q, + ) + return max(q_geometry.num_q_ctas(self.max_seq_len_q), 1) + + @property + def has_single_q_cta(self) -> bool: + """Whether every physical split maps to one logical Q CTA.""" + return ( + self.heads_q_per_kv > 0 and self.max_seq_len_q > 0 and self.num_q_ctas == 1 + ) + + @property + def q_tma_rows_per_cta(self) -> int: + """Q rows represented by the configured tensor-map box.""" + return _q_tma_rows_per_cta( + self.tile_size_q, + self.heads_q_per_kv, + self.groups_tokens_heads_q, + ) + + @property + def q_manual_padding_rows(self) -> int: + """Grouped structural Q rows completed without a TMA instruction.""" + if self.groups_tokens_heads_q: + return self.tile_size_q - self.q_tma_rows_per_cta + return 0 + + @property + def smem_kv_tile_bytes(self) -> int: + """SMEM bytes for one staged K or V tile.""" + return self.tile_size_kv * self.head_dim_kv_stage * self.kv_dtype_bytes + + @property + def head_dim_kv_stage(self) -> int: + """Head-dim slice loaded by one K/V stage.""" + return ( + self.head_dim_per_stage_kv + if self.head_dim_per_stage_kv != 0 + else self.headdim + ) + + @property + def num_head_dim_stages_kv(self) -> int: + """Number of K/V head-dim stages needed to cover headdim.""" + return (self.headdim + self.head_dim_kv_stage - 1) // self.head_dim_kv_stage + + @property + def tmem_o_cols_per_head_dim_stage(self) -> int: + """TMEM O columns owned by one head-dim stage.""" + if self.use_keeps_mma_ab: + return self.head_dim_kv_stage + return self.tile_size_q + + @property + def tmem_o_stage_cols(self) -> int: + """TMEM columns reserved for one logical O pipeline stage.""" + if self.use_keeps_mma_ab: + return self.tmem_o_cols + head_dim_stage_cols = self.tmem_o_cols_per_head_dim_stage + stages_per_tmem_row = max(TMEM_COLUMNS_PER_ROW // head_dim_stage_cols, 1) + return head_dim_stage_cols * min( + self.num_head_dim_stages_kv, stages_per_tmem_row + ) + + def swaps_head_dim_stage_tmem_offset(self, head_dim_stage_idx: int) -> int: + """Return the Swaps O offset for one staged head-dimension slice.""" + if self.use_keeps_mma_ab or self.head_dim_per_stage_kv == 0: + return 0 + stage_cols = self.tmem_o_cols_per_head_dim_stage + stages_per_tmem_row = max(TMEM_COLUMNS_PER_ROW // stage_cols, 1) + return (head_dim_stage_idx // stages_per_tmem_row) * TMEM_ROW_STRIDE + ( + head_dim_stage_idx % stages_per_tmem_row + ) * stage_cols + + def swaps_o_chunk_tmem_offset(self, chunk_idx: int) -> int: + """Return the Swaps O offset for one 64-column correction chunk.""" + if self.use_keeps_mma_ab or self.head_dim_per_stage_kv == 0: + return chunk_idx * TMEM_ROW_STRIDE + chunks_per_head_dim_stage = max(self.head_dim_kv_stage // 64, 1) + head_dim_stage_idx, chunk_idx_in_stage = divmod( + chunk_idx, chunks_per_head_dim_stage + ) + return ( + self.swaps_head_dim_stage_tmem_offset(head_dim_stage_idx) + + chunk_idx_in_stage * TMEM_ROW_STRIDE + ) + + def pv_head_dim_stage_tmem_offset(self, head_dim_stage_idx: int) -> int: + """Return the staged PV destination offset in one logical O tile.""" + if self.head_dim_per_stage_kv == 0: + return 0 + if self.use_keeps_mma_ab: + return head_dim_stage_idx * self.head_dim_kv_stage + return self.swaps_head_dim_stage_tmem_offset(head_dim_stage_idx) + + @property + def smem_p_tile_bytes(self) -> int: + """P stored in SMEM for the SwapsMmaAb BMM2 B operand.""" + return self.tile_size_kv * self.tile_size_q * self.q_dtype_bytes + + @property + def tmem_total_cols(self) -> int: + """Sum of TMEM columns used by the kernel: 2× S, 2× softmax stats, + and o_stages× O accumulators (the factor 2 is the two softmax + instances K0/V0 and K1/V1).""" + if self.use_keeps_mma_ab: + if self.num_insts_kv == 1: + return ( + self.tmem_s_cols + + self.tmem_stats_cols + + self.tmem_p_cols + + self.tmem_o_stage_cols * self.o_stages + ) + return ( + 2 * self.tmem_s_cols + + ( + 2 * self.tmem_stats_cols + if self.keeps_separates_tmem_s_and_stats + else 0 + ) + + self.tmem_o_stage_cols * self.o_stages + ) + return ( + 2 * self.tmem_s_cols + + 2 * self.tmem_stats_cols + + self.tmem_o_stage_cols * self.o_stages + ) + + @property + def tmem_alloc_cols(self) -> int: + """TMEM allocation rounded up to a power of two and at least 32, + matching the hardware TMEM allocator's granularity.""" + return max(32, 1 << (self.tmem_total_cols - 1).bit_length()) + + @property + def threads_per_cta(self) -> int: + """CTA thread count implied by the configured task warp layout.""" + return ( + max( + role.preferred_warp_idx + role.num_warps + for role in _active_warp_roles(self) + ) + * 32 + ) + + # ------------------------------------------------------------------ + # Inferred dtype attributes (derived from q_dtype / kv_dtype / out_dtype) + # ------------------------------------------------------------------ + @property + def q_dtype_bytes(self) -> int: + """Byte width of one Q element (fp16/bf16=2, e4m3=1). + + Also the byte width used by anything that feeds the MMA on the Q/S/P + side (softmax stats, P tile, MMA operand descriptors).""" + return 1 if self.q_dtype == Float8E4M3FN else 2 + + @property + def kv_dtype_bytes(self) -> int: + """Byte width of one K/V element (fp16/bf16=2, e4m3=1).""" + return 1 if self.kv_dtype == Float8E4M3FN else 2 + + @property + def o_dtype_bytes(self) -> int: + """Byte width of one O element (fp16/bf16=2, e4m3=1).""" + return 1 if self.out_dtype == Float8E4M3FN else 2 + + @property + def acc_dtype_bytes(self) -> int: + """Byte width of one accumulator element (fp32=4).""" + return 4 if self.acc_dtype == Float32 else 2 + + @property + def use_bf16_qkv(self) -> bool: + """Whether Q/K/V use BF16 storage and MMA inputs.""" + return self.kv_dtype == BFloat16 + + @property + def use_bf16_output(self) -> bool: + """Whether final O is stored as BF16.""" + return self.out_dtype == BFloat16 + + @property + def use_fp8_qkv(self) -> bool: + """fp8 (E4M3) Q/K/V path: switches MMA kind and P-quantization.""" + return self.kv_dtype == Float8E4M3FN + + @property + def use_fp8_output(self) -> bool: + """Whether final O is stored as FP8 E4M3.""" + return self.out_dtype == Float8E4M3FN + + @property + def use_bf16_separate_partial_o(self) -> bool: + """Whether normalized separate-GMEM partial O uses BF16 storage. + + FP16 output keeps FP16 partials to preserve its mantissa/error + envelope. BF16 and FP8 output use BF16 partials for the wider range. + """ + return self.use_bf16_output or self.use_fp8_output + + @property + def supports_reduction_dtypes(self) -> bool: + """Whether split reduction supports the configured input/output dtypes.""" + return ( + self.q_dtype in (Float16, BFloat16) + and self.out_dtype in (Float16, BFloat16) + ) or ( + self.q_dtype == Float8E4M3FN and self.out_dtype in (Float16, Float8E4M3FN) + ) + + # ------------------------------------------------------------------ + # Feature flags + # ------------------------------------------------------------------ + # Enable persistent scheduling: work tiles are fetched from a CLC response + # queue at runtime instead of mapping one tile per CTA in the launch grid. + # Mutually exclusive with split-KV mode. + use_persistent_scheduler: bool = False + # Split-KV: split the K-sequence across several CTAs that produce partial + # O/stats, with a GMEM reduction epilogue. + use_split_kv: bool = False + # Split-KV fanout: number of CTAs that cooperate on one (batch, head_kv) + # when use_split_kv is enabled. The launcher picks this based on SM count. + splits_kv: int = 1 + # Upper bound of splits_kv across all (batch, head_kv) groups, + # used to size the per-launch partial-O / partial-stats / counter + # scratch GMEM buffers. + max_splits_kv: int = 1 + # Paged-KV cache layout: K/V live in fixed-size pages and the kernel + # follows a logical→physical page index table per request. + use_paged_kv: bool = False + # Page size (tokens per page) when use_paged_kv is enabled. Must be one of + # 16 / 32 / 64 / 128 and must divide the 128-token KV tile. + num_tokens_per_page: int = 32 + # Maximum number of pages per (batch, head_kv) — sizes the page index + # table stride. + max_num_pages_per_seq_kv: int = 1 + # Enable sliding-window-causal masking: KV tiles fully outside + # [seq_len_kv − attention_window_size, seq_len_kv) are skipped entirely + # rather than masked element-wise. + use_sliding_window_causal: bool = False + # Window size W: each Q attends to the last W KV tokens. Ignored unless + # use_sliding_window_causal is enabled. + attention_window_size: int = 0 + # Fixed-launch seq_len_kv after sliding-window trimming (constexpr seen + # by the kernel; populated by _configure_static_sliding_window). Falls + # back to the launch seq_len_kv when the bias-TMA path is not used. + static_seq_len_kv: int = 0 + # Number of leading full KV tiles skipped by the sliding window for the + # fixed-launch path. Added to runtime tile indices in resource code. + static_num_skipped_kv_tiles: int = 0 + # Token offset where the sliding window starts (mod tile boundaries), + # used for partial-tile masking on the leading edge of the window. + static_window_start_idx: int = 0 + # When set, the TMA descriptors are pre-biased to point at the first + # in-window token, so the kernel doesn't have to add the skipped-tiles + # offset at runtime. Trades launch-time flexibility for codegen + # simplicity in the fixed-length case. + use_static_sliding_kv_tma_bias: bool = False + # Attention-sinks: add a per-head "sink" exponent to the softmax + # denominator (extra logit that absorbs probability mass). Requires the + # `attention_sinks` tensor argument at launch. + use_attention_sinks: bool = False + # Optional profile and reduction knobs selected by launcher policy or tests. + use_keeps_mma_ab: bool = False + # Nonzero means each K/V stage covers only this many head-dim columns. + # H256 SwapsMmaAb uses 128-column stages to keep TMA and TMEM layouts valid. + head_dim_per_stage_kv: int = 0 + # Ordered softmax barrier mode: 0 disables, 1 auto-enables for supported + # profiles, and 2 forces the barrier path for targeted validation. + ordered_softmax_barrier_mode: int = 0 + # Named barrier slot used by ordered softmax. Slots below 8 are already + # consumed by the main pipeline resources in these schedules. + softmax_order_barrier_id: int = 8 + # Both softmax task groups participate: 8 warps * 32 lanes. + softmax_order_barrier_threads: int = 256 + use_cluster_smem_reduction: bool = False + use_separate_reduction_kernel: bool = False + # Compile-time attention-mask selection. Public APIs normalize the string + # names in MASK_TYPES to the integer constants used by CuTe DSL branches. + mask_type: int = DENSE + + # ------------------------------------------------------------------ + # Derived resource footprints + # ------------------------------------------------------------------ + @property + def smem_q_tile_elements(self) -> int: + """Return Q elements in one staged SMEM tile.""" + return self.smem_q_tile_bytes // self.q_dtype_bytes + + @property + def use_parallel_separate_reduction(self) -> bool: + """Use the production standalone reducer for separate-GMEM modes.""" + return self.use_separate_reduction_kernel + + @property + def use_parallel_separate_reduction_pdl(self) -> bool: + """Order every production standalone reducer through PDL.""" + return self.use_separate_reduction_kernel + + @property + def parallel_reduction_padded_splits(self) -> int: + """Return the power-of-two split capacity of the reducer cluster.""" + if self.max_splits_kv <= 1: + return 1 + return 1 << (self.max_splits_kv - 1).bit_length() + + @property + def use_compact_parallel_reduction(self) -> bool: + """Use one wide CTA when S2-S4 cannot amortize a CTA cluster.""" + return self.use_separate_reduction_kernel and 2 <= self.max_splits_kv <= 4 + + @property + def parallel_reduction_cluster_size(self) -> int: + """Return reducer CTAs for the padded split capacity.""" + if ( + not self.use_separate_reduction_kernel + or self.use_compact_parallel_reduction + ): + return 1 + return { + 8: 1, + 16: 2, + 32: 4, + 64: 8, + 128: 16, + }.get(self.parallel_reduction_padded_splits, 1) + + @property + def parallel_reduction_splits_per_cta(self) -> int: + """Return split slots owned by each reducer CTA.""" + if self.use_compact_parallel_reduction: + return self.max_splits_kv + return ( + self.parallel_reduction_padded_splits + // self.parallel_reduction_cluster_size + ) + + @property + def parallel_reduction_threads_per_cta(self) -> int: + """Return the thread count selected by the reducer schedule.""" + if self.use_compact_parallel_reduction: + return REDUCTION_THREADS_PER_CTA + return PARALLEL_REDUCTION_THREADS_PER_CTA + + @property + def parallel_reduction_bytes_per_slice(self) -> int: + """Return output bytes covered by one independent reducer group.""" + if self.use_compact_parallel_reduction: + return REDUCTION_BYTES_PER_SLICE + return PARALLEL_REDUCTION_BYTES_PER_SLICE + + @property + def smem_kv_tile_elements(self) -> int: + """Return K or V elements in one staged SMEM tile.""" + return self.smem_kv_tile_bytes // self.kv_dtype_bytes + + @property + def num_softmax_scale_groups(self) -> int: + """Return independent max/sum groups tracked by each softmax lane.""" + if self.use_keeps_mma_ab: + return 1 + return max(self.tile_size_q // 4, 1) + + @property + def num_s_regs_per_thread(self) -> int: + """Return all score values owned by each softmax lane.""" + if self.use_keeps_mma_ab: + if self.tile_size_q == 128: + return self.tile_size_kv + return self.tile_size_kv // 2 + return self.num_softmax_scale_groups * 4 + + @property + def softmax_score_fragment_regs(self) -> int: + """Return the maximum score fragment kept live in registers. + + KV256 owns 128 score values per lane but streams them as four native + 32-register LDTM atoms. Other profiles retain their complete score + fragment, so this property is intentionally distinct from + ``num_s_regs_per_thread`` (the total logical ownership). + """ + if self.tile_size_kv == 256: + return 32 + return self.num_s_regs_per_thread + + @property + def num_softmax_score_fragments(self) -> int: + """Return score fragments used to cover one logical KV tile.""" + return self.num_s_regs_per_thread // self.softmax_score_fragment_regs + + @property + def num_packed_p_regs(self) -> int: + """Return packed P registers stored by each softmax producer lane.""" + if self.use_keeps_mma_ab: + values_per_reg = ( + FP8_VALUES_PER_REG if self.use_fp8_qkv else FP16_VALUES_PER_REG + ) + return max(self.num_s_regs_per_thread // values_per_reg, 1) + q_repeats = max(self.tile_size_q // Q_REPETITION_GROUP_HEADS, 1) + regs_per_repeat = ( + FP8_P_PACKED_REGS_PER_Q_REPEAT + if self.use_fp8_qkv + else FP16_P_PACKED_REGS_PER_Q_REPEAT + ) + return regs_per_repeat * q_repeats + + @property + def tmem_p_cols_per_inst(self) -> int: + """Return the TMEM-P columns owned by one K/V instance. + + Two-instance Keeps publishes one packed row per producer lane, with one + packed 32-bit register slot per TMEM column. Therefore its footprint is + ``num_s_regs_per_thread / values_per_reg == num_packed_p_regs``, rather + than a function of the complete logical KV tile width. Q128/KV128 and + Q64/KV256 both own 128 packed 16-bit P values per lane and need 64 + columns per instance. + """ + if self.uses_two_inst_tmem_p: + return self.num_packed_p_regs + return self.tmem_p_cols + + @property + def num_fp8_output_regs(self) -> int: + """Return packed FP8 output registers owned by each correction lane.""" + return max( + (self.tile_size_q * self.headdim) // FP8_OUTPUT_ELEMENTS_PER_REG_GROUP, + 1, + ) + + @property + def num_fp16_output_regs(self) -> int: + """Return packed 16-bit output registers owned by each correction lane.""" + return max( + (self.tile_size_q * self.headdim) // FP16_OUTPUT_ELEMENTS_PER_REG_GROUP, + 1, + ) + + @property + def keeps_output_f32_regs(self) -> int: + """Return FP32 O registers owned by one Keeps correction lane.""" + if self.tile_size_q == 128 or self.tile_size_kv == 256: + return self.headdim + return self.headdim // 2 + + @property + def correction_barrier_threads(self) -> int: + """Return named-barrier participants for the correction task.""" + return self.correction_num_warps * WARP_THREADS + + @property + def keeps_p_smem_vector_elements(self) -> int: + """Return P elements in one aligned 16-byte SMEM store.""" + return 16 // self.q_dtype_bytes + + @property + def static_local_kv_tiles(self) -> int: + """Return static KV tiles assigned to one split CTA.""" + if not self.use_split_kv: + return self.total_kv_tiles + tiles_per_cta_group = self.splits_kv * self.num_insts_kv + num_groups = ( + self.total_kv_tiles + tiles_per_cta_group - 1 + ) // tiles_per_cta_group + return max(self.num_insts_kv, num_groups * self.num_insts_kv) + + @property + def inferred_kv_stages(self) -> int: + """Return the deepest K/V ring that fits the shared-memory budget.""" + q_dtype_bits = 8 if self.q_dtype == Float8E4M3FN else 16 + q_row_bytes = ( + (q_dtype_bits * self.headdim // BITS_PER_BYTE + Q_ROW_ALIGNMENT_BYTES - 1) + // Q_ROW_ALIGNMENT_BYTES + ) * Q_ROW_ALIGNMENT_BYTES + q_tile_kib = q_row_bytes * self.tile_size_q // BYTES_PER_KIB + kv_budget_kib = min( + MAX_KV_STAGE_SMEM_KIB, + TOTAL_SMEM_BUDGET_KIB - q_tile_kib * self.q_stages, + ) + kv_stage_head_dim = self.head_dim_per_stage_kv or self.headdim + kv_tile_bits = q_dtype_bits * self.tile_size_kv * kv_stage_head_dim + return max( + 1, + kv_budget_kib * BYTES_PER_KIB * BITS_PER_BYTE // kv_tile_bits, + ) + + def validate_dtypes(self) -> None: + """Validate decode input, output, and accumulator dtypes.""" + for name, dtype, supported in ( + ("q_dtype", self.q_dtype, SUPPORTED_IO_DTYPES), + ("kv_dtype", self.kv_dtype, SUPPORTED_IO_DTYPES), + ("out_dtype", self.out_dtype, SUPPORTED_IO_DTYPES), + ("acc_dtype", self.acc_dtype, SUPPORTED_ACC_DTYPES), + ): + if dtype not in supported: + raise ValueError(f"Unsupported {name}: {dtype}") + if self.q_dtype != self.kv_dtype: + raise ValueError( + f"q_dtype ({self.q_dtype}) != kv_dtype ({self.kv_dtype}): " + "mixed Q/KV element types are not supported" + ) + + def validate_boolean_fields(self) -> None: + """Require every boolean config field to carry a real Python bool.""" + for name, config_field in self.__dataclass_fields__.items(): + value = getattr(self, name) + if config_field.type in (bool, "bool") and not isinstance(value, bool): + raise TypeError(f"{name} must be a bool, got {type(value).__name__}") + + def validate_paged_kv_staging_config(self) -> None: + """Validate the selected dense or sparse paged-KV staging geometry.""" + + if not self.use_paged_kv: + raise ValueError("paged-KV staging requires use_paged_kv=True") + validate_page_size(self.num_tokens_per_page) + + if self.use_block_sparse: + if self.tile_size_kv not in (128, 256): + raise ValueError( + "paged block-sparse supports only KV128 or KV256 routes" + ) + atom_size = _block_sparse_kv_atom_size(self.kv_block_size) + if atom_size > self.num_tokens_per_page: + raise ValueError( + "paged block-sparse atom size must not exceed page size" + ) + if self.num_tokens_per_page % atom_size != 0: + raise ValueError( + "paged block-sparse page size must be divisible by atom size" + ) + if self.tile_size_kv == 256 and self.num_tokens_per_page not in ( + 64, + 128, + ): + raise ValueError( + "paged block-sparse KV256 routes require page size 64 or 128" + ) + return + + if self.tile_size_kv <= 0: + raise ValueError("paged-KV tile_size_kv must be positive") + if self.tile_size_kv % self.num_tokens_per_page != 0: + raise ValueError( + "paged-KV num_tokens_per_page must divide tile_size_kv exactly" + ) + pages_per_tile = self.tile_size_kv // self.num_tokens_per_page + if pages_per_tile not in (1, 2, 4, 8, 16): + raise ValueError( + "paged-KV staging supports 1, 2, 4, 8, or 16 pages per KV tile" + ) + if self.page_offsets_num_warps != 1: + raise ValueError( + "paged-KV page-offset staging requires exactly one producer warp" + ) + + def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: + """Validate the qualified host profile for block-sparse.""" + if not self.use_block_sparse: + if self.use_parallel_sparse_kv_loads: + raise ValueError( + "sparse execution policy requires block-sparse attention" + ) + return + + if not (self.q_dtype == self.kv_dtype == self.out_dtype): + raise ValueError("block-sparse requires q_dtype == kv_dtype == out_dtype") + if self.q_dtype not in (Float16, BFloat16): + raise ValueError( + "block-sparse supports only matching Float16 or BFloat16 IO" + ) + + kv_block_size = _validate_sparse_kv_block_size(self.kv_block_size) + selected_q_tile = _select_block_sparse_q_tile_size( + q_block_size=self.q_block_size, + heads_q_per_kv=heads_q_per_kv, + kv_block_size=kv_block_size, + ) + if self.use_paged_kv: + self.validate_paged_kv_staging_config() + if not self.groups_tokens_heads_q: + raise ValueError("block-sparse requires groups_tokens_heads_q=True") + if self.headdim != 128: + raise ValueError("block-sparse requires headdim=128") + if self.tile_size_kv not in (128, 256): + raise ValueError("block-sparse requires tile_size_kv in (128, 256)") + if self.tile_size_kv == 256 and not ( + self.tile_size_q == 64 + and kv_block_size % 64 == 0 + and not self.use_parallel_sparse_kv_loads + ): + raise ValueError( + "block-sparse tile_size_kv=256 requires the Q64 16-bit Keeps " + "profile with coarse KV blocks and one load task" + ) + if self.tile_size_q != selected_q_tile: + raise ValueError( + "block-sparse tile_size_q must match its grouped-Q geometry" + ) + uses_fine_q_tile = selected_q_tile < 64 + uses_fine_kv_blocks = kv_block_size < 64 + if uses_fine_q_tile: + if self.use_keeps_mma_ab: + raise ValueError("fine block-sparse Q tiles require SwapsMmaAb") + elif not self.use_keeps_mma_ab: + raise ValueError("coarse block-sparse Q tiles require KeepsMmaAb") + if uses_fine_kv_blocks and not uses_fine_q_tile: + raise ValueError("fine KV blocks require a SwapsMmaAb Q tile") + if self.use_parallel_sparse_kv_loads and ( + self.use_keeps_mma_ab + or kv_block_size not in (8, 16) + or self.num_insts_kv != 2 + ): + raise ValueError( + "parallel sparse K/V loads require two-instance SwapsMmaAb " + "with KV block size 8 or 16" + ) + if self.use_variable_seqlens_q: + raise ValueError("block-sparse does not support variable-Q sequences") + if self.use_sliding_window_causal: + raise ValueError("block-sparse does not support sliding window attention") + if self.use_attention_sinks: + raise ValueError("block-sparse does not support attention sinks") + if ( + self.use_split_kv + or self.use_cluster_smem_reduction + or self.use_separate_reduction_kernel + ): + raise ValueError("block-sparse does not support split-KV reduction") + + def compile_signature(self) -> tuple[tuple[str, object], ...]: + """Key and reconstruct the batch-dynamic callable in the decode cache. + + The public planner can reuse one compiled topology across batch sizes. + Keeping the complete static config in the cache key prevents a callable + compiled for one scheduler, reduction, or tile layout from being reused + by another. + """ + + return tuple( + (config_field.name, getattr(self, config_field.name)) + for config_field in fields(self) + ) + + @property + def uses_q_desc_ref(self) -> bool: + """Whether QK derives Q's descriptor from shared resource state.""" + return self.use_variable_seqlens_q and self.use_persistent_scheduler + + @property + def has_odd_kv_tail(self) -> bool: + """Whether the K/V tile count leaves an unpaired tail instance.""" + return (self.total_kv_tiles % self.num_insts_kv) != 0 + + @property + def uses_uniform_causal_mask(self) -> bool: + """Whether every Q row in a CTA shares one causal right bound.""" + return self.mask_type == CAUSAL and ( + not self.groups_tokens_heads_q + or self.q_tokens_per_cta == 1 + or self.uses_guarded_fixed_q1_grouped_keeps + ) + + @property + def uses_per_row_causal_mask(self) -> bool: + """Whether grouped Q rows require distinct causal right bounds.""" + return ( + self.mask_type == CAUSAL + and self.groups_tokens_heads_q + and self.q_tokens_per_cta > 1 + and not self.uses_guarded_fixed_q1_grouped_keeps + ) + + @property + def _grouped_keeps_profile_key(self) -> _GroupedKeepsProfileKey: + """Return the dtype, shape, and staging key used by Keeps recipe tables.""" + return ( + self.q_dtype, + self.kv_dtype, + self.out_dtype, + self.headdim, + self.head_dim_per_stage_kv, + self.num_insts_kv, + self.o_stages, + ) + + @property + def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: + """Whether inactive grouped rows are excluded solely at output stores. + + The validated paged FP8 Keeps profiles pack two or four structural Q + tokens into TileQ64/128 while the public decode problem has exactly one + logical token. QK/PV rows are independent and every direct, split + scratch, and final-reduction store already checks row validity, so the + inactive score rows do not need per-KV-tile suppression. + """ + profile = self._grouped_keeps_profile_key + return ( + self.use_keeps_mma_ab + and self.groups_tokens_heads_q + and self.max_seq_len_q == 1 + and not self.use_variable_seqlens_q + and self.q_manual_padding_rows == 0 + and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES + and self.supports_grouped_keeps + # The staged TileQ64 one-instance TMEM-P schedule keeps per-row + # score masking because its generated code is sensitive to that + # control-flow shape. + and not (self.uses_staged_one_inst_tmem_p and self.tile_size_q == 64) + ) + + @property + def uses_guarded_grouped_keeps_output_rows(self) -> bool: + """Whether inactive Keeps rows can be discarded only at publication. + + Keeps MMA, softmax, and PV rows are independent. For a fixed grouped-Q + launch, structural padding and a partial final token group therefore + cannot affect a valid row; direct output, split scratch, and reduction + publication already guard row validity. Keep the staged one-instance + TileQ64/D256 exception on its per-row score-mask path because its + generated schedule is sensitive to that control-flow shape. + """ + profile = self._grouped_keeps_profile_key + return ( + self.use_keeps_mma_ab + and self.groups_tokens_heads_q + and not self.use_variable_seqlens_q + and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES + and self.supports_grouped_keeps + and not (self.uses_staged_one_inst_tmem_p and self.tile_size_q == 64) + ) + + @property + def has_static_dense_full_kv_tiles(self) -> bool: + """Whether static dense KV avoids masking and runtime tile remapping.""" + return ( + not self.use_split_kv + and not self.use_block_sparse + and self.mask_type == DENSE + and not self.use_sliding_window_causal + and self.static_seq_len_kv != 0 + and (self.static_seq_len_kv % self.tile_size_kv) == 0 + ) + + @property + def uses_ordered_softmax_barrier(self) -> bool: + """Whether this profile selects the ordered P0/P1 softmax barrier.""" + if self.tile_size_kv == 256: + # KV256 uses independent four-stage P-fragment pipelines. Ordering + # the two softmax groups would serialize fragment production and + # defeat the intended P/PV overlap. + return False + if self.ordered_softmax_barrier_mode == 2: + return True + return self.ordered_softmax_barrier_mode == 1 and ( + self.headdim == 128 and self.tile_size_q in (32, 64, 128) + ) + + @property + def ordered_softmax_early_release(self) -> bool: + """Whether the P0/P1 baton is released at TMEM store issue. + + For two-instance TMEM-P the partner softmax group's publication + touches only its own registers and TMEM region, so it does not need + this group's store drain, async fence, or pipeline commit. Releasing + at TMEM store issue overlaps that tail with the partner's wakeup while the + exp2 phases stay serialized on the shared MUFU pipes. + """ + return self.uses_ordered_softmax_barrier and self.uses_two_inst_tmem_p + + @property + def resolved_softmax_order_barrier_threads(self) -> int: + """Return the participant count for ordered softmax barriers.""" + heads_q_per_kv = self.heads_q_per_kv or self.tile_size_q + if self.use_keeps_mma_ab: + return (self.softmax0_num_warps + self.softmax1_num_warps) * WARP_THREADS + if ( + self.headdim == 128 + and self.tile_size_q == 32 + and self.tile_size_kv == 128 + and heads_q_per_kv == 32 + and not self.groups_tokens_heads_q + ): + return 128 + return self.softmax_order_barrier_threads + + @property + def uses_nontrivial_grouped_q_layout(self) -> bool: + """Whether grouped Q differs from one complete token per CTA.""" + return self.groups_tokens_heads_q and ( + self.q_tokens_per_cta > 1 or self.q_manual_padding_rows > 0 + ) + + @property + def q_tiles_are_full(self) -> bool: + """Whether every launched Q CTA owns all of its MMA rows.""" + if self.use_variable_seqlens_q: + return False + if self.uses_nontrivial_grouped_q_layout: + return ( + self.q_manual_padding_rows == 0 + and self.max_seq_len_q % self.q_tokens_per_cta == 0 + ) + if self.max_seq_len_q > 1: + return self.heads_q_per_kv % self.tile_size_q == 0 + return self.heads_q_per_kv == self.tile_size_q + + @property + def q_tiles_need_row_mask(self) -> bool: + """Whether softmax must suppress inactive Q rows.""" + if self.use_variable_seqlens_q: + return True + if self.uses_nontrivial_grouped_q_layout or self.max_seq_len_q > 1: + return not self.q_tiles_are_full + return False + + @property + def q_score_rows_need_mask(self) -> bool: + """Whether inactive Q rows must be suppressed in every score tile.""" + return ( + self.q_tiles_need_row_mask + and not self.uses_guarded_grouped_keeps_output_rows + ) + + @property + def uses_q_cta_sliding_union(self) -> bool: + """Whether the causal/window KV union depends on the logical Q CTA.""" + return ( + self.mask_type == CAUSAL + and self.max_seq_len_q > 1 + and not self.has_single_q_cta + ) + + @property + def uses_runtime_q_kv_union(self) -> bool: + """Whether task/resource KV geometry must retain runtime Q metadata. + + Multiple causal Q CTAs have distinct right bounds. A multi-token + sliding-window launch also retains the runtime path even when it fits + in one Q CTA, because the existing static window-prefix metadata is + specialized only for SQ1. + """ + return self.uses_q_cta_sliding_union or ( + self.use_sliding_window_causal and self.max_seq_len_q > 1 + ) + + @property + def uses_tmem_p(self) -> bool: + """Whether Keeps P is materialized in TMEM for BMM2.""" + return self.uses_staged_one_inst_tmem_p or self.uses_two_inst_tmem_p + + @property + def uses_staged_one_inst_tmem_p(self) -> bool: + """Whether P uses the double-buffered D256 TMEM overlay.""" + return ( + self.use_keeps_mma_ab + and self.headdim == 256 + and self.head_dim_per_stage_kv == 128 + and self.num_insts_kv == 1 + and self.o_stages == 1 + ) + + @property + def uses_two_inst_tmem_p(self) -> bool: + """Whether a two-instance Keeps profile uses the TMEM-P overlay. + + Q128/KV128 and sparse Q64/KV128 publish a complete packed-P row per + pipeline token. Q64/KV256 uses the same S-to-P aliasing contract but + streams four independently ready K32 fragments. Dense Q64/KV128 keeps + the base kernel's faster SMEM-P cadence. + """ + # Two-instance Keeps keeps stats outside S, so both static and persistent + # work tiles can overlay P on the consumed S instance. The split K/V + # schedule preserves same-instance PV -> QK order, while Softmax delays + # S release through TMEM store completion and the P-pipeline commit. + return ( + self.use_keeps_mma_ab + and ( + (self.tile_size_q == 128 and self.tile_size_kv == 128) + or (self.tile_size_q == 64 and self.tile_size_kv == 256) + or ( + self.use_block_sparse + and self.tile_size_q == 64 + and self.tile_size_kv == 128 + ) + ) + and self.head_dim_per_stage_kv == 0 + and self.num_insts_kv == 2 + and self.o_stages == 2 + and self.tmem_total_cols <= 512 + ) + + @property + def streams_tmem_p_fragments(self) -> bool: + """Whether P is published as independently ready TMEM fragments.""" + return self.uses_two_inst_tmem_p and self.num_softmax_score_fragments > 1 + + @property + def matches_kv256_task_topology(self) -> bool: + """Whether task roles match KV256's validated 16-warp layout.""" + return all( + getattr(self, field) == expected + for field, expected in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items() + ) + + @property + def uses_rotating_kv256_exchange(self) -> bool: + """Whether this profile selects KV-ring scratch for correction. + + Persistent direct output can overlap the next work tile's first two + K loads with correction by placing its exchange in the third, drained + KV stage. Split-KV and attention sinks retain the fixed exchange because + their tail storage and lifetime differ from direct output. + """ + selects_persistent_kv256 = ( + self.streams_tmem_p_fragments + and self.tile_size_q == 64 + and self.tile_size_kv == 256 + and self.use_persistent_scheduler + ) + if not selects_persistent_kv256: + return False + + has_rotating_kv_ring = ( + self.num_head_dim_stages_kv == 1 + and self.kv_stages == KV_TILE_256_SHARED_FIFO_STAGES + and self.load_num_warps == 1 + ) + has_direct_output_lifetime = not (self.use_split_kv or self.use_attention_sinks) + return has_rotating_kv_ring and has_direct_output_lifetime + + @property + def keeps_separates_tmem_s_and_stats(self) -> bool: + """Whether two-instance Keeps has room for standalone stats tiles.""" + if not ( + self.use_keeps_mma_ab + and self.use_fp8_qkv + and self.tile_size_kv == 128 + and self.head_dim_per_stage_kv == 0 + and self.num_insts_kv == 2 + and self.o_stages == 2 + ): + return False + return ( + 2 * self.tmem_s_cols + + 2 * self.tmem_stats_cols + + self.tmem_o_stage_cols * self.o_stages + <= 512 + ) + + @property + def keeps_stats_via_smem(self) -> bool: + """Whether Keeps softmax->correction stats travel through SMEM. + + When two-instance Keeps cannot give the stats payload standalone + TMEM columns (S/stats/O exceed the 512-column budget), the stats + slot aliases S and a stats-done credit pipeline must gate every QK + re-issue on correction's TMEM stats read. Routing the small per-row + payload through an SMEM ring removes that MMA-side serialization; + the existing softmax-local pipelines already order the handoff. + """ + return self.use_keeps_mma_ab and not self.keeps_separates_tmem_s_and_stats + + @property + def keeps_loop_correction_chunk_regs(self) -> int: + """Return FP32 registers corrected by one Keeps TMEM pair.""" + if self.tile_size_q == 64: + return 32 + return 32 if self.tile_size_q == 128 and self.headdim >= 256 else 8 + + @property + def keeps_loop_correction_stage_layout(self) -> tuple[tuple[int, int, int], ...]: + """Return ``(TMEM offset, half split, chunk count)`` for each O slice.""" + stage_cols = self.head_dim_kv_stage + lane_regs_per_stage = ( + stage_cols + if self.tile_size_q == 128 or self.tile_size_kv == 256 + else stage_cols // 2 + ) + chunk_regs = self.keeps_loop_correction_chunk_regs + assert stage_cols % 2 == 0 + assert lane_regs_per_stage % chunk_regs == 0 + return tuple( + ( + stage_idx * stage_cols, + stage_cols // 2, + lane_regs_per_stage // chunk_regs, + ) + for stage_idx in range(self.num_head_dim_stages_kv) + ) + + @property + def fp8_copy_can_use_full_tile_fast_path(self) -> bool: + """Whether every correction wave owns only in-bounds FP8 bytes.""" + correction_copy_bytes = self.correction_num_warps * WARP_THREADS * 16 + tile_bytes = self.tile_size_q * self.headdim + return self.q_tiles_are_full and tile_bytes % correction_copy_bytes == 0 + + @property + def can_use_cluster_smem_reduction(self) -> bool: + """Whether the configured split profile is eligible for cluster reduction.""" + ungrouped_q_layout = ( + not self.groups_tokens_heads_q + and not self.use_variable_seqlens_q + and self.max_seq_len_q == 1 + ) + grouped_q_layout = ( + self.groups_tokens_heads_q + and self.heads_q_per_kv > 0 + and self.tile_size_q >= self.heads_q_per_kv + ) + return ( + self.use_split_kv + and not self.use_separate_reduction_kernel + and not self.use_persistent_scheduler + and not self.use_keeps_mma_ab + and self.headdim in (64, 128, 256) + and self.tile_size_q in (8, 16, 32) + and (ungrouped_q_layout or grouped_q_layout) + and self.max_splits_kv >= self.splits_kv >= 2 + ) + + @property + def supports_cluster_smem_reduction(self) -> bool: + """Whether cluster reduction is both selected and eligible.""" + return self.use_cluster_smem_reduction and self.can_use_cluster_smem_reduction + + @property + def split_reduction_slice_bytes(self) -> int: + """Return the minimum byte range assigned to one reducer CTA.""" + return self.correction_barrier_threads * SPLIT_REDUCTION_VECTOR_BYTES_PER_THREAD + + @property + def split_reduction_rows_per_slice(self) -> int: + """Return complete partial-O rows covered by one reducer slice.""" + row_bytes = self.headdim * PARTIAL_O_ELEMENT_BYTES + return max(self.split_reduction_slice_bytes // row_bytes, 1) + + @property + def split_reduction_slices_per_cta(self) -> int: + """Return contiguous reducer slices assigned to one owner CTA.""" + rows_per_slice = self.split_reduction_rows_per_slice + num_slices = (self.tile_size_q + rows_per_slice - 1) // rows_per_slice + return max((num_slices + self.splits_kv - 1) // self.splits_kv, 1) + + @property + def cluster_reduction_rows_per_cta(self) -> int: + """Return the slice-aligned row capacity of one split-reduction owner.""" + return self.split_reduction_slices_per_cta * self.split_reduction_rows_per_slice + + @property + def cluster_reduction_num_owner_ctas(self) -> int: + """Return split CTAs that own at least one physical reducer slice.""" + rows_per_slice = self.split_reduction_rows_per_slice + num_slices = (self.tile_size_q + rows_per_slice - 1) // rows_per_slice + return min( + (num_slices + self.split_reduction_slices_per_cta - 1) + // self.split_reduction_slices_per_cta, + self.splits_kv, + ) + + @property + def cluster_max_runtime_partial_rows(self) -> int: + """Maximum slice-aligned ``split x owner-row`` records at runtime.""" + if not self.supports_cluster_smem_reduction: + return self.max_splits_kv * self.cluster_reduction_rows_per_cta + rows_per_slice = self.split_reduction_rows_per_slice + num_slices = (self.tile_size_q + rows_per_slice - 1) // rows_per_slice + return max( + splits * ((num_slices + splits - 1) // splits) * rows_per_slice + for splits in range(1, self.splits_kv + 1) + ) + + @property + def correction_sum_scratch_entries(self) -> int: + """Return correction denominator scratch entries, or zero if unused.""" + return 0 if self.use_keeps_mma_ab else 4 * self.tile_size_q + + @property + def cluster_transaction_bytes(self) -> int: + """Return the byte count expected by each cluster owner barrier.""" + row_bytes = ( + self.headdim * PARTIAL_O_ELEMENT_BYTES + + PARTIAL_STATS_VALUES_PER_ROW * FP32_BYTES + ) + return self.cluster_max_runtime_partial_rows * row_bytes + + @property + def max_runtime_row_split_segments(self) -> int: + """Return the maximum reducer slices assigned to one runtime owner.""" + # Runtime contraction can leave one active owner responsible for the + # complete Q tile. + rows_per_slice = self.split_reduction_rows_per_slice + return max((self.tile_size_q + rows_per_slice - 1) // rows_per_slice, 1) + + @property + def supports_grouped_keeps(self) -> bool: + """Whether a validated config uses a qualified grouped-Keeps recipe.""" + if self.tile_size_kv == 256: + # KV256 reuses the common fixed/packed-Q, page-table, masking, + # persistent scheduler, attention-sink, and GMEM split-publisher + # semantics. Keeps has no cluster-SMEM publisher at either KV tile + # size. + if ( + not self.use_keeps_mma_ab + or not self.groups_tokens_heads_q + or self.tile_size_q != 64 + or self.headdim != 128 + or self.q_dtype not in (Float16, BFloat16) + or not (self.q_dtype == self.kv_dtype == self.out_dtype) + or self.use_cluster_smem_reduction + or not self.matches_kv256_task_topology + ): + return False + direct = not (self.use_split_kv or self.use_separate_reduction_kernel) + if direct: + return True + if not ( + self.use_split_kv + and self.splits_kv > 1 + and self.max_splits_kv >= self.splits_kv + ): + return False + if self.use_separate_reduction_kernel and self.use_variable_seqlens_q: + return self.mask_type == CAUSAL + return self.mask_type == DENSE + if ( + not self.use_keeps_mma_ab + or not self.groups_tokens_heads_q + or self.tile_size_kv != 128 + or self.tile_size_q not in (64, 128) + or self.use_cluster_smem_reduction + ): + return False + + profile = self._grouped_keeps_profile_key + + # Keep block-sparse qualification separate from the dense/paged + # profile matrix below. Its structural, masking, and reduction + # constraints are validated separately; both scheduler modes use the + # same qualified recipe keys. + if self.use_block_sparse: + return profile in _BLOCK_SPARSE_GROUPED_KEEPS_PROFILES + + direct = not (self.use_split_kv or self.use_separate_reduction_kernel) + + if profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES: + fixed_q1_ratio32 = self.max_seq_len_q == 1 and self.heads_q_per_kv == 32 + fixed_grouped_q = self.max_seq_len_q > 1 + return ( + (fixed_q1_ratio32 or fixed_grouped_q) + and not self.use_variable_seqlens_q + and self.use_paged_kv + and self.num_tokens_per_page == 32 + and self.mask_type == CAUSAL + and not any( + ( + self.use_cluster_smem_reduction, + self.use_sliding_window_causal, + self.use_attention_sinks, + ) + ) + and ( + direct + or ( + self.use_split_kv + and self.splits_kv > 1 + and self.max_splits_kv >= self.splits_kv + ) + ) + ) + + if profile in _GROUPED_KEEPS_STATIC_ONLY_PROFILES: + return ( + self.tile_size_q == 64 + and direct + and self.mask_type == DENSE + and not any( + ( + self.use_variable_seqlens_q, + self.use_persistent_scheduler, + self.use_paged_kv, + self.use_sliding_window_causal, + self.use_attention_sinks, + ) + ) + ) + if profile != _GROUPED_KEEPS_MAIN_PROFILE: + return False + + if self.tile_size_q == 128: + return direct and not any( + ( + self.use_persistent_scheduler, + self.use_paged_kv, + self.use_sliding_window_causal, + self.use_attention_sinks, + ) + ) + + if self.use_paged_kv or self.use_sliding_window_causal: + return ( + direct + and self.mask_type == CAUSAL + and not self.use_persistent_scheduler + and not self.use_attention_sinks + and not (self.use_paged_kv and self.use_sliding_window_causal) + and ( + not self.use_sliding_window_causal or self.attention_window_size > 0 + ) + ) + + if self.use_persistent_scheduler: + return ( + direct + and not self.use_attention_sinks + and (self.use_variable_seqlens_q or self.mask_type == DENSE) + ) + if direct: + return True + + if not ( + self.use_split_kv + and self.splits_kv > 1 + and self.max_splits_kv >= self.splits_kv + ): + return False + + if self.use_separate_reduction_kernel and self.use_variable_seqlens_q: + return not self.use_attention_sinks or self.mask_type == CAUSAL + return self.mask_type == DENSE and not self.use_attention_sinks + + +SUPPORTED_IO_DTYPES = {Float16, BFloat16, Float8E4M3FN} +SUPPORTED_ACC_DTYPES = {Float32} + + +def _decode_config_items(source: object): + """Return candidate config key/value pairs from a mapping or args-like object.""" + if isinstance(source, Mapping): + return source.items() + namespace = getattr(source, "__dict__", None) + if namespace is not None: + return namespace.items() + field_names = FmhaDecodeConfig.__dataclass_fields__ + return ( + (name, getattr(source, name)) for name in field_names if hasattr(source, name) + ) + + +def _iter_config_sources(source: object): + """Yield config sources in precedence order for direct config mutation.""" + if source is None: + return + if isinstance(source, (tuple, list)): + for item in source: + yield from _iter_config_sources(item) + return + yield source + + +def _apply_config_source(cfg: FmhaDecodeConfig, source: object) -> set[str]: + """Apply explicit, correctly typed config fields and return names touched.""" + field_names = FmhaDecodeConfig.__dataclass_fields__ + explicit_fields: set[str] = set() + for source_item in _iter_config_sources(source): + for key, value in _decode_config_items(source_item): + if key == "use_causal_spec_decoding": + raise ValueError( + "use_causal_spec_decoding was removed; use " + "mask_type='dense' or mask_type='causal' instead" + ) + if key == "single_token_q_per_cta": + raise ValueError( + "single_token_q_per_cta was removed; select the Q layout " + "with groups_tokens_heads_q" + ) + if key == "mask_type": + if value is not None: + explicit_fields.add(key) + # Normalize mask strings after all sources have been inspected + # and the sliding-window request is known. + continue + if value is None or key == "headdim" or key not in field_names: + continue + if key in ("splits_kv", "max_splits_kv") and int(value) <= 0: + continue + setattr(cfg, key, value) + explicit_fields.add(key) + cfg.validate_boolean_fields() + return explicit_fields + + +def _resolve_explicit_split_controls( + cfg: FmhaDecodeConfig, + *, + explicit_fields: set[str], + splits_kv: int, + max_splits_kv: int | None, +) -> tuple[int, int | None]: + """Merge public and config-source split controls without losing intent.""" + if "splits_kv" in explicit_fields: + if splits_kv > 0 and splits_kv != cfg.splits_kv: + raise ValueError( + "conflicting splits_kv selections: public API requested " + f"{splits_kv}, while config overrides requested {cfg.splits_kv}" + ) + if splits_kv <= 0: + splits_kv = cfg.splits_kv + if "max_splits_kv" in explicit_fields: + if ( + max_splits_kv is not None + and max_splits_kv > 0 + and max_splits_kv != cfg.max_splits_kv + ): + raise ValueError( + "conflicting max_splits_kv selections: public API requested " + f"{max_splits_kv}, while config overrides requested " + f"{cfg.max_splits_kv}" + ) + if max_splits_kv is None or max_splits_kv <= 0: + max_splits_kv = cfg.max_splits_kv + return splits_kv, max_splits_kv + + +def _mask_type_from_config_source(source: object) -> str | int | None: + """Return the last non-None mask selection from config sources.""" + selected: str | int | None = None + for source_item in _iter_config_sources(source): + for key, value in _decode_config_items(source_item): + if key == "mask_type" and value is not None: + selected = value + return selected + + +def groups_tokens_heads_q_from_config_source(source: object) -> bool | None: + """Return the last explicit Q-grouping selection from config sources.""" + selected = None + for source_item in _iter_config_sources(source): + for key, value in _decode_config_items(source_item): + if key != "groups_tokens_heads_q" or value is None: + continue + if not isinstance(value, bool): + raise TypeError( + f"groups_tokens_heads_q must be a bool, got {type(value).__name__}" + ) + selected = value + return selected + + +def _apply_mask_type_config( + cfg: FmhaDecodeConfig, + *, + source: object, + mask_type: str | int | None, + sliding_window_causal: bool, + explicit_fields: set[str], +) -> None: + """Resolve public and config-source mask selections into a constexpr id.""" + source_mask_type = _mask_type_from_config_source(source) + if mask_type is not None: + explicit_fields.add("mask_type") + if mask_type is not None and source_mask_type is not None: + public_mask_type = normalize_mask_type( + mask_type, sliding_window_causal=sliding_window_causal + ) + config_mask_type = normalize_mask_type( + source_mask_type, sliding_window_causal=sliding_window_causal + ) + if public_mask_type != config_mask_type: + raise ValueError( + "conflicting mask_type selections: public API requested " + f"{mask_type_name(public_mask_type)!r}, while config overrides " + f"requested {mask_type_name(config_mask_type)!r}" + ) + cfg.mask_type = public_mask_type + return + cfg.mask_type = normalize_mask_type( + mask_type if mask_type is not None else source_mask_type, + sliding_window_causal=sliding_window_causal, + ) + + +def _set_if_implicit( + cfg: FmhaDecodeConfig, + field_name: str, + value: ConfigValue, + explicit_fields: set[str], +) -> None: + """Set a derived default only when the caller did not provide the field.""" + if field_name not in explicit_fields: + setattr(cfg, field_name, value) + + +def _finalize_static_decode_config( + cfg: FmhaDecodeConfig, + explicit_fields: set[str], +) -> None: + """Fill dtype-dependent, profile-dependent, and SMEM-derived defaults.""" + cfg.validate_boolean_fields() + cfg.validate_dtypes() + + use_keeps_mma_ab = cfg.use_keeps_mma_ab + if not use_keeps_mma_ab and cfg.headdim > 128: + _set_if_implicit(cfg, "head_dim_per_stage_kv", 128, explicit_fields) + _set_if_implicit(cfg, "num_insts_kv", 2, explicit_fields) + + if use_keeps_mma_ab: + tile_size_q = cfg.tile_size_q if "tile_size_q" in explicit_fields else 64 + tile_size_kv = cfg.tile_size_kv if "tile_size_kv" in explicit_fields else 128 + if cfg.headdim > 128: + _set_if_implicit(cfg, "num_insts_kv", 1, explicit_fields) + _set_if_implicit(cfg, "head_dim_per_stage_kv", 128, explicit_fields) + _set_if_implicit(cfg, "o_stages", 1, explicit_fields) + _set_if_implicit(cfg, "tile_size_q", tile_size_q, explicit_fields) + if cfg.tile_size_kv == 256: + # Materialize the selected profile without overriding caller + # tuning. The common validator below decides whether the resulting + # effective configuration is supported by the kernel. + for field_name, value in _KV_TILE_256_PHYSICAL_DEFAULTS.items(): + _set_if_implicit(cfg, field_name, value, explicit_fields) + for field_name, value in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items(): + _set_if_implicit(cfg, field_name, value, explicit_fields) + else: + _set_if_implicit(cfg, "tmem_s_cols", tile_size_kv, explicit_fields) + _set_if_implicit( + cfg, + "tmem_p_cols", + tile_size_kv // 2, + explicit_fields, + ) + _set_if_implicit(cfg, "tmem_o_cols", cfg.headdim, explicit_fields) + _set_if_implicit(cfg, "mma_tile_m_bmm1", tile_size_q, explicit_fields) + _set_if_implicit(cfg, "mma_tile_n_bmm1", tile_size_kv, explicit_fields) + _set_if_implicit(cfg, "mma_tile_m_bmm2", tile_size_q, explicit_fields) + _set_if_implicit( + cfg, + "mma_tile_n_bmm2", + cfg.head_dim_per_stage_kv or cfg.headdim, + explicit_fields, + ) + if cfg.num_insts_kv == 1: + # One-inst static profiles use a compact 12-warp layout. Persistent + # profiles keep MMA/scheduler/page-or-padding/load together in WG2 + # and add a work-queue-aware padding WG3 to preserve the existing + # 16-warp CLC and CTA-barrier contract. + _set_if_implicit(cfg, "correction_warp_idx", 4, explicit_fields) + _set_if_implicit(cfg, "mma_warp_idx", 8, explicit_fields) + _set_if_implicit(cfg, "page_offsets_warp_idx", 9, explicit_fields) + _set_if_implicit(cfg, "load_warp_idx", 11, explicit_fields) + _set_if_implicit(cfg, "scheduler_warp_idx", 9, explicit_fields) + _set_if_implicit(cfg, "clc_load_warp_idx", 11, explicit_fields) + if cfg.tile_size_q == 128 and cfg.tile_size_kv != 256: + _set_if_implicit(cfg, "q_stages", 1, explicit_fields) + _set_if_implicit(cfg, "ordered_softmax_barrier_mode", 1, explicit_fields) + + if cfg.tile_size_kv != 256 and "kv_stages" not in explicit_fields: + cfg.kv_stages = cfg.inferred_kv_stages + + # Split-KV mode forbids persistent scheduling. Canonicalize here so + # downstream consumers can gate on use_persistent_scheduler alone without + # having to repeat the (... and not use_split_kv) check. + if cfg.use_split_kv: + cfg.use_persistent_scheduler = False + + +def _validate_kv256_static_config(cfg: FmhaDecodeConfig) -> None: + """Validate the effective KV256 profile after implicit defaults are filled.""" + if cfg.tile_size_kv != 256: + return + + def _require_python_int(field_name: str) -> int: + value = getattr(cfg, field_name) + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError( + f"KV256 {field_name} must be a Python integer, " + f"got {type(value).__name__}" + ) + return value + + for field_name, expected in _KV_TILE_256_PHYSICAL_DEFAULTS.items(): + actual = _require_python_int(field_name) + if field_name in _KV_TILE_256_TUNABLE_FIELDS: + continue + if actual != expected: + raise ValueError(f"KV256 requires {field_name}={expected}, got {actual}") + for field_name, expected in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items(): + actual = _require_python_int(field_name) + if actual != expected: + raise ValueError(f"KV256 requires {field_name}={expected}, got {actual}") + + if cfg.q_stages <= 0 or cfg.kv_stages <= 0: + raise ValueError( + "KV256 shared-memory pipeline requires positive q_stages and " + f"kv_stages, got q_stages={cfg.q_stages}, " + f"kv_stages={cfg.kv_stages}" + ) + pipeline_smem_bytes = ( + cfg.q_stages * cfg.smem_q_tile_bytes + cfg.kv_stages * cfg.smem_kv_tile_bytes + ) + pipeline_smem_budget_bytes = TOTAL_SMEM_BUDGET_KIB * BYTES_PER_KIB + if pipeline_smem_bytes > pipeline_smem_budget_bytes: + raise ValueError( + "KV256 shared-memory pipeline exceeds the static SMEM budget: " + f"q_stages={cfg.q_stages}, kv_stages={cfg.kv_stages} require " + f"{pipeline_smem_bytes} bytes, limit is " + f"{pipeline_smem_budget_bytes} bytes" + ) + if ( + cfg.use_persistent_scheduler + and not cfg.use_split_kv + and not cfg.use_attention_sinks + and cfg.kv_stages != KV_TILE_256_SHARED_FIFO_STAGES + ): + raise ValueError( + "persistent KV256 requires kv_stages=" + f"{KV_TILE_256_SHARED_FIFO_STAGES} for the rotating shared-KV " + f"exchange, got {cfg.kv_stages}" + ) + if not cfg.supports_grouped_keeps: + raise ValueError( + "KV256 currently supports only the qualified Q64 FP16/BF16/D128 " + "grouped Keeps profile" + ) + + +@dataclass(frozen=True) +class _WarpRole: + """One active task's configured warp placement.""" + + name: str + index_field: str + preferred_warp_idx: int + num_warps: int + is_padding: bool = False + + +def _append_padding_warp_roles(cfg: FmhaDecodeConfig, roles: list[_WarpRole]) -> None: + """Append padding warps to the tail of each warp group.""" + num_warps_per_wg = [0] * MAX_WARP_GROUPS + total_num_wgs = 0 + for role in roles: + if role.num_warps <= 0: + raise ValueError(f"{role.name} must use at least one warp") + preferred_wg = role.preferred_warp_idx // 4 + preferred_wg_end = (role.preferred_warp_idx + role.num_warps - 1) // 4 + if preferred_wg != preferred_wg_end: + raise ValueError( + f"{role.name} spans warp groups {preferred_wg} and {preferred_wg_end}; " + "each task role must fit within one warp group" + ) + if preferred_wg < 0 or preferred_wg >= MAX_WARP_GROUPS: + raise ValueError( + f"{role.name} is assigned to unsupported warp group {preferred_wg}" + ) + num_warps_per_wg[preferred_wg] += role.num_warps + total_num_wgs = max(total_num_wgs, preferred_wg + 1) + for wg_idx, num_warps in enumerate(num_warps_per_wg): + if num_warps > 4: + raise ValueError( + f"warp group {wg_idx} has {num_warps} active warps; at most 4 are allowed" + ) + + if cfg.use_persistent_scheduler: + total_num_wgs = MAX_WARP_GROUPS + + for wg_idx, active_warps in enumerate(num_warps_per_wg): + padding_index_field = f"wg{wg_idx}_padding_warp_idx" + padding_count_field = f"wg{wg_idx}_padding_num_warps" + padding_warp_idx = wg_idx * 4 + active_warps + padding_num_warps = 4 - active_warps if wg_idx < total_num_wgs else 0 + setattr(cfg, padding_index_field, padding_warp_idx) + setattr(cfg, padding_count_field, padding_num_warps) + if padding_num_warps > 0: + roles.append( + _WarpRole( + f"wg{wg_idx}_padding", + padding_index_field, + padding_warp_idx, + padding_num_warps, + is_padding=True, + ) + ) + + +def _active_warp_roles(cfg: FmhaDecodeConfig) -> list[_WarpRole]: + """Return the warp roles instantiated by the current kernel profile.""" + roles = [ + _WarpRole( + "softmax0", + "softmax0_warp_idx", + cfg.softmax0_warp_idx, + cfg.softmax0_num_warps, + ), + _WarpRole( + "correction", + "correction_warp_idx", + cfg.correction_warp_idx, + cfg.correction_num_warps, + ), + _WarpRole("mma", "mma_warp_idx", cfg.mma_warp_idx, cfg.mma_num_warps), + ] + if cfg.num_insts_kv != 1: + roles.append( + _WarpRole( + "softmax1", + "softmax1_warp_idx", + cfg.softmax1_warp_idx, + cfg.softmax1_num_warps, + ) + ) + + if cfg.use_persistent_scheduler: + roles.extend( + ( + _WarpRole( + "scheduler", + "scheduler_warp_idx", + cfg.scheduler_warp_idx, + cfg.scheduler_num_warps, + ), + _WarpRole( + "load", + "clc_load_warp_idx", + cfg.clc_load_warp_idx, + cfg.load_num_warps, + ), + ) + ) + else: + roles.append( + _WarpRole( + "load", + "load_warp_idx", + cfg.load_warp_idx, + cfg.load_num_warps, + ) + ) + + if cfg.use_paged_kv and not cfg.use_block_sparse: + roles.append( + _WarpRole( + "page_offsets", + "page_offsets_warp_idx", + cfg.page_offsets_warp_idx, + cfg.page_offsets_num_warps, + ) + ) + _append_padding_warp_roles(cfg, roles) + return roles + + +def _finalize_warp_roles(cfg: FmhaDecodeConfig) -> None: + """Compact active warp indices and add paddings.""" + warp_roles = sorted( + _active_warp_roles(cfg), + key=lambda role: ( + role.preferred_warp_idx // 4, + role.is_padding, + role.preferred_warp_idx, + ), + ) + + next_warp_idx = 0 + for role in warp_roles: + setattr(cfg, role.index_field, next_warp_idx) + next_warp_idx += role.num_warps + + +def _make_static_decode_config( + headdim: int = 128, + args: object | None = None, + *, + mask_type: str | int | None = None, + sliding_window_causal: bool = False, +) -> FmhaDecodeConfig: + """Build a static FmhaDecodeConfig by mutating a default config object. + + ``args`` may be a parser namespace, harness dictionary, fully specified + config object, or a tuple/list of those sources. Only ``FmhaDecodeConfig`` + fields are read from it. ``None`` values mean "keep the default"; mappings + are the appropriate source when omitted fields should remain implicit. + """ + cfg = FmhaDecodeConfig(headdim=headdim) + explicit_fields = _apply_config_source(cfg, args) + _apply_mask_type_config( + cfg, + source=args, + mask_type=mask_type, + sliding_window_causal=(sliding_window_causal or cfg.use_sliding_window_causal), + explicit_fields=explicit_fields, + ) + _finalize_static_decode_config(cfg, explicit_fields) + _validate_kv256_static_config(cfg) + _finalize_warp_roles(cfg) + return cfg + + +def _swaps_tile_fields_for_heads( + num_heads_q: int, + num_heads_kv: int, + *, + groups_tokens_heads_q: bool, +) -> dict[str, int]: + """Return SwapsMmaAb tile metadata for a GQA head ratio.""" + h_r = num_heads_q // num_heads_kv + if groups_tokens_heads_q: + tile_size_q = next((tile_q for tile_q in (8, 16, 32) if h_r <= tile_q), 0) + if tile_size_q == 0: + raise ValueError( + "default groups_tokens_heads_q=True supports Hq/Hkv <= 32; set " + "groups_tokens_heads_q=False to use ungrouped SwapsMmaAb head bands" + ) + else: + tile_size_q = 8 if h_r <= 8 else 16 + return { + "tile_size_q": tile_size_q, + "tmem_s_cols": tile_size_q, + "tmem_o_cols": tile_size_q, + "mma_tile_n_bmm1": tile_size_q, + "mma_tile_n_bmm2": tile_size_q, + } + + +def cluster_smem_reduction_partial_smem_bytes( + *, + max_splits_kv: int, + tile_size_q: int, + headdim: int, + splits_kv: int | None = None, + correction_num_warps: int = 4, +) -> int: + """Bytes staged per owner CTA for every runtime cluster split prefix. + + For active count ``s``, each owner stages ``s`` times its slice-aligned row + band. The maximum across all prefixes covers non-divisor contractions. + """ + if max_splits_kv <= 0: + return 0 + configured_splits_kv = splits_kv if splits_kv is not None else max_splits_kv + if configured_splits_kv <= 0: + return 0 + slice_bytes = ( + correction_num_warps * WARP_THREADS * SPLIT_REDUCTION_VECTOR_BYTES_PER_THREAD + ) + row_bytes = headdim * PARTIAL_O_ELEMENT_BYTES + rows_per_slice = max(slice_bytes // row_bytes, 1) + num_slices = (tile_size_q + rows_per_slice - 1) // rows_per_slice + max_runtime_partial_rows = max( + active_splits + * max((num_slices + active_splits - 1) // active_splits, 1) + * rows_per_slice + for active_splits in range(1, configured_splits_kv + 1) + ) + return max_runtime_partial_rows * ( + headdim * PARTIAL_O_ELEMENT_BYTES + PARTIAL_STATS_VALUES_PER_ROW * FP32_BYTES + ) + + +def cluster_smem_reduction_unsupported_reason( + *, + max_splits_kv: int, + splits_kv: int | None = None, + tile_size_q: int, + headdim: int, + correction_num_warps: int = 4, + cluster_dim_x: int = 1, + max_partial_smem_bytes: int = MAX_CLUSTER_PARTIAL_SMEM_BYTES, +) -> str: + """Return why cluster SMEM reduction must be rejected, or "" if supported.""" + if max_splits_kv <= 1: + return "" + if max_splits_kv * cluster_dim_x > MAX_CLUSTER_DIM_X: + return ( + "splits_kv * clusterDimX exceeds the cluster-size limit of " + f"{MAX_CLUSTER_DIM_X}" + ) + partial_smem_bytes = cluster_smem_reduction_partial_smem_bytes( + max_splits_kv=max_splits_kv, + splits_kv=splits_kv, + tile_size_q=tile_size_q, + headdim=headdim, + correction_num_warps=correction_num_warps, + ) + if partial_smem_bytes > max_partial_smem_bytes: + return ( + "cluster leader-DSMEM partial staging would use " + f"{partial_smem_bytes} bytes, above the conservative " + f"{max_partial_smem_bytes}-byte limit" + ) + return "" + + +def compute_runtime_active_splits_kv( + *, + valid_k: int, + tile_size_kv: int, + num_insts_kv: int, + configured_splits_kv: int, +) -> int: + """Host mirror of the device runtime split-prefix calculation.""" + if valid_k < 0: + raise ValueError("valid_k must be non-negative") + if tile_size_kv <= 0: + raise ValueError("tile_size_kv must be positive") + if num_insts_kv <= 0: + raise ValueError("num_insts_kv must be positive") + if configured_splits_kv <= 0: + raise ValueError("configured_splits_kv must be positive") + total_kv_tiles = (valid_k + tile_size_kv - 1) // tile_size_kv + groups_per_split = (total_kv_tiles + configured_splits_kv * num_insts_kv - 1) // ( + configured_splits_kv * num_insts_kv + ) + local_kv_tiles = max(groups_per_split * num_insts_kv, num_insts_kv) + return (total_kv_tiles + local_kv_tiles - 1) // local_kv_tiles + + +def _max_splits_kv_by_work( + *, + seq_len_kv: int, + tile_size_kv: int, + num_insts_kv: int, + max_splits_kv: int | None = None, +) -> int: + """Return the fanout cap that retains useful KV work per CTA.""" + tile_size_per_cta_kv = tile_size_kv * num_insts_kv * MIN_LOOP_ITERS_PER_SPLIT + max_by_seq = max( + 1, + (seq_len_kv + tile_size_per_cta_kv - 1) // tile_size_per_cta_kv, + ) + if max_splits_kv is not None and max_splits_kv > 0: + max_by_seq = min(max_by_seq, max_splits_kv) + return max_by_seq + + +def enumerate_auto_splits_kv( + *, + seq_len_kv: int, + batch_size: int, + num_heads_kv: int, + tile_size_kv: int, + num_insts_kv: int, + num_q_tiles: int, + service_capacity: int, +) -> tuple[int, ...]: + """Enumerate direct and useful split fanouts for an under-filled Q grid. + + All one-wave fanouts and the first capacity-crossing fanout participate in + the empirical score. The latter is important when a partially filled wave + cannot be completed by any uniform integer fanout. + """ + if num_q_tiles <= 0: + raise ValueError("num_q_tiles must be positive") + if service_capacity <= 0: + raise ValueError("service_capacity must be positive") + max_by_work = _max_splits_kv_by_work( + seq_len_kv=seq_len_kv, + tile_size_kv=tile_size_kv, + num_insts_kv=num_insts_kv, + ) + base_grid = max(1, batch_size * num_heads_kv * num_q_tiles) + if base_grid >= service_capacity or max_by_work <= 1: + return (1,) + first_full_wave_fanout = (service_capacity + base_grid - 1) // base_grid + max_considered = min( + max_by_work, + max(first_full_wave_fanout, 2), + ) + return tuple(range(1, max_considered + 1)) + + +def select_splits_kv( + *, + seq_len_kv: int, + batch_size: int, + num_heads_kv: int, + tile_size_kv: int, + num_insts_kv: int, + num_q_tiles: int = 1, + service_capacity: int | None = None, + requested_splits_kv: int = -1, + max_splits_kv: int | None = None, +) -> int: + """ + Select a Q-grid-aware split-KV fanout. + + Every legal TileQ is considered with its actual number of Q CTAs. The + automatic fanout fills otherwise idle cluster-size-one service slots while + retaining at least ``MIN_LOOP_ITERS_PER_SPLIT`` KV iterations per CTA. + Positive caller fanouts remain pinned subject only to the KV-work cap. + """ + if num_q_tiles <= 0: + raise ValueError("num_q_tiles must be positive") + max_by_seq = _max_splits_kv_by_work( + seq_len_kv=seq_len_kv, + tile_size_kv=tile_size_kv, + num_insts_kv=num_insts_kv, + max_splits_kv=max_splits_kv, + ) + if requested_splits_kv > 0: + return max(1, min(max_by_seq, requested_splits_kv)) + + if service_capacity is None: + hardware_info = utils.HardwareInfo() + service_capacity = hardware_info.get_device_multiprocessor_count() + # B200 fallback when the runtime SM query is unavailable. + service_capacity = ( + FALLBACK_SM_COUNT_B200 if service_capacity <= 0 else service_capacity + ) + if service_capacity <= 0: + raise ValueError("service_capacity must be positive") + base_grid = max(1, batch_size * num_heads_kv * num_q_tiles) + return max(1, min(max_by_seq, max(service_capacity // base_grid, 1))) + + +SPLIT_KV_MODES = ( + "disabled", + "gmem_reduction", + "gmem_reduction_with_separate_kernel", + "cluster_smem_reduction", +) + + +def _select_auto_launch_mode( + *, + batch_size: int, + num_heads_kv: int, + seq_len_kv: int, + num_q_tiles: int = 1, + tile_size_kv: int = AUTO_LAUNCH_TILE_SIZE_KV, + persistent_min_waves: int = 1, + persistent_min_tiles_per_cta: int = 1, +) -> str: + """Pick the launch mode that best matches the kernel's parallelism budget. + + The kernel can run in three launch modes, each suited to a different + occupancy regime. + + Returns one of: + + ``"gmem_reduction"`` + Split the K/V sequence across several CTAs (Flash-Decoding GMEM + reduction). Chosen when the static grid sits under one SM wave + (``waves < 1``) *and* each CTA has at least 2,048 padded K/V tokens + to dwarf the GMEM reduction overhead. Legal split fanout is still + bounded independently by the minimum loop iterations. At ``b=1`` + there are only ``num_heads_kv`` CTAs in the static grid, which + fills a few percent of an SM-rich device; splitting K/V across + tens of CTAs unlocks the remaining bandwidth. + + ``"persistent"`` + Switch to the CLC dynamic persistent scheduler whenever the direct + launch contains more than ``persistent_min_waves`` resident CTA + waves and at least ``persistent_min_tiles_per_cta`` K/V tiles per + task. Persistence has no launch work to eliminate within one wave; + callers with heavier per-task scheduling may require more work to + amortize that overhead. + + ``"static"`` + Everything else. The default static grid is a good fit when the + launch already saturates the device with substantial per-CTA + work. + """ + if seq_len_kv <= 0 or batch_size <= 0 or num_heads_kv <= 0 or num_q_tiles <= 0: + return "static" + hardware_info = utils.HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + sm_count = FALLBACK_SM_COUNT_B200 if sm_count <= 0 else sm_count + ctas = batch_size * num_heads_kv * num_q_tiles + waves = ctas / sm_count + tiles_per_cta = (seq_len_kv + tile_size_kv - 1) // tile_size_kv + kv_tokens_per_cta = tiles_per_cta * tile_size_kv + if waves < 1 and kv_tokens_per_cta >= SPLIT_KV_MIN_TOKENS_PER_CTA: + return "gmem_reduction" + if ( + ctas > persistent_min_waves * sm_count + and tiles_per_cta >= persistent_min_tiles_per_cta + ): + return "persistent" + return "static" + + +def get_max_active_clusters_for_cluster_size(cluster_size: int) -> int: + """Query cluster occupancy after establishing CUDA's primary context. + + ``HardwareInfo`` compiles and retains a tiny occupancy-query module on its + first call. If that first call precedes CUDA context creation, retrying can + reuse an invalid module handle. Establish the context before constructing + ``HardwareInfo`` so a fresh Python process is reliable as well. + """ + cuda_context_ready = False + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.init() + torch.empty(0, device="cuda") + cuda_context_ready = True + except (ImportError, RuntimeError): + pass + try: + return utils.HardwareInfo().get_max_active_clusters(cluster_size) + except RuntimeError: + if cuda_context_ready: + raise + return max(FALLBACK_SM_COUNT_B200 // max(cluster_size, 1), 1) + + +def select_one_wave_cluster_split( + *, + initial_splits_kv: int, + minimum_splits_kv: int = 2, + num_launched_clusters: int, + tile_size_q: int, + headdim: int, + correction_num_warps: int = 4, +) -> int | None: + """Return the largest legal cluster split count whose clusters fit one wave. + + The ordinary split heuristic maximizes mainloop parallelism without + considering cluster residency. If that cluster size does not fit the + entire logical grid concurrently, search smaller cluster sizes instead of + falling back immediately to GMEM reduction. Every candidate goes through + the same cluster-size and leader-DSMEM gates; no problem-shape whitelist is + involved. ``minimum_splits_kv`` can pin an explicit fanout while retaining + the same eligibility checks. + """ + if initial_splits_kv < 2 or num_launched_clusters <= 0: + return None + lower_bound = max(minimum_splits_kv, 2) + if lower_bound > initial_splits_kv: + return None + for candidate in range(initial_splits_kv, lower_bound - 1, -1): + cluster_reason = cluster_smem_reduction_unsupported_reason( + max_splits_kv=candidate, + splits_kv=candidate, + tile_size_q=tile_size_q, + headdim=headdim, + correction_num_warps=correction_num_warps, + ) + if cluster_reason: + continue + max_active_clusters = get_max_active_clusters_for_cluster_size(candidate) + if max_active_clusters > 0 and num_launched_clusters <= max_active_clusters: + return candidate + return None + + +def validate_sliding_window_args( + sliding_window_causal: bool, attention_window_size: int +) -> None: + """Validate the sliding-window causal command-line arguments.""" + if sliding_window_causal and attention_window_size <= 0: + raise ValueError( + "attention_window_size must be positive when sliding_window_causal is enabled" + ) + + +def validate_causal_decode_lengths( + *, + seq_len_q: int, + seq_len_kv: int, + mask_type: int, +) -> None: + """Reject Q/KV lengths that cannot represent causal speculative decode.""" + if mask_type == CAUSAL and seq_len_q > seq_len_kv: + raise ValueError( + "causal decode requires seq_len_q <= seq_len_kv; " + f"got seq_len_q={seq_len_q}, seq_len_kv={seq_len_kv}" + ) + + +def _effective_seq_len_for_sliding( + seq_len_kv: int, + sliding_window_causal: bool, + attention_window_size: int, + tile_size_kv: int = 128, +) -> int: + """Return the effective KV length after sliding-window tile skipping.""" + if not sliding_window_causal: + return seq_len_kv + skipped_tiles = max(seq_len_kv - attention_window_size, 0) // tile_size_kv + return seq_len_kv - skipped_tiles * tile_size_kv + + +def _validate_split_kv_mode(mode: str) -> str: + """Validate and return a canonical split-KV launch mode.""" + if mode not in SPLIT_KV_MODES: + allowed = ", ".join(SPLIT_KV_MODES) + raise ValueError( + f"Unsupported split_kv_mode: {mode}. Expected one of: {allowed}" + ) + return mode + + +def _apply_swaps_tile_config( + cfg: FmhaDecodeConfig, + *, + explicit_fields: set[str], + num_heads_q: int, + num_heads_kv: int, +) -> None: + """Fill SwapsMmaAb tile fields while preserving explicit user fields.""" + if cfg.use_keeps_mma_ab: + return + + tile_fields = _swaps_tile_fields_for_heads( + num_heads_q, + num_heads_kv, + groups_tokens_heads_q=cfg.groups_tokens_heads_q, + ) + for key, value in tile_fields.items(): + _set_if_implicit(cfg, key, value, explicit_fields) + + if "tile_size_q" not in explicit_fields: + return + + for key in ( + "tmem_s_cols", + "tmem_o_cols", + "mma_tile_n_bmm1", + "mma_tile_n_bmm2", + ): + _set_if_implicit(cfg, key, cfg.tile_size_q, explicit_fields) + + +_MMA_SELECTION_FIELDS = { + "use_keeps_mma_ab", + "tile_size_q", + "tmem_s_cols", + "tmem_p_cols", + "tmem_o_cols", + "mma_tile_m_bmm1", + "mma_tile_n_bmm1", + "mma_tile_m_bmm2", + "mma_tile_n_bmm2", +} + +_LAUNCH_SELECTION_FIELDS = { + "use_split_kv", + "splits_kv", + "max_splits_kv", + "use_separate_reduction_kernel", + "use_cluster_smem_reduction", + "use_persistent_scheduler", +} + + +def _try_apply_auto_kv256_profile( + cfg: FmhaDecodeConfig, + *, + q_candidate: GroupedQMmaCandidate | None, + explicit_fields: set[str], + auto_tuner: bool, + split_kv_mode: str, + num_heads_q: int, + num_heads_kv: int, + splits_kv: int, + max_splits_kv: int | None, +) -> bool: + """Promote KV128 after the Q cost model selects an exact Q64 Keeps tile. + + Every other automatic Q result retains the default KV128 tile. Explicit + MMA, KV-tile, or launch policies remain caller-controlled. Device + compatibility is enforced by the public PrimTS wrapper, alongside the + other decode profiles, rather than duplicated in config selection. + """ + selection_is_unpinned = ( + auto_tuner + and split_kv_mode == "disabled" + and splits_kv == -1 + and max_splits_kv is None + and not (_MMA_SELECTION_FIELDS & explicit_fields) + and not (_LAUNCH_SELECTION_FIELDS & explicit_fields) + and "tile_size_kv" not in explicit_fields + ) + profile_is_eligible = ( + selection_is_unpinned + and q_candidate is not None + and q_candidate.variant == "keeps_mma_ab" + and q_candidate.tile_size_q == 64 + and cfg.use_paged_kv + and cfg.groups_tokens_heads_q + and cfg.headdim == 128 + and cfg.q_dtype in (Float16, BFloat16) + and cfg.q_dtype == cfg.kv_dtype == cfg.out_dtype + and num_heads_q // num_heads_kv <= 64 + ) + if not profile_is_eligible: + return False + + cfg.tile_size_kv = 256 + return True + + +def _apply_grouped_q_mma_candidate( + cfg: FmhaDecodeConfig, + candidate: GroupedQMmaCandidate, +) -> None: + """Apply one internally selected grouped-Q MMA tile to ``cfg``.""" + cfg.use_keeps_mma_ab = candidate.variant == "keeps_mma_ab" + cfg.tile_size_q = candidate.tile_size_q + if cfg.use_keeps_mma_ab: + return + cfg.tmem_s_cols = candidate.tile_size_q + cfg.tmem_o_cols = candidate.tile_size_q + cfg.mma_tile_n_bmm1 = candidate.tile_size_q + cfg.mma_tile_n_bmm2 = candidate.tile_size_q + + +def _resolve_grouped_q_launch_candidates( + cfg: FmhaDecodeConfig, + candidate: GroupedQMmaCandidate, + *, + explicit_fields: set[str], + seq_len_q: int, + seq_len_kv: int, + batch_size: int, + num_heads_q: int, + num_heads_kv: int, + service_capacity: int, +) -> tuple[GroupedQLaunchCandidate, ...]: + """Resolve one TileQ with Q-grid-aware direct and split recipes. + + The under-filled Q grid always participates in split selection. If the + GMEM split profile is unsupported, the legal direct recipe remains in the + candidate set. + """ + probe = deepcopy(cfg) + # TileQ is always scored on the common KV128 baseline. This includes an + # explicitly requested final KV width: explicit materialization remains a + # downstream constraint and must not feed back into the Q winner. + probe.tile_size_kv = 128 + _apply_grouped_q_mma_candidate(probe, candidate) + cost_tile_size_kv = 128 + if ( + candidate.variant == "keeps_mma_ab" + and candidate.tile_size_q == 64 + and probe.headdim == 128 + and probe.q_dtype == BFloat16 + and probe.q_dtype == probe.kv_dtype == probe.out_dtype + ): + # TileQ selection uses a common KV128 cost basis and is independent of + # the later KV-tile decision. BF16 Q64 has no final KV128 profile, but + # its Q geometry and modeled cost are identical to the qualified FP16 + # logical candidate. The KV selector materializes and validates the + # actual BF16 profile only after the Q winner is known. + probe.q_dtype = Float16 + probe.kv_dtype = Float16 + probe.out_dtype = Float16 + try: + _finalize_static_decode_config( + probe, + explicit_fields | {"tile_size_q"}, + ) + cost_num_insts_kv = probe.num_insts_kv + _validate_profile_support( + cfg=probe, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + split_kv_mode="disabled", + ) + except ValueError: + return () + + split_candidates = enumerate_auto_splits_kv( + seq_len_kv=seq_len_kv, + batch_size=batch_size, + num_heads_kv=num_heads_kv, + tile_size_kv=cost_tile_size_kv, + num_insts_kv=cost_num_insts_kv, + num_q_tiles=candidate.q_tiles, + service_capacity=service_capacity, + ) + recipes = [] + for candidate_splits_kv in split_candidates: + if candidate_splits_kv > 1: + split_probe = deepcopy(probe) + split_probe.use_split_kv = True + split_probe.splits_kv = candidate_splits_kv + split_probe.max_splits_kv = candidate_splits_kv + try: + _validate_profile_support( + cfg=split_probe, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + split_kv_mode="gmem_reduction", + ) + except ValueError: + continue + recipes.append( + make_grouped_q_launch_candidate( + candidate, + splits_kv=candidate_splits_kv, + seq_len_kv=seq_len_kv, + tile_size_kv=cost_tile_size_kv, + num_insts_kv=cost_num_insts_kv, + batch_size=batch_size, + num_heads_kv=num_heads_kv, + service_capacity=service_capacity, + ) + ) + return tuple(recipes) + + +def _apply_auto_grouped_q_mma_config( + cfg: FmhaDecodeConfig, + *, + explicit_fields: set[str], + auto_tuner: bool, + split_kv_mode: str, + seq_len_q: int, + seq_len_kv: int, + batch_size: int, + num_heads_q: int, + num_heads_kv: int, + splits_kv: int, + max_splits_kv: int | None, +) -> GroupedQLaunchCandidate | None: + """Select a legal fixed multi-Q MMA and KV-split recipe when unpinned. + + SQ1, packed/variable Q, explicit launch modes, ungrouped layouts, and + caller-provided MMA fields retain the existing path. Explicit fanout + controls bypass this selector together with explicit launch and MMA fields. + """ + if ( + not auto_tuner + or split_kv_mode != "disabled" + or seq_len_q <= 1 + or cfg.use_variable_seqlens_q + or not cfg.groups_tokens_heads_q + or not cfg.use_paged_kv + or cfg.num_tokens_per_page != 32 + or cfg.mask_type != CAUSAL + or cfg.use_sliding_window_causal + or cfg.use_attention_sinks + or bool(_MMA_SELECTION_FIELDS & explicit_fields) + or bool(_LAUNCH_SELECTION_FIELDS & explicit_fields) + or splits_kv != -1 + or max_splits_kv is not None + ): + return None + + candidates = enumerate_grouped_q_mma_candidates( + heads_q_per_kv=num_heads_q // num_heads_kv, + seq_len_q=seq_len_q, + ) + if not candidates: + return None + service_capacity = get_max_active_clusters_for_cluster_size(1) + supported = tuple( + recipe + for candidate in candidates + for recipe in _resolve_grouped_q_launch_candidates( + cfg, + candidate, + explicit_fields=explicit_fields, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + batch_size=batch_size, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + service_capacity=service_capacity, + ) + ) + if not supported: + return None + + has_underfilled_q_grid = any( + batch_size * num_heads_kv * recipe.mma.q_tiles < service_capacity + for recipe in supported + ) + if has_underfilled_q_grid: + selected = select_grouped_q_launch_candidate(supported) + else: + # Every legal Q grid already fills the machine, so splitting cannot + # expose otherwise-idle SMs. Compare direct recipes with the same + # mainloop-aware score instead of rewarding narrow tiles that reread + # the complete KV sequence for each additional Q tile. + direct = tuple(recipe for recipe in supported if recipe.splits_kv == 1) + if not direct: + return None + selected = select_grouped_q_direct_wave_candidate(direct) + _apply_grouped_q_mma_candidate(cfg, selected.mma) + if selected.splits_kv == 1 and selected.base_ctas > service_capacity: + # CLC persistence pays for work discovery by reusing one resident CTA + # wave. Select it only when the direct launch has more than one wave; + # a single-wave grid has no launch work for persistence to eliminate. + persistent_probe = deepcopy(cfg) + persistent_probe.use_persistent_scheduler = True + try: + _finalize_static_decode_config( + persistent_probe, + explicit_fields | {"tile_size_q"}, + ) + _validate_profile_support( + cfg=persistent_probe, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + split_kv_mode="disabled", + ) + except ValueError: + # Some public profiles support the selected direct grouped tile but + # not its CLC variant. Keep the valid static recipe rather than + # turning a successful automatic selection into a late rejection. + pass + else: + cfg.use_persistent_scheduler = True + # Keeps static defaults otherwise canonicalize an implicit tile to 64. + # Record only these local derived choices so finalization preserves them; + # the pre-Q snapshot used by the downstream KV selector remains unpinned. + explicit_fields.add("tile_size_q") + return selected + + +def _apply_default_q_grouping( + cfg: FmhaDecodeConfig, + *, + explicit_fields: set[str], +) -> None: + """Enable token/head grouping unless the caller explicitly opts out.""" + if "groups_tokens_heads_q" not in explicit_fields: + cfg.groups_tokens_heads_q = True + + +def _apply_layout_config( + cfg: FmhaDecodeConfig, + *, + qkv_layout: str, + num_tokens_per_page: int, + seq_len_kv: int, +) -> str: + """Apply contiguous/paged-KV layout fields and return the canonical layout.""" + qkv_layout = normalize_qkv_layout(qkv_layout) + if qkv_layout == "pagedKv": + validate_page_size(num_tokens_per_page) + cfg.use_paged_kv = True + cfg.num_tokens_per_page = num_tokens_per_page + cfg.max_num_pages_per_seq_kv = ( + seq_len_kv + num_tokens_per_page - 1 + ) // num_tokens_per_page + return qkv_layout + + +def _apply_feature_config( + cfg: FmhaDecodeConfig, + *, + explicit_fields: set[str], + seq_len_q: int, + num_heads_q: int, + num_heads_kv: int, + sliding_window_causal: bool, + attention_window_size: int, + use_attention_sinks: bool, +) -> None: + """Apply feature flags that affect config construction.""" + if sliding_window_causal: + cfg.use_sliding_window_causal = True + cfg.attention_window_size = attention_window_size + if use_attention_sinks: + cfg.use_attention_sinks = True + if seq_len_q > 1 or cfg.use_variable_seqlens_q or cfg.groups_tokens_heads_q: + _set_if_implicit(cfg, "max_seq_len_q", seq_len_q, explicit_fields) + _set_if_implicit( + cfg, "heads_q_per_kv", num_heads_q // num_heads_kv, explicit_fields + ) + + +def _should_auto_select_launch_mode( + cfg: FmhaDecodeConfig, + *, + auto_tuner: bool, + split_kv_mode: str, + seq_len_q: int, +) -> bool: + """Return whether automatic launch-mode selection is allowed for this shape.""" + single_query = seq_len_q == 1 + grouped_query = cfg.groups_tokens_heads_q and seq_len_q > 1 + return ( + auto_tuner + and split_kv_mode == "disabled" + and not cfg.use_persistent_scheduler + and not cfg.use_attention_sinks + and ( + single_query + or grouped_query + or cfg.use_variable_seqlens_q + or cfg.use_sliding_window_causal + ) + ) + + +def _num_q_tiles_for_launch(cfg: FmhaDecodeConfig) -> int: + """Return the physical Q-CTA multiplicity for launch occupancy.""" + q_geometry = make_q_tile_geometry( + rows_per_cta=cfg.tile_size_q, + heads_q_per_kv=cfg.heads_q_per_kv, + groups_tokens_heads_q=cfg.groups_tokens_heads_q, + ) + return max(q_geometry.num_q_ctas(cfg.max_seq_len_q), 1) + + +def _apply_auto_launch_mode( + cfg: FmhaDecodeConfig, + *, + auto_tuner: bool, + split_kv_mode: str, + batch_size: int, + num_heads_q: int, + num_heads_kv: int, + seq_len_kv: int, + seq_len_q: int, +) -> str: + """Apply the static/persistent/split-KV launch heuristic when it is allowed.""" + if not _should_auto_select_launch_mode( + cfg, + auto_tuner=auto_tuner, + split_kv_mode=split_kv_mode, + seq_len_q=seq_len_q, + ): + return split_kv_mode + + mode = _select_auto_launch_mode( + batch_size=batch_size, + num_heads_kv=num_heads_kv, + seq_len_kv=seq_len_kv, + num_q_tiles=_num_q_tiles_for_launch(cfg), + tile_size_kv=cfg.tile_size_kv, + ) + if (cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal) and mode == ( + "gmem_reduction" + ): + # Runtime Q offsets and sliding-window bounds are compatible with CLC + # work discovery, but they deliberately remain nonsplit. Underfilled + # grids therefore stay direct while grids above one resident wave use + # the same structural persistence rule as fixed-Q decode. + return split_kv_mode + if mode not in ("gmem_reduction", "persistent"): + return split_kv_mode + + # An explicitly selected MMA profile can bypass the joint selector while + # still leaving launch mode automatic. Probe the complete derived launch + # before committing it, so an unsupported Keeps split/persistent recipe + # falls back to the caller's valid direct profile. + probe = deepcopy(cfg) + if mode == "gmem_reduction": + _apply_split_kv_config( + probe, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + batch_size=batch_size, + num_heads_kv=num_heads_kv, + split_kv_mode=mode, + splits_kv=-1, + max_splits_kv=None, + sliding_window_causal=False, + attention_window_size=0, + ) + else: + probe.use_persistent_scheduler = True + try: + _validate_profile_support( + cfg=probe, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + split_kv_mode=(mode if mode == "gmem_reduction" else "disabled"), + ) + except ValueError: + return split_kv_mode + + if mode == "persistent": + cfg.use_persistent_scheduler = True + return split_kv_mode + if mode == "gmem_reduction": + return mode + return split_kv_mode + + +def _apply_split_kv_config( + cfg: FmhaDecodeConfig, + *, + seq_len_q: int, + seq_len_kv: int, + batch_size: int, + num_heads_kv: int, + split_kv_mode: str, + splits_kv: int, + max_splits_kv: int | None, + sliding_window_causal: bool, + attention_window_size: int, +) -> None: + """Resolve split-KV fanout and config flags for the selected reduction mode.""" + if split_kv_mode == "disabled": + return + + heuristic_seq_len_kv = _effective_seq_len_for_sliding( + seq_len_kv, + sliding_window_causal, + attention_window_size, + cfg.tile_size_kv, + ) + selected_splits_kv = select_splits_kv( + seq_len_kv=heuristic_seq_len_kv, + batch_size=batch_size, + num_heads_kv=num_heads_kv, + tile_size_kv=cfg.tile_size_kv, + num_insts_kv=cfg.num_insts_kv, + num_q_tiles=_num_q_tiles_for_launch(cfg), + requested_splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + ) + if selected_splits_kv <= 1: + return + + cfg.use_split_kv = True + cfg.splits_kv = selected_splits_kv + cfg.max_splits_kv = selected_splits_kv + cfg.use_persistent_scheduler = False + if split_kv_mode == "gmem_reduction_with_separate_kernel": + cfg.use_separate_reduction_kernel = True + elif split_kv_mode == "cluster_smem_reduction": + cfg.use_cluster_smem_reduction = True + + +_AUTO_SPLIT_KV_REDUCTION_MODES = ( + "gmem_reduction", + "gmem_reduction_with_separate_kernel", + "cluster_smem_reduction", +) + + +def _config_with_split_kv_mode( + cfg: FmhaDecodeConfig, split_kv_mode: str +) -> FmhaDecodeConfig: + """Return ``cfg`` with exactly one split-KV reduction mode selected.""" + + return replace( + cfg, + use_cluster_smem_reduction=split_kv_mode == "cluster_smem_reduction", + use_separate_reduction_kernel=( + split_kv_mode == "gmem_reduction_with_separate_kernel" + ), + ) + + +def _select_auto_split_kv_reduction_mode( + cfg: FmhaDecodeConfig, + *, + seq_len_q: int, + batch_size: int, + num_heads_q: int, + num_heads_kv: int, + preserve_exact_cluster_fanout: bool, +) -> tuple[FmhaDecodeConfig, str]: + """Select the first supported reduction mode from the shared policy. + + FlashInfer retains its general one-wave cluster search for legacy auto-derived + fanouts. Jointly scored recipes and caller-provided fanouts remain exact, + because changing their split count after selection would invalidate the + TileQ/split comparison or the caller's request. + """ + + modes = select_split_kv_modes( + family="fmha_decode", + topology="1cta", + tile_size_q=cfg.tile_size_q, + head_dim=cfg.headdim, + head_dim_per_cta_v=None, + split_kv=cfg.splits_kv, + available_modes=_AUTO_SPLIT_KV_REDUCTION_MODES, + ) + for split_kv_mode in modes: + trial = _config_with_split_kv_mode(cfg, split_kv_mode) + if split_kv_mode == "cluster_smem_reduction": + if not trial.can_use_cluster_smem_reduction: + continue + q_geometry = make_q_tile_geometry( + rows_per_cta=trial.tile_size_q, + heads_q_per_kv=num_heads_q // num_heads_kv, + groups_tokens_heads_q=trial.groups_tokens_heads_q, + ) + num_launched_clusters = ( + q_geometry.num_q_ctas(trial.max_seq_len_q) * num_heads_kv * batch_size + ) + cluster_split = select_one_wave_cluster_split( + initial_splits_kv=trial.splits_kv, + minimum_splits_kv=( + trial.splits_kv if preserve_exact_cluster_fanout else 2 + ), + num_launched_clusters=num_launched_clusters, + tile_size_q=trial.tile_size_q, + headdim=trial.headdim, + correction_num_warps=trial.correction_num_warps, + ) + if cluster_split is None: + continue + trial.splits_kv = cluster_split + trial.max_splits_kv = cluster_split + try: + _validate_profile_support( + cfg=trial, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + split_kv_mode=split_kv_mode, + ) + except ValueError: + continue + return trial, split_kv_mode + + # Inline GMEM was the prior automatic behavior. Final validation preserves + # its canonical error if the base profile itself is unsupported. + return _config_with_split_kv_mode(cfg, "gmem_reduction"), "gmem_reduction" + + +def _validate_profile_support( + *, + cfg: FmhaDecodeConfig, + seq_len_q: int, + num_heads_q: int, + num_heads_kv: int, + split_kv_mode: str, +) -> None: + """Reject unsupported profile combinations before kernel compilation.""" + + headdim = cfg.headdim + use_keeps_mma_ab = cfg.use_keeps_mma_ab + use_groups_tokens_heads_q = cfg.groups_tokens_heads_q + tile_size_q = cfg.tile_size_q + cfg.validate_boolean_fields() + _validate_kv256_static_config(cfg) + if cfg.mask_type not in (DENSE, CAUSAL): + raise ValueError("mask_type must be DENSE or CAUSAL") + if cfg.use_paged_kv and not cfg.use_block_sparse: + cfg.validate_paged_kv_staging_config() + if num_heads_q % num_heads_kv != 0: + raise ValueError("fmha_decode requires num_heads_q divisible by num_heads_kv") + heads_q_per_kv = num_heads_q // num_heads_kv + if cfg.heads_q_per_kv != heads_q_per_kv: + raise ValueError( + "heads_q_per_kv metadata must equal num_heads_q / num_heads_kv" + ) + cfg.validate_block_sparse_profile(heads_q_per_kv=heads_q_per_kv) + if use_groups_tokens_heads_q: + make_q_tile_geometry( + rows_per_cta=tile_size_q, + heads_q_per_kv=heads_q_per_kv, + groups_tokens_heads_q=True, + ) + # Preserve the legacy fixed FP8 Q8/D128 profile. The D64/D256 extensions + # are deliberately limited to grouped page-32 decode; broader dtypes, + # head ratios, and variable-Q layouts remain unqualified. + qualified_fp8_q8_separate_reduction_supported = ( + not use_keeps_mma_ab + and cfg.use_split_kv + and cfg.use_separate_reduction_kernel + and not cfg.use_cluster_smem_reduction + and cfg.q_dtype == Float8E4M3FN + and cfg.kv_dtype == Float8E4M3FN + and ( + (headdim == 128 and cfg.out_dtype == Float8E4M3FN) + or ( + use_groups_tokens_heads_q + and cfg.use_paged_kv + and cfg.num_tokens_per_page == 32 + and ( + (headdim == 64 and cfg.out_dtype == Float8E4M3FN) + or (headdim == 256 and cfg.out_dtype == Float16) + ) + ) + ) + and tile_size_q == 8 + and heads_q_per_kv == 8 + and seq_len_q == 1 + and cfg.max_seq_len_q == 1 + and not cfg.use_variable_seqlens_q + and cfg.splits_kv == cfg.max_splits_kv + and 2 <= cfg.max_splits_kv <= 128 + ) + if cfg.use_separate_reduction_kernel: + cluster_size = cfg.parallel_reduction_cluster_size + splits_per_cta = cfg.parallel_reduction_splits_per_cta + padded_splits = cfg.parallel_reduction_padded_splits + default_cluster_size = { + 8: 1, + 16: 2, + 32: 4, + 64: 8, + 128: 16, + }.get(padded_splits) + compact_topology_supported = ( + cfg.use_compact_parallel_reduction + and cluster_size == 1 + and splits_per_cta == cfg.max_splits_kv + ) + clustered_topology_supported = ( + not cfg.use_compact_parallel_reduction + and default_cluster_size == cluster_size + and splits_per_cta in (2, 4, 8) + and cluster_size * splits_per_cta == padded_splits + ) + if not ( + cfg.use_split_kv + and cfg.splits_kv == cfg.max_splits_kv + and 2 <= cfg.max_splits_kv <= 128 + and (compact_topology_supported or clustered_topology_supported) + ): + raise ValueError( + "separate GMEM reduction requires equal static " + "splits_kv/max_splits_kv in [2,128] using the production " + "compact or padded clustered topology" + ) + if ( + cluster_size > 1 + and get_max_active_clusters_for_cluster_size(cluster_size) <= 0 + ): + raise ValueError( + "parallel separate reduction cluster size " + f"{cluster_size} is not supported on this device" + ) + supports_grouped_keeps = cfg.supports_grouped_keeps + if cfg.tile_size_kv != 128 and not ( + cfg.tile_size_kv == 256 and supports_grouped_keeps + ): + raise ValueError( + "wide KeepsMmaAb is enabled only for the qualified FP16/BF16/D128 " + "KV256 native warp-specialized profile" + ) + if use_keeps_mma_ab and use_groups_tokens_heads_q: + if not supports_grouped_keeps: + raise ValueError( + "grouped KeepsMmaAb currently supports only validated narrow " + "profiles; pass groups_tokens_heads_q=False to use a supported " + "ungrouped Keeps profile" + ) + if cfg.use_variable_seqlens_q: + # Packed Q supports the broad grouped Swaps matrix plus the narrow + # grouped Keeps direct profile validated above. + if use_keeps_mma_ab and not ( + use_groups_tokens_heads_q and supports_grouped_keeps + ): + raise ValueError( + "packed variable-Q KeepsMmaAb requires its supported grouped " + "static direct profile" + ) + if not use_groups_tokens_heads_q and cfg.use_cluster_smem_reduction: + raise ValueError( + "packed ungrouped SwapsMmaAb does not support cluster SMEM reduction; " + "use grouped Q or a GMEM reduction mode" + ) + if use_keeps_mma_ab: + if headdim not in (64, 128, 256): + raise ValueError("fmha_decode keepsMmaAb supports headdim=64, 128, or 256") + if tile_size_q not in (64, 128): + raise ValueError("fmha_decode keepsMmaAb supports tile_size_q=64 or 128") + effective_head_dim_stage = cfg.head_dim_per_stage_kv + effective_num_insts_kv = cfg.num_insts_kv + effective_o_stages = cfg.o_stages + if effective_num_insts_kv == 1 and ( + headdim != 256 or effective_head_dim_stage != 128 or effective_o_stages != 1 + ): + raise ValueError( + "one-instance KeepsMmaAb is enabled only for the staged " + "headDim=256 profile with head_dim_per_stage_kv=128 and " + "o_stages=1" + ) + if headdim == 256 and ( + effective_head_dim_stage != 128 + or effective_num_insts_kv != 1 + or effective_o_stages != 1 + ): + raise ValueError( + "fmha_decode keepsMmaAb headDim=256 requires " + "head_dim_per_stage_kv=128, num_insts_kv=1, and o_stages=1" + ) + if headdim != 256 and effective_head_dim_stage != 0: + raise ValueError( + "split head_dim_per_stage_kv keepsMmaAb profiles are enabled " + "only for headDim=256" + ) + if not use_groups_tokens_heads_q and heads_q_per_kv != tile_size_q: + raise ValueError( + "fmha_decode keepsMmaAb requires numHeadsQPerKv == tile_size_q" + ) + if cfg.q_dtype == Float8E4M3FN and cfg.out_dtype not in ( + Float16, + Float8E4M3FN, + ): + raise ValueError( + "fmha_decode keepsMmaAb fp8 qkv path supports fp16 or fp8 output" + ) + use_split_kv = split_kv_mode != "disabled" or cfg.use_split_kv + if use_split_kv: + if cfg.q_dtype not in (Float16, BFloat16, Float8E4M3FN): + raise ValueError( + "split-KV keepsMmaAb profiles support only fp16, bf16, or fp8 qkv" + ) + if cfg.use_cluster_smem_reduction: + raise ValueError( + "fmha_decode does not support cluster SMEM reduction with keepsMmaAb" + ) + use_separate_reduction_kernel = cfg.use_separate_reduction_kernel + separate_reduction_q_layout_supported = ( + not use_groups_tokens_heads_q and heads_q_per_kv == tile_size_q + ) or use_groups_tokens_heads_q + separate_reduction_unstaged_supported = ( + headdim == 128 + and tile_size_q in (64, 128) + and effective_head_dim_stage == 0 + and effective_num_insts_kv == 2 + and effective_o_stages == 2 + ) + separate_reduction_h256_supported = ( + headdim == 256 + and tile_size_q == 128 + and effective_head_dim_stage == 128 + and effective_num_insts_kv == 1 + and effective_o_stages == 1 + ) + fixed_fp8_new_keeps_separate_reduction_supported = ( + (tile_size_q, headdim) in ((64, 64), (64, 256), (128, 64)) + and not use_groups_tokens_heads_q + and seq_len_q == 1 + and cfg.max_seq_len_q == 1 + and not cfg.use_variable_seqlens_q + and cfg.use_paged_kv + and cfg.num_tokens_per_page == 32 + and cfg.q_dtype == Float8E4M3FN + and cfg.kv_dtype == Float8E4M3FN + and cfg.out_dtype == (Float16 if headdim == 256 else Float8E4M3FN) + and cfg.splits_kv == cfg.max_splits_kv + and 2 <= cfg.max_splits_kv <= 128 + ) + separate_reduction_profile_supported = ( + separate_reduction_unstaged_supported + or separate_reduction_h256_supported + or fixed_fp8_new_keeps_separate_reduction_supported + ) + if use_separate_reduction_kernel and not ( + separate_reduction_profile_supported + and separate_reduction_q_layout_supported + and cfg.supports_reduction_dtypes + and use_split_kv + ): + raise ValueError( + "separate reduction keepsMmaAb profiles require " + "the established D128/Q64-Q128 or D256/Q128 profiles, or a " + "fixed FP8/page-32 D64/Q64-Q128 or D256/Q64 profile; valid " + "reduction dtypes, Q layout, and static split-KV are required" + ) + return + use_split_kv = split_kv_mode != "disabled" or cfg.use_split_kv + # SwapsMmaAb tensor maps provide OOB fill for a final partial head band. + # Keep the ungrouped one-token-per-CTA control available at the same tile Q + # as grouped profiles so explicit ungrouped launches retain the MMA shape. + effective_head_dim_stage = cfg.head_dim_per_stage_kv + effective_num_insts_kv = cfg.num_insts_kv + if headdim == 256: + if effective_head_dim_stage != 128 or effective_num_insts_kv != 2: + raise ValueError( + "fmha_decode SwapsMmaAb headDim=256 requires " + "head_dim_per_stage_kv=128 and num_insts_kv=2" + ) + elif effective_head_dim_stage != 0: + raise ValueError( + "split head_dim_per_stage_kv SwapsMmaAb profiles are enabled only " + "for headDim=256" + ) + if cfg.use_cluster_smem_reduction: + # Single source of truth for structural eligibility plus dtype and + # SMEM-budget checks. Grouped cluster is an explicit static choice. + if not (cfg.supports_reduction_dtypes and cfg.can_use_cluster_smem_reduction): + raise ValueError( + "cluster SMEM reduction requires static SwapsMmaAb headDim in " + "{64,128,256}, TileSizeQ in {8,16,32}, " + "either ungrouped single-token or complete-token grouped Q, " + "at least two split CTAs, " + "and fp16/bf16 or fp8 qkv with fp16/fp8 output" + ) + cluster_reason = cluster_smem_reduction_unsupported_reason( + max_splits_kv=cfg.max_splits_kv, + splits_kv=cfg.splits_kv, + tile_size_q=cfg.tile_size_q, + headdim=headdim, + correction_num_warps=cfg.correction_num_warps, + ) + if cluster_reason: + raise ValueError(f"cluster SMEM reduction rejected: {cluster_reason}") + use_separate_reduction_kernel = cfg.use_separate_reduction_kernel + separate_reduction_supported = ( + use_split_kv + and cfg.supports_reduction_dtypes + and (seq_len_q == 1 or cfg.use_variable_seqlens_q or use_groups_tokens_heads_q) + and (tile_size_q in (16, 32) or qualified_fp8_q8_separate_reduction_supported) + and headdim >= 64 + ) + if use_separate_reduction_kernel and not separate_reduction_supported: + raise ValueError( + "separate reduction SwapsMmaAb profiles require fixed SQ=1, " + "fixed grouped Q, or packed variable Q; tile_size_q in {16,32} " + "or a fixed FP8 Q8/HqPerKv8 profile (legacy D128 with FP8 output, " + "plus grouped paged-KV/page-32 D64 with FP8 output or D256 with " + "FP16 output); valid reduction dtypes and static split-KV are required" + ) + if tile_size_q and tile_size_q not in (8, 16, 32): + raise ValueError("fmha_decode SwapsMmaAb supports tile_size_q in {8,16,32}") + if headdim not in (64, 128, 256): + raise ValueError( + "fmha_decode SwapsMmaAb supports headDim in " + "{64,128,256}. headDim=256 uses the staged profile with " + "head_dim_per_stage_kv=128 and num_insts_kv=2." + ) + + +def normalize_qkv_layout(qkv_layout: str) -> str: + """Normalize CLI aliases to the canonical QKV layout names.""" + normalized = qkv_layout.strip().lower() + if normalized in ("contiguous", "contiguouskv", "dense"): + return "contiguousKv" + if normalized in ("paged", "pagedkv", "page-index", "page_index"): + return "pagedKv" + raise ValueError(f"Unsupported qkv_layout: {qkv_layout}") + + +def validate_page_size(num_tokens_per_page: int) -> None: + """Validate a paged-KV page size against supported tile shapes.""" + if num_tokens_per_page not in (16, 32, 64, 128): + raise ValueError("num_tokens_per_page must be one of 16, 32, 64, or 128") + if 128 % num_tokens_per_page != 0: + raise ValueError("num_tokens_per_page must divide the 128-token KV tile") + + +def make_decode_config( + headdim: int = 128, + args: object | None = None, + *, + seq_len_q: int = 1, + seq_len_kv: int | None = None, + batch_size: int | None = None, + num_heads_q: int | None = None, + num_heads_kv: int | None = None, + qkv_dtype: type = Float16, + o_dtype: type = Float16, + qkv_layout: str = "contiguousKv", + num_tokens_per_page: int = 32, + split_kv_mode: str = "disabled", + splits_kv: int = -1, + max_splits_kv: int | None = None, + sliding_window_causal: bool = False, + attention_window_size: int = 0, + mask_type: str | None = None, + use_attention_sinks: bool = False, + auto_tuner: bool = True, +) -> FmhaDecodeConfig: + """Build the static decode kernel config and apply auto-selection policy. + + When no launch-shape inputs are supplied, this is the simple static config + constructor: it reads only ``FmhaDecodeConfig`` fields from ``args`` and + applies static profile defaults. When launch-shape inputs are supplied, it + additionally runs the decode kernel-selection workflow below. + + Workflow: + 1. Create a default ``FmhaDecodeConfig`` and apply caller-supplied config + fields directly onto it, remembering which profile fields were explicit. + 2. Fill profile-derived defaults: dtype fields, paged-KV metadata, the + dense/causal mask, sliding-window flags, attention-sink flags, default + grouped-Q metadata for fixed or packed launches, and the SMEM-derived + KV stage count. For an unpinned fixed multi-Q paged-causal page-32 + launch, the Q selector first enumerates Swaps8/16/32 and Keeps64/128 + tiles and every useful direct/split recipe through the first + capacity-crossing fanout. + Production-valid recipes minimize the empirical TileQ + mainloop-plus-reduction proxy using their actual Q-grid CTA waves. + Only a selected Q64 Keeps recipe may then promote a qualified FP16/BF16 + D128 launch from KV128 to KV256; every other automatic Q tile retains + KV128. The final KV width re-derives launch policy instead of reusing a + fanout scored for KV128. Explicit policies remain caller-controlled. + TileQ128 remains automatic over TileQ64 only for staged D256. SQ1 is + outside this grouped-Q cost model and therefore remains KV128. + 3. Shapes outside that qualified Q/launch selector retain the general launch + policy: under-filled fixed-Q long-sequence grids use split-KV GMEM + reduction, direct grids above one resident wave use persistent + scheduling, and the rest stay static. Packed-Q and sliding-window grids + remain nonsplit but use the same structural persistence boundary. An + unsupported automatic mode falls back to direct. + 4. If split-KV is selected or requested, compute the split fanout from the + effective KV length, requested split count, max split cap, SM count, and + ``batch_size * num_heads_kv * q_tiles`` physical grid size. + 5. Choose the reduction mode for an automatic split-KV launch: prefer a + legal one-wave cluster configuration, otherwise use the standalone GMEM + reducer with inline reduction as a support fallback. A jointly scored + recipe keeps its exact fanout; the legacy launch path may search + downward for a one-wave cluster split. + 6. Validate the final profile combination before returning the config. + + Attention-sink paths skip automatic launch-mode selection. Explicit launch + modes are still validated. + """ + shape_values = (seq_len_kv, batch_size, num_heads_q, num_heads_kv) + if all(value is None for value in shape_values): + return _make_static_decode_config( + headdim, + args, + mask_type=mask_type, + sliding_window_causal=sliding_window_causal, + ) + if any(value is None for value in shape_values): + raise ValueError( + "seq_len_kv, batch_size, num_heads_q, and num_heads_kv are all " + "required for decode kernel selection" + ) + + validate_sliding_window_args(sliding_window_causal, attention_window_size) + cfg = FmhaDecodeConfig(headdim=headdim) + explicit_fields = _apply_config_source(cfg, args) + splits_kv, max_splits_kv = _resolve_explicit_split_controls( + cfg, + explicit_fields=explicit_fields, + splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + ) + _apply_mask_type_config( + cfg, + source=args, + mask_type=mask_type, + sliding_window_causal=(sliding_window_causal or cfg.use_sliding_window_causal), + explicit_fields=explicit_fields, + ) + split_kv_mode = _validate_split_kv_mode(split_kv_mode) + # Whether the launch mode was left to the auto-tuner rather than explicitly + # chosen by the caller. Only auto-derived modes are eligible for the cluster + # promotion below. + launch_mode_was_auto = split_kv_mode == "disabled" + _fanout_was_explicit = splits_kv > 0 + if num_heads_q <= 0 or num_heads_kv <= 0: + raise ValueError("fmha_decode head counts must be positive") + if num_heads_q % num_heads_kv != 0: + raise ValueError("fmha_decode requires num_heads_q divisible by num_heads_kv") + _apply_default_q_grouping( + cfg, + explicit_fields=explicit_fields, + ) + cfg.q_dtype = qkv_dtype + cfg.kv_dtype = qkv_dtype + cfg.out_dtype = o_dtype + + qkv_layout = _apply_layout_config( + cfg, + qkv_layout=qkv_layout, + num_tokens_per_page=num_tokens_per_page, + seq_len_kv=seq_len_kv, + ) + _apply_feature_config( + cfg, + explicit_fields=explicit_fields, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + sliding_window_causal=sliding_window_causal, + attention_window_size=attention_window_size, + use_attention_sinks=use_attention_sinks, + ) + if not cfg.use_variable_seqlens_q and cfg.max_seq_len_q != seq_len_q: + raise ValueError( + "fixed-Q max_seq_len_q must equal seq_len_q; use " + "use_variable_seqlens_q for a runtime-varying Q length" + ) + # Resolve Q first. The KV selector consumes the multi-Q cost-model result + # instead of forcing a Q64 profile ahead of the established Q policy. + auto_selection_explicit_fields = set(explicit_fields) + selected_grouped_q_recipe = _apply_auto_grouped_q_mma_config( + cfg, + explicit_fields=explicit_fields, + auto_tuner=auto_tuner, + split_kv_mode=split_kv_mode, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + batch_size=batch_size, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + ) + if selected_grouped_q_recipe is None: + _apply_swaps_tile_config( + cfg, + explicit_fields=explicit_fields, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + ) + kv_tile_was_promoted = _try_apply_auto_kv256_profile( + cfg, + q_candidate=( + selected_grouped_q_recipe.mma + if selected_grouped_q_recipe is not None + else None + ), + explicit_fields=auto_selection_explicit_fields, + auto_tuner=auto_tuner, + split_kv_mode=split_kv_mode, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + ) + if kv_tile_was_promoted: + # The selected Q tile remains valid, but its KV128 launch fanout, score, + # and possible persistent side effect do not. Re-derive launch policy + # below from the final KV256 work granularity. + cfg.use_persistent_scheduler = False + selected_grouped_q_recipe = None + _finalize_static_decode_config(cfg, explicit_fields) + if not cfg.use_variable_seqlens_q: + validate_causal_decode_lengths( + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + mask_type=cfg.mask_type, + ) + + # Auto launch-mode selection. Only kicks in if the caller has not already + # opted into a specific mode. Packed-Q and sliding-window shapes may select + # CLC persistence, but remain nonsplit. + if selected_grouped_q_recipe is not None: + if selected_grouped_q_recipe.splits_kv > 1: + split_kv_mode = selected_grouped_q_recipe.split_kv_mode + splits_kv = selected_grouped_q_recipe.splits_kv + max_splits_kv = selected_grouped_q_recipe.splits_kv + elif not (_LAUNCH_SELECTION_FIELDS & explicit_fields): + if split_kv_mode == "disabled" and splits_kv > 1: + split_kv_mode = "gmem_reduction" + elif splits_kv <= 0: + split_kv_mode = _apply_auto_launch_mode( + cfg, + auto_tuner=auto_tuner, + split_kv_mode=split_kv_mode, + batch_size=batch_size, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + seq_len_kv=seq_len_kv, + seq_len_q=seq_len_q, + ) + + reduction_mode_preconfigured = ( + cfg.use_cluster_smem_reduction or cfg.use_separate_reduction_kernel + ) + auto_split_kv_selected = ( + auto_tuner + and launch_mode_was_auto + and split_kv_mode == "gmem_reduction" + and not reduction_mode_preconfigured + and not cfg.use_variable_seqlens_q + and not cfg.use_sliding_window_causal + and not cfg.use_attention_sinks + ) + + _apply_split_kv_config( + cfg, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + batch_size=batch_size, + num_heads_kv=num_heads_kv, + split_kv_mode=split_kv_mode, + splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + sliding_window_causal=sliding_window_causal, + attention_window_size=attention_window_size, + ) + + if auto_split_kv_selected and cfg.use_split_kv: + cfg, split_kv_mode = _select_auto_split_kv_reduction_mode( + cfg, + seq_len_q=seq_len_q, + batch_size=batch_size, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + preserve_exact_cluster_fanout=( + _fanout_was_explicit or selected_grouped_q_recipe is not None + ), + ) + + _validate_kv256_static_config(cfg) + _finalize_warp_roles(cfg) + + _validate_profile_support( + cfg=cfg, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + split_kv_mode=split_kv_mode, + ) + return cfg diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py new file mode 100644 index 000000000000..68686febfbbb --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integer constants shared by the FMHA decode TS implementation. + +Keep non-obvious integer constants here with their rationale so config, +resource, and reduction code can use named values without duplicating comments. +""" + +# B200 has 148 SMs. Use this only when the runtime SM query is unavailable, +# so auto split-KV selection remains deterministic in offline/test flows. +FALLBACK_SM_COUNT_B200 = 148 + +# Shared-memory budget constants are in KiB because profile sizing is based on +# the hardware SMEM carveout. KV staging is capped below the full +# 218 KiB budget so Q staging, page-offset staging, and scratch can coexist. +TOTAL_SMEM_BUDGET_KIB = 218 +MAX_KV_STAGE_SMEM_KIB = 144 +BYTES_PER_KIB = 1024 + +# The supported BF16 M64N256 profile uses a 64-KiB shared K/V stage. Three +# stages occupy 192 KiB; its 16-KiB Q stage and small metadata/barrier +# allocations fit in the remaining SM100 budget. Persistent direct-output tail +# correction rotates a compact 35,840-byte exchange payload over one drained +# 64-KiB stage in this ring; split-KV keeps its fixed full exchange allocation. +# Keep this exact-profile override separate from the conservative, +# topology-independent MAX_KV_STAGE_SMEM_KIB inference above. +KV_TILE_256_SHARED_FIFO_STAGES = 3 + +# The four semantic K64 atoms are stored in the physical K slots consumed by +# the two interleaved QK instructions in this order. +KV_TILE_256_K_SLOT_FOR_SEMANTIC_ATOM = (0, 2, 1, 3) + +# Keep the old maximum as the exponent reference while a new maximum is at +# most eight log2 units larger. This avoids an output-correction round without +# letting an intermediate probability exceed 2**8; the softmax identity is +# unchanged apart from normal finite-precision rounding. +KV_TILE_256_RESCALE_THRESHOLD_LOG2 = 8.0 + +# A launch bound makes ptxas honor warpgroup ``setmaxnreg`` allocations, but +# the resulting register hand-off has a fixed cost. Paired B200 measurements +# show that it is amortized once a Q64/KV256 CTA processes at least 32 tiles +# (8K dense KV tokens); shorter loops are as fast or faster without it. +KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES = 32 + +# TMA-swizzled Q rows are padded to 128 B before computing how many KV stages +# fit in the remaining SMEM budget. +Q_ROW_ALIGNMENT_BYTES = 128 +BITS_PER_BYTE = 8 + +# Conservative per-CTA budget for cluster distributed-SMEM reduction leader +# staging. The reducer-owner CTA holds every split's partial O/stats for its +# row band; keep staging below this cap to avoid dynamic SMEM overflow. +CLUSTER_PARTIAL_SMEM_LIMIT_KIB = 96 +MAX_CLUSTER_PARTIAL_SMEM_BYTES = CLUSTER_PARTIAL_SMEM_LIMIT_KIB * BYTES_PER_KIB + +# Fused GMEM/cluster partial O is staged in 16-bit elements, and each row's stats +# are a float2 pair (max, sum). The separate-GMEM workspace contract instead +# uses one FP32 log2-LSE scalar and normalized 16-bit O. +PARTIAL_O_ELEMENT_BYTES = 2 +PARTIAL_STATS_VALUES_PER_ROW = 2 +SEPARATE_REDUCTION_LSE_VALUES_PER_ROW = 1 +FP32_BYTES = 4 + +# Hardware clusterDim.x limit for this decode reduction layout. +MAX_CLUSTER_DIM_X = 16 + +# One split-KV CTA should cover at least two loop iterations so reduction +# overhead does not dominate tiny per-split K ranges. +MIN_LOOP_ITERS_PER_SPLIT = 2 + +# The maximum number of warp groups per CTA. +MAX_WARP_GROUPS = 4 + +# Default used only when the launch helper has no resolved decode config. +# Config-aware FMHA and block-sparse callers pass their selected KV tile. +AUTO_LAUNCH_TILE_SIZE_KV = 128 + +# Split-KV is worthwhile below one static SM wave only when each CTA still owns +# enough K/V work to amortize GMEM reduction. The B200-qualified +# crossover is 2,048 tokens: 16 KV128 tiles or 8 KV256 tiles. +SPLIT_KV_MIN_TOKENS_PER_CTA = 2_048 + +# Two interleaved K/V instances form the decode cadence: instance 0 is the +# first K/P/V stream in each loop iteration and instance 1 is the second. These +# integer tags are passed through TS work calls and used in constexpr branches. +KV_INST0 = 0 +KV_INST1 = 1 + +# Compact K/V selector used by shared SMEM resources. Keep this as an integer +# contract because the JIT work-call plumbing expects constexpr scalar values. +KV_KIND_K = 0 +KV_KIND_V = 1 + +# One hardware warp. Barrier participant counts are expressed as warps * lanes. +WARP_THREADS = 32 + +# TMEM column layout for staged SwapsMmaAb O. A TMEM row holds 256 columns, and +# the tcgen05 descriptor encodes a 16-row jump in the high 16 bits. +TMEM_COLUMNS_PER_ROW = 256 +TMEM_ROW_STRIDE = 16 << 16 + +# Number of scalar softmax/output values packed in one register for the +# supported element widths. +PACKED_REGISTER_BYTES = 4 +FP8_VALUES_PER_REG = 4 +FP16_VALUES_PER_REG = 2 + +# Per-lane register ownership denominators for packed output fragments. +FP8_OUTPUT_ELEMENTS_PER_REG_GROUP = 512 +FP16_OUTPUT_ELEMENTS_PER_REG_GROUP = 256 + +# SwapsMmaAb maps up to eight Q heads into one q-repetition group. +Q_REPETITION_GROUP_HEADS = 8 + +# Packed P register count per q-repetition in the SwapsMmaAb path. +FP8_P_PACKED_REGS_PER_Q_REPEAT = 2 +FP16_P_PACKED_REGS_PER_Q_REPEAT = 4 + +# Standalone reducer CTA shape. Each thread owns a 16-byte vector, so one CTA +# reduces one contiguous 8 KiB slice of the partial-O buffer. +REDUCTION_THREADS_PER_CTA = 512 +REDUCTION_BYTES_PER_THREAD = 16 +REDUCTION_BYTES_PER_SLICE = REDUCTION_THREADS_PER_CTA * REDUCTION_BYTES_PER_THREAD + +# Clustered standalone reducer shape. One 128-thread CTA covers a contiguous +# 2 KiB partial-O slice with one 16-byte vector per thread. Each cluster rank +# owns a compile-time 2, 4, or 8 split slots. Loads are batched in groups of at +# most four; padded split slots remain neutral and never form GMEM pointers. +PARALLEL_REDUCTION_THREADS_PER_CTA = 128 +PARALLEL_REDUCTION_BYTES_PER_SLICE = ( + PARALLEL_REDUCTION_THREADS_PER_CTA * REDUCTION_BYTES_PER_THREAD +) +PARALLEL_REDUCTION_LOAD_BATCH = 4 +PARALLEL_REDUCTION_FINAL_REDUCERS = 4 + +# Each reduction thread produces an 8-element O vector backed by four packed +# 16-bit registers. +OUTPUT_VALUES_PER_THREAD = 8 +FP8_PACKED_OUTPUT_REGS_PER_THREAD = 2 +PACKED_OUTPUT_REGS_PER_THREAD = 4 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py new file mode 100644 index 000000000000..929696c30d3e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py @@ -0,0 +1,3020 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FMHA decode TS kernel assembly. + +Assembles all resources, tasks, pipeline configs, and the dependency graph +into a TaskManager for the SwapsMmaAb decode kernel. + +SwapsMmaAb, also shortened to swapsAb, names the MMA layout choice that maps +the logical attention operands onto the opposite MMA A/B roles from the +textbook QK form. BMM1 issues K as the MMA A operand and Q as the MMA B +operand, producing S = K * Q^T so KV tokens occupy the MMA M axis and the GQA +head group occupies the small N axis. BMM2 follows the same convention with V +as A and P as B. + +Entry points: + - build_decode_task_manager() — pure Python, validation only (no GPU) + - FmhaDecodeTs — GPU kernel class with @cute.jit + @cute.kernel +""" + +import math + +import cutlass +import cutlass.experimental.cuda as cuda +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +from cuda.bindings import driver as cuda_drv +from cutlass import Float32, Int32, Int64 +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.enums import PipelineType +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + SmemAllocator, + TmemAllocator, +) +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + PipelineConfig, + TileSchedulerConfig, + WorkQueue, +) +from cutlass.experimental.task_scheduling.task import Task +from cutlass.experimental.task_scheduling.task_manager import TaskManager + +from ..._block_sparse.common import ( + _block_sparse_kv_atom_size, + _prepared_kv_routes_are_block_aligned, +) +from ..._block_sparse.prepared import _BlockSparseRouteLayout +from ..tensor_map import ( + create_tensor_map_ragged_from_tensor, + create_tensor_map_tiled, + create_tensor_map_tiled_from_view, +) +from .fmha_decode_config import FmhaDecodeConfig +from .fmha_decode_constants import ( + KV_KIND_K, + KV_KIND_V, + KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES, +) +from .fmha_decode_resources import ( + SmemBlockSparseKvMetadataResource, + SmemBlockSparseSoftmaxMetadataResource, + SmemKvTileResource, + SmemKvResource, + SmemPageOffsetsKvResource, + SmemQResource, + TmemCorrResource, + TmemOResource, + SmemPResource, + TmemSResource, + TmemStatsDoneResource, + TmemSoftmaxGlobalResource, + TmemSoftmaxLocalResource, + TmemSoftmaxOrderResource, +) +from .fmha_decode_resources.helpers_common import ( + _q_group_token_base, + _q_seq_bounds, +) +from .fmha_decode_resources.helpers_kv_tile_idx import _runtime_active_splits_kv +from .fmha_decode_tasks import ( + PackedDecodeWorkQueue, + ScheduleTokenThrottleResource, + SmemKvReuseCreditResource, + create_block_sparse_load_tasks_per_inst, + create_correction_task, + create_correction_task_one_inst_qkv, + create_load_task, + create_load_task_one_inst_qkv, + create_load_task_split_kv, + create_mma_task, + create_mma_task_one_inst_qkv, + create_mma_task_split_kv, + create_page_offsets_task, + create_page_offsets_task_one_inst_qkv, + create_page_offsets_task_split_kv, + create_padding_task, + create_scheduler_task, + create_softmax0_task, + create_softmax1_task, +) + +from .reduction import ( # noqa: F401 + decode_gen_separate_reduction_kernel, + fmha_decode_separate_reduction_launch, +) + +_PERSISTENT_SCHEDULE_TOKEN_STAGES = 2 + + +def _block_sparse_bshd_tma_strides( + *, + q_seq: cutlass.Integer | int, + h_q: cutlass.Integer | int, + h_k: cutlass.Integer | int, + s_k: cutlass.Integer | int, + d: cutlass.Integer | int, +) -> tuple[ + tuple[cutlass.Integer | int, ...], + tuple[cutlass.Integer | int, ...], +]: + """Build BSHD TensorMap strides in 16-byte units using Int64 math. + + The raw TensorMap API omits the implicit contiguous stride and takes the + remaining strides in 16-byte units. Block-sparse attention supports only + 16-bit Q/K/V with headDim=128, so one stride unit contains eight elements. + Keep every returned value in Int64: the outer batch stride can exceed the + signed Int32 range even though each public tensor dimension is Int32. + """ + + elements_per_stride_unit = 8 + d_units = Int64(d // elements_per_stride_unit) + h_r = h_q // h_k + return ( + ( + d_units, + Int64(h_r) * d_units, + Int64(h_q) * d_units, + Int64(q_seq) * Int64(h_q) * d_units, + ), + ( + Int64(h_k) * d_units, + d_units, + Int64(s_k) * Int64(h_k) * d_units, + ), + ) + + +def _resolve_block_sparse_per_inst_load_topology( + cfg: FmhaDecodeConfig, + *, + use_clc_dynamic: bool, +) -> tuple[tuple[int, int], tuple[int, int] | None] | None: + """Reuse idle WG3 padding warps for two independent sparse load streams. + + The return value contains the two load warp indices and an optional + residual padding task ``(warp_idx, num_warps)``. ``None`` identifies a + noncanonical override for which the caller keeps the common load task. + """ + + if cfg.load_num_warps != 1: + return None + padding_warps = tuple( + range( + cfg.wg3_padding_warp_idx, + cfg.wg3_padding_warp_idx + cfg.wg3_padding_num_warps, + ) + ) + if use_clc_dynamic: + # Load1 consumes the only otherwise-idle warp in WG3. + if cfg.scheduler_num_warps != 1 or len(padding_warps) != 1: + return None + role_warps = ( + cfg.mma_warp_idx, + cfg.scheduler_warp_idx, + cfg.clc_load_warp_idx, + padding_warps[0], + ) + if len(set(role_warps)) != len(role_warps): + return None + if len({warp_idx // 4 for warp_idx in role_warps}) != 1: + return None + return (cfg.clc_load_warp_idx, padding_warps[0]), None + + if len(padding_warps) != 2: + return None + role_warps = (cfg.mma_warp_idx, cfg.load_warp_idx, *padding_warps) + if len(set(role_warps)) != len(role_warps): + return None + if len({warp_idx // 4 for warp_idx in role_warps}) != 1: + return None + return (cfg.load_warp_idx, padding_warps[0]), (padding_warps[1], 1) + + +def _stages_page_ids_per_tile(uses_paired_page_offset_resources: bool) -> bool: + """Return whether each published page-offset stage owns exactly one tile.""" + + # Shared resources retain an aligned 32-ID window so all lanes issue one + # coalesced page-table transaction. Paired K0/K1 and V0/V1 resources use + # exact per-tile stages so a pair may safely cross a 32-ID boundary. + return uses_paired_page_offset_resources + + +def _compute_decode_gen_loop_domain(total_kv_tiles: int, num_insts_kv: int) -> int: + """Number of post-head steady-state iterations. + + HEAD consumes the first `num_insts_kv` K tiles. The remaining tiles are + processed in staggered groups of `num_insts_kv`, so odd tail groups still + need one final loop iteration, matching the schedule's pull-down behavior. + """ + remaining_kv_tiles = max(total_kv_tiles - num_insts_kv, 0) + return (remaining_kv_tiles + num_insts_kv - 1) // num_insts_kv + + +def _compute_total_kv_tiles(seq_len_kv: int, tile_size_kv: int) -> int: + """Number of KV tiles needed for a fixed-length launch.""" + return (seq_len_kv + tile_size_kv - 1) // tile_size_kv + + +def _decode_min_blocks_per_mp(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: + """Return the launch bound needed for dynamic register reallocation. + + ``setmaxnreg`` needs a kernel-entry occupancy bound before ptxas can infer + the initial per-thread register allocation. KV256 only pays that fixed + hand-off cost for a long enough mainloop; the established profiles below + retain their existing unconditional launch bounds. + """ + kv256_reallocation = ( + cfg.tile_size_kv == 256 + and _compute_total_kv_tiles(seq_len_kv, cfg.tile_size_kv) + >= KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES + ) + return int( + kv256_reallocation + or cfg.tile_size_q == 8 + or (cfg.tile_size_q == 16 and cfg.q_dtype_bytes == 1) + or (cfg.use_keeps_mma_ab and cfg.tile_size_q == 128) + ) + + +def _compute_static_num_skipped_kv_tiles(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: + """Return full leading KV tiles skipped by a static sliding window.""" + if not cfg.use_sliding_window_causal or cfg.max_seq_len_q > 1: + return 0 + return max(seq_len_kv - cfg.attention_window_size, 0) // cfg.tile_size_kv + + +def _compute_static_window_start_idx(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: + """Return the token index where a static sliding window begins.""" + if not cfg.use_sliding_window_causal or cfg.max_seq_len_q > 1: + return 0 + return max(seq_len_kv - cfg.attention_window_size, 0) + + +def _configure_static_sliding_window( + cfg: FmhaDecodeConfig, seq_len_kv: int, bias_kv_tma: bool = False +) -> int: + """Populate fixed-length sliding metadata and return effective seqLenKv.""" + skipped_tiles = _compute_static_num_skipped_kv_tiles(cfg, seq_len_kv) + skipped_tokens = skipped_tiles * cfg.tile_size_kv + window_start_idx = _compute_static_window_start_idx(cfg, seq_len_kv) + effective_seq_len_kv = seq_len_kv - skipped_tokens + cfg.use_static_sliding_kv_tma_bias = bias_kv_tma and skipped_tiles > 0 + cfg.static_seq_len_kv = ( + effective_seq_len_kv if cfg.use_static_sliding_kv_tma_bias else seq_len_kv + ) + cfg.static_num_skipped_kv_tiles = ( + 0 if cfg.use_static_sliding_kv_tma_bias else skipped_tiles + ) + cfg.static_window_start_idx = ( + window_start_idx - skipped_tokens + if cfg.use_static_sliding_kv_tma_bias + else window_start_idx + ) + return effective_seq_len_kv + + +def _compute_local_kv_tiles(cfg: FmhaDecodeConfig, total_kv_tiles: int) -> int: + """KV tiles covered by each CtaKv in split-KV mode.""" + if not cfg.use_split_kv: + return total_kv_tiles + tiles_per_cta_group = cfg.splits_kv * cfg.num_insts_kv + num_groups = (total_kv_tiles + tiles_per_cta_group - 1) // tiles_per_cta_group + return max( + cfg.num_insts_kv, + num_groups * cfg.num_insts_kv, + ) + + +def _build_decode_gen_schedule( + cfg: FmhaDecodeConfig, + total_kv_tiles: int | Int32, + scale_softmax_log2: Float32 | None = None, + o_ptr: cute.Pointer | None = None, + output_scale: Float32 | None = None, + partial_o_ptr: cute.Pointer | None = None, + partial_stats_ptr: cute.Pointer | None = None, + split_kv_counter_ptr: cute.Pointer | None = None, + attention_sinks_ptr: cute.Pointer | None = None, + seqlens_kv: cute.Pointer | None = None, + cu_seqlens_q: cute.Pointer | None = None, + max_seq_len_kv: int | Int32 = 0, + corr_max_seq_len_kv: int | Int32 | None = None, + num_heads_kv: Int32 | None = None, + h_r: Int32 | None = None, + tma_desc_q: cutlass.Pointer | None = None, + tma_desc_k: cutlass.Pointer | None = None, + tma_desc_v: cutlass.Pointer | None = None, + tma_desc_k_atom: cutlass.Pointer | None = None, + tma_desc_v_atom: cutlass.Pointer | None = None, + page_idx_kv: cute.Pointer | None = None, + h_k_idx: Int32 | None = None, + b_idx: Int32 | None = None, + q_group_idx: Int32 | None = None, + q_token_offset: Int32 | None = None, + seq_len_q: Int32 | None = None, + active_splits_kv: Int32 | None = None, + static_full_split_prefix: bool = False, + tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None = None, + clc_response_ptr: cute.Pointer | None = None, + use_variable_seqlens_kv: bool = False, + use_native_paged_kv: bool = False, + use_static_native_seqlens_kv: bool = False, + block_tables: cute.Pointer | None = None, + block_table_capacity: Int32 | None = None, + block_table_row_stride: Int64 | None = None, + sparse_row_route_offsets: cute.Pointer | None = None, + sparse_row_route_counts: cute.Pointer | None = None, + sparse_route_metadata: cute.Pointer | None = None, +) -> tuple[ + list[Task], + dict[MemoryResource, list[MemoryResource]], + dict[tuple[MemoryResource, MemoryResource], set[str]], + SmemAllocator, + TmemAllocator, + list[MemoryResource], +]: + """Build all resources, tasks, and dep graph. + + Parameters + ---------- + cfg : FmhaDecodeConfig + Kernel configuration. + total_kv_tiles : int or Int32 + Total number of KV tiles (seqLenKv / tileSizeKv). + tma_desc_q/k/v : TMA descriptor pointers (None for validation-only mode). + h_k_idx, b_idx, q_group_idx : Grid coordinates (None for validation-only mode). + corr_max_seq_len_kv : Bound passed to TmemCorrResource; defaults to + ``max_seq_len_kv``. The GPU kernel uses the constexpr ``seq_len_kv`` + here (full sequence) while other resources use the static/varlen + runtime value. + + Returns + ------- + tuple + (task_list, dep_graph, dma labels, smem_allocator, tmem_allocator, + eager_init_resources) + """ + if cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1 and not cfg.uses_tmem_p: + raise ValueError( + "one-instance KeepsMmaAb is enabled only for the staged headDim=256 " + "profile with head_dim_per_stage_kv=128 and o_stages=1" + ) + if use_native_paged_kv and not cfg.use_paged_kv: + raise ValueError("native paged-KV ABI requires cfg.use_paged_kv=True") + if use_native_paged_kv and ( + block_tables is None + or block_table_capacity is None + or block_table_row_stride is None + ): + raise ValueError( + "native paged-KV ABI requires block tables, capacity, and row stride" + ) + if cfg.use_paged_kv: + cfg.validate_paged_kv_staging_config() + if cfg.use_block_sparse: + if tma_desc_q is not None: + if sparse_row_route_offsets is None: + raise ValueError( + "sparse_row_route_offsets is required for block-sparse kernel " + "construction" + ) + if sparse_row_route_counts is None: + raise ValueError( + "sparse_row_route_counts is required for block-sparse kernel " + "construction" + ) + if sparse_route_metadata is None: + raise ValueError( + "sparse_route_metadata is required for block-sparse kernel " + "construction" + ) + segment_tensormaps = { + "tma_desc_k_atom": tma_desc_k_atom, + "tma_desc_v_atom": tma_desc_v_atom, + } + for name, descriptor in segment_tensormaps.items(): + if descriptor is None: + raise ValueError( + f"{name} is required for block-sparse kernel construction" + ) + if num_heads_kv is None: + raise ValueError( + "num_heads_kv is required for block-sparse kernel construction" + ) + if corr_max_seq_len_kv is None: + corr_max_seq_len_kv = max_seq_len_kv + if h_k_idx is None: + h_k_idx = Int32(0) + if b_idx is None: + b_idx = Int32(0) + if q_group_idx is None: + q_group_idx = Int32(0) + if q_token_offset is None: + q_token_offset = Int32(0) + if seq_len_q is None: + seq_len_q = Int32(cfg.max_seq_len_q) + + WARP_SIZE = 32 + Agent = pipeline.Agent + cta_layout = (1, 1, 1, 1) + + # ------------------------------------------------------------------ + # Cooperative groups + # ------------------------------------------------------------------ + tma_producer = pipeline.CooperativeGroup(Agent.Thread) + page_offsets_grp = pipeline.CooperativeGroup( + Agent.Thread, cfg.page_offsets_num_warps * WARP_SIZE + ) + load_grp = pipeline.CooperativeGroup(Agent.Thread, cfg.load_num_warps * WARP_SIZE) + umma_hw = pipeline.CooperativeGroup(Agent.Thread) + # The staged one-instance S/P overlay uses this group for overwrite credit. + mma_grp = pipeline.CooperativeGroup(Agent.Thread, cfg.mma_num_warps * WARP_SIZE) + softmax0_grp = pipeline.CooperativeGroup( + Agent.Thread, cfg.softmax0_num_warps * WARP_SIZE + ) + softmax1_grp = pipeline.CooperativeGroup( + Agent.Thread, cfg.softmax1_num_warps * WARP_SIZE + ) + correction_grp = pipeline.CooperativeGroup( + Agent.Thread, cfg.correction_num_warps * WARP_SIZE + ) + scheduler_grp = pipeline.CooperativeGroup( + Agent.Thread, cfg.scheduler_num_warps * WARP_SIZE + ) + + # ------------------------------------------------------------------ + # Pipeline configs + # ------------------------------------------------------------------ + # Leave barrier_ptr unset so SmemAllocator packs every pipeline barrier + # into the unified block. Separate barrier arrays create an alignment gap + # before the 1024-byte-aligned data block and overflow near-capacity Q128. + use_paged_kv = cfg.use_paged_kv + use_dense_page_offsets = use_paged_kv and not cfg.use_block_sparse + use_one_inst_qkv = cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1 + one_inst_tmem_stages = 2 if use_one_inst_qkv else 1 + one_inst_kv_stages = cfg.num_head_dim_stages_kv if use_one_inst_qkv else 1 + use_distributed_split_kv_stages = not use_one_inst_qkv + if cfg.tile_size_q == 128 and use_distributed_split_kv_stages: + # Q128's four instruction-local K0/K1/V0/V1 rings need equal depth. + # Round the inferred aggregate budget down to a complete balanced set; + # on the FP8 decode profile this is 2/2/2/2, matching the roughly + # 165-KiB staged footprint of the corresponding reference profile. + balanced_total_stages = max( + (cfg.kv_stages // (2 * cfg.num_insts_kv)) * cfg.num_insts_kv, + cfg.num_insts_kv, + ) + split_total_k_stages = balanced_total_stages + split_total_v_stages = balanced_total_stages + else: + split_total_k_stages = ( + max(cfg.kv_stages // 2, cfg.num_insts_kv) + if use_distributed_split_kv_stages + else cfg.num_insts_kv + ) + split_total_v_stages = ( + max(cfg.kv_stages - split_total_k_stages, cfg.num_insts_kv) + if use_distributed_split_kv_stages + else cfg.num_insts_kv + ) + split_k0_stages = ( + one_inst_kv_stages + if use_one_inst_qkv + else max((split_total_k_stages + cfg.num_insts_kv - 1) // cfg.num_insts_kv, 1) + ) + split_k1_stages = ( + 1 + if use_one_inst_qkv + else max((split_total_k_stages + cfg.num_insts_kv - 2) // cfg.num_insts_kv, 1) + ) + split_v0_stages = ( + one_inst_kv_stages + if use_one_inst_qkv + else max((split_total_v_stages + cfg.num_insts_kv - 1) // cfg.num_insts_kv, 1) + ) + split_v1_stages = ( + 1 + if use_one_inst_qkv + else max((split_total_v_stages + cfg.num_insts_kv - 2) // cfg.num_insts_kv, 1) + ) + use_ordered_softmax_barrier = ( + not use_one_inst_qkv and cfg.uses_ordered_softmax_barrier + ) + # A two-inst Keeps profile can use the deeper shared K/V FIFO when stats + # are standalone and P remains in SMEM. Keep instruction-local FIFOs when + # stats or TMEM-P alias S: their overwrite-credit cadence is tied to each + # instruction. Dense Swaps uses the shared FIFO, including staged H256. + # Sparse KV128 keeps instruction-local rings in either MMA orientation; + # sparse KV256 reuses its only feasible three-stage shared data ring while + # retaining instruction-local route metadata. The load warp issues V(route + # R) before replacing that metadata with route R+1, so its lifetime remains + # independent of the K/V data-ring depth. + # With cfg.keeps_stats_via_smem the stats-alias justification no longer + # applies, but the shared FIFO still causes a material Q128 regression, so + # the instruction-local FIFO gate remains part of that kernel policy. + use_per_inst_kv_resources = (cfg.use_block_sparse and cfg.tile_size_kv != 256) or ( + cfg.use_keeps_mma_ab + and cfg.tile_size_kv != 256 + and (not cfg.keeps_separates_tmem_s_and_stats or cfg.uses_two_inst_tmem_p) + ) + # B8/B16 issue enough fine-grained TMA copies to benefit from reusing a + # padding warp as a second issuer. The host policy applies one KV-side + # crossover across all two-instance Swaps Q tiles. + supports_per_inst_block_sparse_load_tasks = ( + cfg.use_block_sparse + and cfg.use_parallel_sparse_kv_loads + and not cfg.use_keeps_mma_ab + and cfg.kv_block_size in (8, 16) + and cfg.num_insts_kv == 2 + ) + use_separate_kv_page_offset_resources = ( + use_dense_page_offsets and use_per_inst_kv_resources and not use_one_inst_qkv + ) + # Paired resources publish independent K0/K1 and V0/V1 stages. Shared + # split-KV retains the aligned 32-ID representation for its optional + # native held-window path. + stage_page_ids_per_tile = _stages_page_ids_per_tile( + use_separate_kv_page_offset_resources, + ) + + smem_q_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.q_stages, + num_bytes=cfg.smem_q_tile_bytes, + producer_group=tma_producer, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.kv_stages, + num_bytes=cfg.smem_kv_tile_bytes, + producer_group=tma_producer, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_k0_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=split_k0_stages, + num_bytes=cfg.smem_kv_tile_bytes, + producer_group=tma_producer, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_k1_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=split_k1_stages, + num_bytes=cfg.smem_kv_tile_bytes, + producer_group=tma_producer, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_v0_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=split_v0_stages, + num_bytes=cfg.smem_kv_tile_bytes, + producer_group=tma_producer, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_v1_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=split_v1_stages, + num_bytes=cfg.smem_kv_tile_bytes, + producer_group=tma_producer, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + + def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: + """Create the async page-offsets pipeline for the selected stage count.""" + if num_stages is None: + num_stages = cfg.page_offsets_stages + return PipelineConfig( + num_stages=num_stages, + num_bytes=0, + producer_group=page_offsets_grp, + consumer_group=load_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + + smem_page_offsets_cfg = None + smem_page_offsets_v_cfg = None + if use_dense_page_offsets: + page_offsets_stages = ( + 3 if use_separate_kv_page_offset_resources else cfg.page_offsets_stages + ) + smem_page_offsets_cfg = _make_page_offsets_cfg(page_offsets_stages) + if use_separate_kv_page_offset_resources: + smem_page_offsets_v_cfg = _make_page_offsets_cfg(page_offsets_stages) + sparse_softmax_metadata0_cfg = None + sparse_softmax_metadata1_cfg = None + if cfg.use_block_sparse: + # Two stages are sufficient for the split-ring cadence: one route can + # await Softmax while Load publishes the next route for the same inst. + sparse_softmax_metadata0_cfg = PipelineConfig( + num_stages=2, + num_bytes=0, + producer_group=load_grp, + consumer_group=softmax0_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + sparse_softmax_metadata1_cfg = PipelineConfig( + num_stages=2, + num_bytes=0, + producer_group=load_grp, + consumer_group=softmax1_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + # tmem_s0/s1, smem_p0/p1, tmem_o, softmax_local cfgs go through direct + # PipelineConfig() so advance_on_wait=True can be set (factories don't + # expose it). + tmem_s0_cfg = PipelineConfig( + num_stages=one_inst_tmem_stages, + num_bytes=0, + producer_group=umma_hw, + consumer_group=softmax0_grp, + pipeline_type=PipelineType.UmmaAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + tmem_s1_cfg = PipelineConfig( + num_stages=1, + num_bytes=0, + producer_group=umma_hw, + consumer_group=softmax1_grp, + pipeline_type=PipelineType.UmmaAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_p0_cfg = None + smem_p1_cfg = None + if not cfg.streams_tmem_p_fragments: + smem_p0_cfg = PipelineConfig( + num_stages=one_inst_tmem_stages, + num_bytes=0, + producer_group=softmax0_grp, + consumer_group=umma_hw, + pipeline_type=PipelineType.AsyncUmma, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_p1_cfg = PipelineConfig( + num_stages=one_inst_tmem_stages, + num_bytes=0, + producer_group=softmax1_grp, + consumer_group=umma_hw, + pipeline_type=PipelineType.AsyncUmma, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + tmem_o_cfg = PipelineConfig( + num_stages=cfg.o_stages, + num_bytes=0, + producer_group=umma_hw, + consumer_group=correction_grp, + pipeline_type=PipelineType.UmmaAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + softmax_local0_cfg = PipelineConfig( + num_stages=one_inst_tmem_stages, + num_bytes=0, + producer_group=softmax0_grp, + consumer_group=correction_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + + softmax_local1_cfg = PipelineConfig( + num_stages=1, + num_bytes=0, + producer_group=softmax1_grp, + consumer_group=correction_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + + # Two-instance Keeps keeps stats outside S and orders each same-instance PV + # before the next QK, so it needs no stats-done credit. + # The staged one-instance path needs an overwrite-credit gate across its + # double-buffered S/P overlay: correction returns the stage credit before + # MMA can reissue QK into those columns. + stats_done0_cfg = None + stats_done1_cfg = None + resource_dependency_graph: dict[MemoryResource, list[MemoryResource]] + if use_one_inst_qkv: + stats_done0_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=one_inst_tmem_stages, + producer_group=mma_grp, + consumer_group=correction_grp, + cta_layout_vmnk=cta_layout, + ) + + # tmemSoftmaxGlobal, tmemCorr: no pipeline (pipeline_config=None) + + # ------------------------------------------------------------------ + # Create resources + # ------------------------------------------------------------------ + work_queue = None + schedule_token_throttle = None + smem_kv_reuse_credit = None + # CLC remains the single persistent policy for every supported topology. + # The stock static WorkQueue advances and decodes coordinates separately + # in every task, which regresses multi-wave decode workloads. CLC computes + # each schedule token once on the scheduler warp and broadcasts it to the workers. + use_clc_dynamic = cfg.use_persistent_scheduler + per_inst_block_sparse_load_topology = None + if supports_per_inst_block_sparse_load_tasks: + # Dual issuers are a performance choice. If a future recipe has no + # compatible idle-warp placement, the common one-warp load task remains + # correct and consumes the same disjoint K/V metadata pipelines. + per_inst_block_sparse_load_topology = ( + _resolve_block_sparse_per_inst_load_topology( + cfg, + use_clc_dynamic=use_clc_dynamic, + ) + ) + use_per_inst_block_sparse_load_tasks = ( + per_inst_block_sparse_load_topology is not None + ) + if use_clc_dynamic: + num_consumer_threads = 16 * WARP_SIZE + wq_pipeline_config = PipelineConfig.create_clc_fetch_async_pipeline_cfg( + num_stages=_PERSISTENT_SCHEDULE_TOKEN_STAGES, + num_bytes=16, + producer_group=pipeline.CooperativeGroup(Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + Agent.Thread, num_consumer_threads + ), + cta_layout_vmnk=cta_layout, + ) + work_queue_kwargs = { + "tile_scheduler_config": ( + TileSchedulerConfig.create_clc_dynamic_persistent_tile_scheduler_params( + tile_scheduler_params=tile_sched_params, + response_ptr=clc_response_ptr, + ) + ), + "pipeline_config": wq_pipeline_config, + "name": "work_queue", + } + if cfg.use_variable_seqlens_q: + work_queue = PackedDecodeWorkQueue( + cfg=cfg, + cu_seqlens_q=cu_seqlens_q, + **work_queue_kwargs, + ) + else: + work_queue = WorkQueue(**work_queue_kwargs) + schedule_token_throttle = ScheduleTokenThrottleResource( + pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( + num_stages=_PERSISTENT_SCHEDULE_TOKEN_STAGES, + producer_group=load_grp, + consumer_group=scheduler_grp, + cta_layout_vmnk=cta_layout, + ), + name="schedule_token_throttle", + ) + if cfg.uses_rotating_kv256_exchange: + smem_kv_reuse_credit = SmemKvReuseCreditResource( + cfg=cfg, + pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=load_grp, + consumer_group=correction_grp, + cta_layout_vmnk=cta_layout, + ), + name="smem_kv_reuse_credit", + ) + smem_q = SmemQResource( + pipeline_config=smem_q_cfg, + cfg=cfg, + tma_desc_q=tma_desc_q, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + q_token_offset=q_token_offset, + seq_len_q=seq_len_q, + name="smemQ", + ) + # Native callers provide one canonical sequence-length tensor alongside + # their fixed page table, so native mode reuses the existing variable-length + # domain, split, sliding-window, and masking paths. + use_runtime_seqlens_kv = use_variable_seqlens_kv or ( + use_native_paged_kv and not use_static_native_seqlens_kv + ) + kv_seqlens = seqlens_kv if use_runtime_seqlens_kv else None + smem_page_offsets = None + smem_page_offsets_v = None + if use_dense_page_offsets: + smem_page_offsets = SmemPageOffsetsKvResource( + pipeline_config=smem_page_offsets_cfg, + cfg=cfg, + stage_page_ids_per_tile=stage_page_ids_per_tile, + page_idx_kv=page_idx_kv, + seqlens_kv=kv_seqlens, + use_native_paged_kv=use_native_paged_kv, + block_tables=block_tables, + block_table_row_stride=block_table_row_stride, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + name=( + "smemPageOffsetsKvK" + if use_separate_kv_page_offset_resources + else "smemPageOffsetsKv" + ), + ) + if use_separate_kv_page_offset_resources: + smem_page_offsets_v = SmemPageOffsetsKvResource( + pipeline_config=smem_page_offsets_v_cfg, + cfg=cfg, + stage_page_ids_per_tile=stage_page_ids_per_tile, + page_idx_kv=page_idx_kv, + seqlens_kv=kv_seqlens, + use_native_paged_kv=use_native_paged_kv, + block_tables=block_tables, + block_table_row_stride=block_table_row_stride, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + name="smemPageOffsetsKvV", + ) + sparse_kv_metadata0 = None + sparse_kv_metadata1 = None + sparse_softmax_metadata0 = None + sparse_softmax_metadata1 = None + if cfg.use_block_sparse: + prepared_route_layout = _BlockSparseRouteLayout.create( + kv_route_size=cfg.tile_size_kv, + kv_block_size=cfg.kv_block_size, + has_token_bits=cfg.use_kv_valid_bits, + route_metadata_capacity=0, + num_rows=1, + page_size=cfg.num_tokens_per_page if cfg.use_paged_kv else None, + ) + sparse_kv_metadata0 = SmemBlockSparseKvMetadataResource( + pipeline_config=None, + cfg=cfg, + inst_id=0, + route_metadata=sparse_route_metadata, + route_layout=prepared_route_layout, + tma_oob_origin=max_seq_len_kv, + name="smemBlockSparseKvMetadata0", + ) + sparse_kv_metadata1 = SmemBlockSparseKvMetadataResource( + pipeline_config=None, + cfg=cfg, + inst_id=1, + route_metadata=sparse_route_metadata, + route_layout=prepared_route_layout, + tma_oob_origin=max_seq_len_kv, + name="smemBlockSparseKvMetadata1", + ) + sparse_softmax_metadata0 = SmemBlockSparseSoftmaxMetadataResource( + pipeline_config=sparse_softmax_metadata0_cfg, + cfg=cfg, + inst_id=0, + route_metadata=sparse_route_metadata, + route_layout=prepared_route_layout, + name="smemBlockSparseSoftmaxMetadata0", + ) + sparse_softmax_metadata1 = SmemBlockSparseSoftmaxMetadataResource( + pipeline_config=sparse_softmax_metadata1_cfg, + cfg=cfg, + inst_id=1, + route_metadata=sparse_route_metadata, + route_layout=prepared_route_layout, + name="smemBlockSparseSoftmaxMetadata1", + ) + smem_kv = None + smem_k0 = None + smem_k1 = None + smem_v0 = None + smem_v1 = None + if use_per_inst_kv_resources: + smem_k0 = SmemKvTileResource( + pipeline_config=smem_k0_cfg, + cfg=cfg, + tma_desc_k=tma_desc_k, + tma_desc_v=tma_desc_v, + tma_desc_k_atom=tma_desc_k_atom, + tma_desc_v_atom=tma_desc_v_atom, + sparse_kv_metadata=sparse_kv_metadata0, + page_offsets_kv=smem_page_offsets, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + inst_id=0, + kv_kind=KV_KIND_K, + name="smemK0", + ) + smem_k1 = SmemKvTileResource( + pipeline_config=smem_k1_cfg, + cfg=cfg, + tma_desc_k=tma_desc_k, + tma_desc_v=tma_desc_v, + tma_desc_k_atom=tma_desc_k_atom, + tma_desc_v_atom=tma_desc_v_atom, + sparse_kv_metadata=sparse_kv_metadata1, + page_offsets_kv=smem_page_offsets, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + inst_id=1, + kv_kind=KV_KIND_K, + name="smemK1", + ) + smem_v0 = SmemKvTileResource( + pipeline_config=smem_v0_cfg, + cfg=cfg, + tma_desc_k=tma_desc_k, + tma_desc_v=tma_desc_v, + tma_desc_k_atom=tma_desc_k_atom, + tma_desc_v_atom=tma_desc_v_atom, + sparse_kv_metadata=sparse_kv_metadata0, + page_offsets_kv=smem_page_offsets_v or smem_page_offsets, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + inst_id=0, + kv_kind=KV_KIND_V, + name="smemV0", + ) + smem_v1 = SmemKvTileResource( + pipeline_config=smem_v1_cfg, + cfg=cfg, + tma_desc_k=tma_desc_k, + tma_desc_v=tma_desc_v, + tma_desc_k_atom=tma_desc_k_atom, + tma_desc_v_atom=tma_desc_v_atom, + sparse_kv_metadata=sparse_kv_metadata1, + page_offsets_kv=smem_page_offsets_v or smem_page_offsets, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + inst_id=1, + kv_kind=KV_KIND_V, + name="smemV1", + ) + else: + smem_kv = SmemKvResource( + pipeline_config=smem_kv_cfg, + cfg=cfg, + tma_desc_k=tma_desc_k, + tma_desc_v=tma_desc_v, + tma_desc_k_atom=tma_desc_k_atom, + tma_desc_v_atom=tma_desc_v_atom, + sparse_kv_metadata0=sparse_kv_metadata0, + sparse_kv_metadata1=sparse_kv_metadata1, + page_offsets_kv=smem_page_offsets, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + name="smemKv", + ) + + tmem_s0 = TmemSResource( + inst_id=0, + pipeline_config=tmem_s0_cfg, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_r=h_r, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + sync_barrier_id=0, + name="tmemS0", + ) + tmem_s1 = TmemSResource( + inst_id=1, + pipeline_config=tmem_s1_cfg, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_r=h_r, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + sync_barrier_id=1, + name="tmemS1", + ) + # Packed persistent QK derives the descriptor from Q's just-waited + # consumer stage, avoiding a routed HEAD-to-LOOP descriptor value across + # the guarded work-tile region. Fixed/static schedules keep their existing + # explicit descriptor route. + tmem_s0.q_ref = smem_q + tmem_s1.q_ref = smem_q + + smem_p0 = SmemPResource( + inst_id=0, + pipeline_config=smem_p0_cfg, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + use_variable_seqlens_kv=use_runtime_seqlens_kv, + name="smemP0", + ) + smem_p1 = SmemPResource( + inst_id=1, + pipeline_config=smem_p1_cfg, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + use_variable_seqlens_kv=use_runtime_seqlens_kv, + name="smemP1", + ) + + tmem_o = TmemOResource( + pipeline_config=tmem_o_cfg, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + name="tmemO", + ) + + tmem_softmax_local0 = TmemSoftmaxLocalResource( + inst_id=0, + pipeline_config=softmax_local0_cfg, + cfg=cfg, + name="tmemSoftmaxLocal0", + ) + tmem_softmax_local1 = TmemSoftmaxLocalResource( + inst_id=1, + pipeline_config=softmax_local1_cfg, + cfg=cfg, + name="tmemSoftmaxLocal1", + ) + tmem_stats_done0 = ( + TmemStatsDoneResource( + pipeline_config=stats_done0_cfg, + name="tmemStatsDone0", + ) + if stats_done0_cfg is not None + else None + ) + tmem_stats_done1 = ( + TmemStatsDoneResource( + pipeline_config=stats_done1_cfg, + name="tmemStatsDone1", + ) + if stats_done1_cfg is not None + else None + ) + tmem_softmax_global0 = TmemSoftmaxGlobalResource( + inst_id=0, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + sum_barrier_id=2, + name="tmemSoftmaxGlobal0", + ) + tmem_softmax_global1 = TmemSoftmaxGlobalResource( + inst_id=1, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + sum_barrier_id=3, + name="tmemSoftmaxGlobal1", + ) + tmem_softmax_order = ( + TmemSoftmaxOrderResource(cfg=cfg, name="tmemSoftmaxOrder") + if use_ordered_softmax_barrier + else None + ) + smem_p0.tmem_s_ref = tmem_s0 + smem_p1.tmem_s_ref = tmem_s1 + smem_p0.tmem_o_ref = tmem_o + smem_p1.tmem_o_ref = tmem_o + tmem_softmax_global0.p_ref = smem_p0 + tmem_softmax_global1.p_ref = smem_p1 + tmem_softmax_global0.tmem_s_ref = tmem_s0 + tmem_softmax_global1.tmem_s_ref = tmem_s1 + + tmem_corr0 = TmemCorrResource( + inst_id=0, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + o_ptr=o_ptr, + partial_o_ptr=partial_o_ptr, + partial_stats_ptr=partial_stats_ptr, + split_kv_counter_ptr=split_kv_counter_ptr, + attention_sinks_ptr=attention_sinks_ptr, + seqlens_kv=kv_seqlens, + max_seq_len_kv=corr_max_seq_len_kv, + num_heads_kv=num_heads_kv, + h_r=h_r, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + q_token_offset=q_token_offset, + seq_len_q=seq_len_q, + active_splits_kv=active_splits_kv, + static_full_split_prefix=static_full_split_prefix, + name="tmemCorr0", + ) + tmem_corr1 = TmemCorrResource( + inst_id=1, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + o_ptr=o_ptr, + partial_o_ptr=partial_o_ptr, + partial_stats_ptr=partial_stats_ptr, + split_kv_counter_ptr=split_kv_counter_ptr, + attention_sinks_ptr=attention_sinks_ptr, + seqlens_kv=kv_seqlens, + max_seq_len_kv=corr_max_seq_len_kv, + num_heads_kv=num_heads_kv, + h_r=h_r, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + q_token_offset=q_token_offset, + seq_len_q=seq_len_q, + active_splits_kv=active_splits_kv, + static_full_split_prefix=static_full_split_prefix, + name="tmemCorr1", + ) + tmem_corr1.smem_p0_ref = smem_p0 + tmem_corr1.smem_p1_ref = smem_p1 + tmem_corr0.tmem_o_ref = tmem_o + tmem_corr1.tmem_o_ref = tmem_o + tmem_corr0.softmax_local0_ref = tmem_softmax_local0 + tmem_corr1.softmax_local0_ref = tmem_softmax_local0 + if use_one_inst_qkv: + tmem_corr0.softmax_local1_ref = None + tmem_corr1.softmax_local1_ref = None + else: + tmem_corr0.softmax_local1_ref = tmem_softmax_local1 + tmem_corr1.softmax_local1_ref = tmem_softmax_local1 + + # ------------------------------------------------------------------ + # Domain computation + # ------------------------------------------------------------------ + # HEAD handles the first 2 K tiles, LOOP advances the staggered + # Qk/Pv steady-state, and TAIL drains the final V wave. For odd tile + # counts, the final wave still requires one additional loop iteration so + # that inst0 can process the last K/V tile. + local_kv_tiles = _compute_local_kv_tiles(cfg, total_kv_tiles) + loop_domain = _compute_decode_gen_loop_domain(local_kv_tiles, cfg.num_insts_kv) + + load_domain = loop_domain + mma_domain = loop_domain + softmax_domain = loop_domain + 1 + corr_domain = loop_domain + + # ------------------------------------------------------------------ + # Create tasks + # ------------------------------------------------------------------ + task_runtime_kwargs = { + "seqlens_kv": kv_seqlens, + "max_seq_len_kv": max_seq_len_kv, + "seq_len_q": seq_len_q, + "sparse_row_route_offsets": sparse_row_route_offsets, + "sparse_row_route_counts": sparse_row_route_counts, + "num_heads_kv": num_heads_kv, + } + if use_one_inst_qkv: + load_tasks = ( + create_load_task_one_inst_qkv( + smem_q, + smem_k0, + smem_v0, + work_queue, + schedule_token_throttle, + cfg, + domain=load_domain, + smem_page_offsets=smem_page_offsets, + domain_bias=0, + warp_idx=cfg.clc_load_warp_idx if use_clc_dynamic else None, + **task_runtime_kwargs, + ), + ) + elif use_per_inst_kv_resources: + if use_per_inst_block_sparse_load_tasks: + assert per_inst_block_sparse_load_topology is not None + load_warp_indices, _ = per_inst_block_sparse_load_topology + load_tasks = create_block_sparse_load_tasks_per_inst( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + work_queue, + schedule_token_throttle, + cfg, + domain=load_domain, + sparse_kv_metadata0=sparse_kv_metadata0, + sparse_kv_metadata1=sparse_kv_metadata1, + sparse_softmax_metadata0=sparse_softmax_metadata0, + sparse_softmax_metadata1=sparse_softmax_metadata1, + domain_bias=0, + warp_indices=load_warp_indices, + **task_runtime_kwargs, + ) + else: + load_tasks = ( + create_load_task_split_kv( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + work_queue, + schedule_token_throttle, + cfg, + domain=load_domain, + smem_page_offsets=smem_page_offsets, + smem_page_offsets_v=smem_page_offsets_v, + sparse_kv_metadata0=sparse_kv_metadata0, + sparse_kv_metadata1=sparse_kv_metadata1, + sparse_softmax_metadata0=sparse_softmax_metadata0, + sparse_softmax_metadata1=sparse_softmax_metadata1, + domain_bias=0, + warp_idx=cfg.clc_load_warp_idx if use_clc_dynamic else None, + **task_runtime_kwargs, + ), + ) + else: + load_tasks = ( + create_load_task( + smem_q, + smem_kv, + work_queue, + schedule_token_throttle, + smem_kv_reuse_credit, + cfg, + domain=load_domain, + domain_bias=0, + warp_idx=cfg.clc_load_warp_idx if use_clc_dynamic else None, + smem_page_offsets=smem_page_offsets, + sparse_kv_metadata0=sparse_kv_metadata0, + sparse_kv_metadata1=sparse_kv_metadata1, + sparse_softmax_metadata0=sparse_softmax_metadata0, + sparse_softmax_metadata1=sparse_softmax_metadata1, + **task_runtime_kwargs, + ), + ) + page_offsets_task = None + if use_dense_page_offsets: + page_offsets_warp_idx = cfg.page_offsets_warp_idx + if use_separate_kv_page_offset_resources: + page_offsets_task = create_page_offsets_task_split_kv( + smem_page_offsets, + smem_page_offsets_v, + work_queue, + cfg, + domain=load_domain, + domain_bias=0, + warp_idx=page_offsets_warp_idx, + num_warps=cfg.page_offsets_num_warps, + block_table_capacity=( + block_table_capacity if use_native_paged_kv else None + ), + **task_runtime_kwargs, + ) + else: + page_offsets_task_fn = ( + create_page_offsets_task_one_inst_qkv + if use_one_inst_qkv + else create_page_offsets_task + ) + page_offsets_task = page_offsets_task_fn( + smem_page_offsets, + work_queue, + cfg, + domain=load_domain, + domain_bias=0, + warp_idx=page_offsets_warp_idx, + num_warps=cfg.page_offsets_num_warps, + block_table_capacity=( + block_table_capacity if use_native_paged_kv else None + ), + **task_runtime_kwargs, + ) + if use_one_inst_qkv: + mma_task = create_mma_task_one_inst_qkv( + smem_q, + smem_k0, + smem_v0, + tmem_s0, + smem_p0, + tmem_o, + work_queue, + cfg, + tmem_stats_done=tmem_stats_done0, + domain=mma_domain, + domain_bias=0, + **task_runtime_kwargs, + ) + elif use_per_inst_kv_resources: + mma_task = create_mma_task_split_kv( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue, + cfg, + domain=mma_domain, + tmem_stats_done0=tmem_stats_done0, + tmem_stats_done1=tmem_stats_done1, + domain_bias=0, + **task_runtime_kwargs, + ) + else: + mma_task = create_mma_task( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue, + cfg, + domain=mma_domain, + domain_bias=0, + **task_runtime_kwargs, + ) + softmax0_task = create_softmax0_task( + tmem_s0, + tmem_softmax_local0, + smem_p0, + tmem_softmax_global0, + tmem_softmax_order, + sparse_softmax_metadata0, + work_queue, + cfg, + domain=softmax_domain, + domain_bias=1, + **task_runtime_kwargs, + ) + softmax1_task = None + if not use_one_inst_qkv: + softmax1_task = create_softmax1_task( + tmem_s1, + tmem_softmax_local1, + smem_p1, + tmem_softmax_global1, + tmem_softmax_order, + sparse_softmax_metadata1, + work_queue, + cfg, + domain=softmax_domain, + domain_bias=1, + **task_runtime_kwargs, + ) + if use_one_inst_qkv: + correction_task = create_correction_task_one_inst_qkv( + tmem_softmax_local0, + tmem_o, + tmem_corr0, + work_queue, + cfg, + tmem_stats_done=tmem_stats_done0, + domain=corr_domain, + domain_bias=0, + **task_runtime_kwargs, + ) + else: + correction_task = create_correction_task( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + work_queue, + smem_kv_reuse_credit, + cfg, + domain=corr_domain, + tmem_stats_done0=tmem_stats_done0, + tmem_stats_done1=tmem_stats_done1, + domain_bias=0, + **task_runtime_kwargs, + ) + if use_per_inst_block_sparse_load_tasks: + assert per_inst_block_sparse_load_topology is not None + _, residual_padding = per_inst_block_sparse_load_topology + padding_warp_ranges = ( + (residual_padding,) if residual_padding is not None else () + ) + else: + padding_warp_ranges = ( + (cfg.wg0_padding_warp_idx, cfg.wg0_padding_num_warps), + (cfg.wg1_padding_warp_idx, cfg.wg1_padding_num_warps), + (cfg.wg2_padding_warp_idx, cfg.wg2_padding_num_warps), + (cfg.wg3_padding_warp_idx, cfg.wg3_padding_num_warps), + ) + padding_tasks = [ + create_padding_task( + cfg, + work_queue, + warp_idx=warp_idx, + num_warps=num_warps, + ) + for warp_idx, num_warps in padding_warp_ranges + if num_warps > 0 + ] + scheduler_task = None + if use_clc_dynamic: + scheduler_task = create_scheduler_task(work_queue, schedule_token_throttle, cfg) + + task_list = [] + if page_offsets_task is not None: + task_list.append(page_offsets_task) + task_list.extend(load_tasks) + if use_one_inst_qkv and not use_clc_dynamic: + task_list.extend([correction_task, mma_task]) + task_list.append(softmax0_task) + else: + task_list.append(softmax0_task) + if softmax1_task is not None: + task_list.append(softmax1_task) + task_list.extend([correction_task, mma_task]) + if scheduler_task is not None: + task_list.append(scheduler_task) + task_list.extend(padding_tasks) + # ------------------------------------------------------------------ + # Resource dependency graph + # ------------------------------------------------------------------ + smem_kv_deps = [] + if smem_page_offsets is not None: + smem_kv_deps.append(smem_page_offsets) + smem_k_deps = list(smem_kv_deps) + smem_v_deps = ( + [smem_page_offsets_v] if smem_page_offsets_v is not None else list(smem_kv_deps) + ) + if use_one_inst_qkv: + resource_dependency_graph = { + smem_q: [], + smem_k0: smem_kv_deps, + smem_v0: smem_kv_deps, + tmem_s0: [smem_k0, smem_q], + smem_p0: [tmem_s0], + tmem_softmax_local0: [tmem_s0], + tmem_softmax_global0: [tmem_s0], + tmem_o: [smem_p0, smem_v0], + tmem_corr0: [tmem_softmax_local0, tmem_o], + } + elif use_per_inst_kv_resources: + resource_dependency_graph = { + **( + { + sparse_kv_metadata0: [], + sparse_kv_metadata1: [], + } + if sparse_kv_metadata0 is not None + else {} + ), + **( + { + sparse_softmax_metadata0: [sparse_kv_metadata0], + sparse_softmax_metadata1: [sparse_kv_metadata1], + } + if sparse_softmax_metadata0 is not None + else {} + ), + smem_q: [], + smem_k0: smem_k_deps + + ([sparse_kv_metadata0] if sparse_kv_metadata0 is not None else []), + smem_k1: smem_k_deps + + ([sparse_kv_metadata1] if sparse_kv_metadata1 is not None else []), + smem_v0: smem_v_deps + + ([sparse_kv_metadata0] if sparse_kv_metadata0 is not None else []), + smem_v1: smem_v_deps + + ([sparse_kv_metadata1] if sparse_kv_metadata1 is not None else []), + tmem_s0: [smem_k0, smem_q], + tmem_s1: [smem_k1, smem_q], + smem_p0: [tmem_s0], + smem_p1: [tmem_s1], + tmem_softmax_local0: [tmem_s0], + tmem_softmax_local1: [tmem_s1], + tmem_softmax_global0: [tmem_s0], + tmem_softmax_global1: [tmem_s1], + tmem_o: [smem_p0, smem_p1, smem_v0, smem_v1], + tmem_corr0: [tmem_softmax_local0, tmem_o], + tmem_corr1: [tmem_softmax_local0, tmem_softmax_local1, tmem_o], + } + else: + resource_dependency_graph = { + **( + { + sparse_kv_metadata0: [], + sparse_kv_metadata1: [], + sparse_softmax_metadata0: [sparse_kv_metadata0], + sparse_softmax_metadata1: [sparse_kv_metadata1], + } + if sparse_kv_metadata0 is not None + else {} + ), + smem_q: [], + smem_kv: smem_kv_deps + + ( + [sparse_kv_metadata0, sparse_kv_metadata1] + if sparse_kv_metadata0 is not None + else [] + ), + tmem_s0: [smem_kv, smem_q], + tmem_s1: [smem_kv, smem_q], + smem_p0: [tmem_s0], + smem_p1: [tmem_s1], + tmem_softmax_local0: [tmem_s0], + tmem_softmax_local1: [tmem_s1], + tmem_softmax_global0: [tmem_s0], + tmem_softmax_global1: [tmem_s1], + tmem_o: [smem_p0, smem_p1, smem_kv], + tmem_corr0: [tmem_softmax_local0, tmem_o], + tmem_corr1: [tmem_softmax_local0, tmem_softmax_local1, tmem_o], + } + if sparse_softmax_metadata0 is not None: + resource_dependency_graph[smem_p0].append(sparse_softmax_metadata0) + resource_dependency_graph[tmem_softmax_local0].append(sparse_softmax_metadata0) + resource_dependency_graph[tmem_softmax_global0].append(sparse_softmax_metadata0) + assert sparse_softmax_metadata1 is not None + resource_dependency_graph[smem_p1].append(sparse_softmax_metadata1) + resource_dependency_graph[tmem_softmax_local1].append(sparse_softmax_metadata1) + resource_dependency_graph[tmem_softmax_global1].append(sparse_softmax_metadata1) + if tmem_stats_done0 is not None: + resource_dependency_graph[tmem_s0].append(tmem_stats_done0) + resource_dependency_graph[tmem_stats_done0] = [tmem_softmax_local0] + if tmem_stats_done1 is not None: + resource_dependency_graph[tmem_s1].append(tmem_stats_done1) + resource_dependency_graph[tmem_stats_done1] = [tmem_softmax_local1] + if cutlass.const_expr(use_ordered_softmax_barrier): + resource_dependency_graph[tmem_softmax_order] = [tmem_s0] + resource_dependency_graph[smem_p1] = [ + *resource_dependency_graph[smem_p1], + tmem_softmax_order, + ] + if smem_page_offsets is not None: + resource_dependency_graph[smem_page_offsets] = [] + if smem_page_offsets_v is not None: + resource_dependency_graph[smem_page_offsets_v] = [] + if work_queue is not None: + for deps in resource_dependency_graph.values(): + deps.append(work_queue) + resource_dependency_graph[work_queue] = ( + [work_queue, schedule_token_throttle] + if schedule_token_throttle is not None + else ([work_queue] if use_clc_dynamic else []) + ) + if schedule_token_throttle is not None: + resource_dependency_graph[schedule_token_throttle] = [work_queue] + if smem_kv_reuse_credit is not None: + # A self-edge models the one-slot ownership token: Load produces it + # for the current tile and Correction consumes it before the next Load. + resource_dependency_graph[smem_kv_reuse_credit] = [smem_kv_reuse_credit] + dma_consumer_release_labels: dict[ + tuple[MemoryResource, MemoryResource], set[str] + ] = {} + if smem_page_offsets is not None: + if use_one_inst_qkv: + dma_consumer_release_labels.update( + { + (smem_page_offsets, smem_k0): {"read_offsets_k0"}, + (smem_page_offsets, smem_v0): {"read_offsets_v0"}, + } + ) + elif use_per_inst_kv_resources: + if smem_page_offsets_v is not None: + dma_consumer_release_labels.update( + { + (smem_page_offsets, smem_k0): {"read_offsets_k0"}, + (smem_page_offsets, smem_k1): {"read_offsets_k1"}, + (smem_page_offsets_v, smem_v0): {"read_offsets_v0"}, + (smem_page_offsets_v, smem_v1): {"read_offsets_v1"}, + } + ) + else: + dma_consumer_release_labels.update( + { + (smem_page_offsets, smem_k0): {"read_offsets_k0"}, + (smem_page_offsets, smem_k1): {"read_offsets_k1"}, + (smem_page_offsets, smem_v0): {"read_offsets_v0"}, + (smem_page_offsets, smem_v1): {"read_offsets_v1"}, + } + ) + else: + dma_consumer_release_labels[(smem_page_offsets, smem_kv)] = { + "cache_page_ids" if cfg.num_head_dim_stages_kv > 1 else "read_offsets" + } + if smem_kv is not None: + dma_consumer_release_labels.update( + { + (smem_kv, tmem_s0): {"k_desc_0"}, + (smem_kv, tmem_s1): {"k_desc_1"}, + (smem_kv, tmem_o): {"v_desc_0", "v_desc_1"}, + } + ) + + # ------------------------------------------------------------------ + # SMEM / TMEM allocators + # ------------------------------------------------------------------ + smem_allocator = SmemAllocator() + if work_queue is not None: + smem_allocator.add_resource(work_queue) + if schedule_token_throttle is not None: + smem_allocator.add_resource(schedule_token_throttle) + if smem_kv_reuse_credit is not None: + smem_allocator.add_resource(smem_kv_reuse_credit) + smem_allocator.add_resource(smem_q) + if smem_page_offsets is not None: + smem_allocator.add_resource(smem_page_offsets) + if smem_page_offsets_v is not None: + smem_allocator.add_resource(smem_page_offsets_v) + if sparse_kv_metadata0 is not None: + smem_allocator.add_resource(sparse_kv_metadata0) + smem_allocator.add_resource(sparse_kv_metadata1) + if sparse_softmax_metadata0 is not None: + smem_allocator.add_resource(sparse_softmax_metadata0) + smem_allocator.add_resource(sparse_softmax_metadata1) + if use_one_inst_qkv: + smem_allocator.add_resource(smem_k0) + smem_allocator.add_resource(smem_v0) + elif use_per_inst_kv_resources: + smem_allocator.add_resource(smem_k0) + smem_allocator.add_resource(smem_k1) + smem_allocator.add_resource(smem_v0) + smem_allocator.add_resource(smem_v1) + else: + smem_allocator.add_resource(smem_kv) + smem_allocator.add_resource(smem_p0) + if not use_one_inst_qkv: + smem_allocator.add_resource(smem_p1) + smem_allocator.add_resource(tmem_s0) + if not use_one_inst_qkv: + smem_allocator.add_resource(tmem_s1) + smem_allocator.add_resource(tmem_o) + smem_allocator.add_resource(tmem_softmax_local0) + if not use_one_inst_qkv: + smem_allocator.add_resource(tmem_softmax_local1) + smem_allocator.add_resource(tmem_softmax_global0) + if not use_one_inst_qkv: + smem_allocator.add_resource(tmem_softmax_global1) + smem_allocator.add_resource(tmem_corr0) + if not use_one_inst_qkv: + smem_allocator.add_resource(tmem_corr1) + if cfg.tile_size_kv == 256: + # KV256 direct-output correction rotates one compact 35,840-byte + # payload through the shared 192-KiB K/V ring. Split-KV retains the + # fixed full exchange. Neither path increases the CTA SMEM footprint. + smem_allocator.add_alias_group( + [ + smem_kv.get_smem_requirements(), + tmem_corr1.get_smem_requirements(), + ] + ) + smem_allocator.add_tmem_ptr( + SmemAllocation("fmha_tmem_ptr_i32", dtype=cutlass.Int32, alignment=4) + ) + smem_allocator.compute_layout() + + tmem_allocator = TmemAllocator() + if cfg.use_keeps_mma_ab: + if use_one_inst_qkv: + tmem_allocator.add_resource(tmem_s0) + else: + # Build the two instruction-local phases from the resources that + # actually use TMEM. Depending on the profile, P can overlay S + # or live in SMEM, and stats can be standalone TMEM or an SMEM + # handoff. Empty resources must not become scheduler aliases. + for tmem_s, p in ( + (tmem_s0, smem_p0), + (tmem_s1, smem_p1), + ): + p_requirements = p.get_tmem_requirements() + if p_requirements: + tmem_allocator.add_alias_group( + [tmem_s.get_tmem_requirements(), p_requirements] + ) + else: + tmem_allocator.add_resource(tmem_s) + # Register standalone stats only after both S/P phases, preserving + # the established allocation order. SMEM-backed stats are a no-op. + tmem_allocator.add_resource(tmem_softmax_local0) + tmem_allocator.add_resource(tmem_softmax_local1) + else: + tmem_allocator.add_resource(tmem_s0) + tmem_allocator.add_resource(tmem_s1) + tmem_allocator.add_resource(tmem_softmax_local0) + tmem_allocator.add_resource(tmem_softmax_local1) + tmem_allocator.add_resource(tmem_o) + tmem_allocator.compute_layout() + if cfg.use_keeps_mma_ab and not use_one_inst_qkv and not cfg.uses_tmem_p: + # Two-inst Keeps with SMEM P (currently Q64) must retain the historical + # standalone-first S/O layout after unused stats aliases are removed. + s0_alloc = tmem_s0.get_tmem_requirements()[0] + s1_alloc = tmem_s1.get_tmem_requirements()[0] + o_alloc = tmem_o.get_tmem_requirements()[0] + stats0_requirements = tmem_softmax_local0.get_tmem_requirements() + stats1_requirements = tmem_softmax_local1.get_tmem_requirements() + if stats0_requirements: + assert len(stats0_requirements) == len(stats1_requirements) == 1 + stats0_requirements[0].offset = 0 + stats1_requirements[0].offset = cfg.tmem_stats_cols + o_alloc.offset = 2 * cfg.tmem_stats_cols + else: + assert not stats1_requirements + o_alloc.offset = 0 + s0_alloc.offset = o_alloc.offset + o_alloc.num_columns + s1_alloc.offset = s0_alloc.offset + cfg.tmem_s_cols + assert s1_alloc.offset + cfg.tmem_s_cols == cfg.tmem_total_cols + elif cfg.uses_two_inst_tmem_p: + # P reuses each independent S region after Softmax consumes QK. Stats + # are either standalone TMEM or an SMEM handoff. + s0_alloc = tmem_s0.get_tmem_requirements()[0] + s1_alloc = tmem_s1.get_tmem_requirements()[0] + p0_alloc = smem_p0.get_tmem_requirements()[0] + p1_alloc = smem_p1.get_tmem_requirements()[0] + o_alloc = tmem_o.get_tmem_requirements()[0] + if cfg.tile_size_kv == 256: + # KV256 keeps O in the low 256 columns and overlays packed P on + # each S region from its first column. Softmax streams K32 + # fragments in order, so every 16-column P store only overwrites + # scores that have already been consumed. Starting P after the + # nominal stats columns would instead clobber the next unread S + # fragment; KV256 keeps its softmax stats in SMEM. + o_alloc.offset = 0 + s0_alloc.offset = 2 * cfg.tmem_o_stage_cols + s1_alloc.offset = s0_alloc.offset + cfg.tmem_s_cols + p0_alloc.offset = s0_alloc.offset + p1_alloc.offset = s1_alloc.offset + else: + # Re-state the intended phase offsets after layout so every + # resource observes the same S/P alias. Stats remain standalone + # when the whole allocation fits; otherwise they use SMEM. Keep + # the historical gap before P so this modeling cleanup does not + # change runtime addresses. + s0_alloc.offset = 0 + s1_alloc.offset = cfg.tmem_s_cols + p0_alloc.offset = ( + 0 if cfg.keeps_separates_tmem_s_and_stats else cfg.tmem_stats_cols + ) + p1_alloc.offset = cfg.tmem_s_cols + ( + 0 if cfg.keeps_separates_tmem_s_and_stats else cfg.tmem_stats_cols + ) + if cfg.keeps_separates_tmem_s_and_stats: + stats0_alloc = tmem_softmax_local0.get_tmem_requirements()[0] + stats1_alloc = tmem_softmax_local1.get_tmem_requirements()[0] + stats0_alloc.offset = 2 * cfg.tmem_s_cols + stats1_alloc.offset = 2 * cfg.tmem_s_cols + cfg.tmem_stats_cols + o_alloc.offset = 2 * (cfg.tmem_s_cols + cfg.tmem_stats_cols) + else: + o_alloc.offset = 2 * cfg.tmem_s_cols + expected_p_cols = cfg.tmem_p_cols_per_inst + assert p0_alloc.num_columns == p1_alloc.num_columns == expected_p_cols + assert s0_alloc.offset <= p0_alloc.offset + assert p0_alloc.offset + expected_p_cols <= s0_alloc.offset + cfg.tmem_s_cols + assert s1_alloc.offset <= p1_alloc.offset + assert p1_alloc.offset + expected_p_cols <= s1_alloc.offset + cfg.tmem_s_cols + assert ( + max( + s0_alloc.offset + cfg.tmem_s_cols, + s1_alloc.offset + cfg.tmem_s_cols, + o_alloc.offset + cfg.tmem_o_stage_cols * cfg.o_stages, + ) + == cfg.tmem_total_cols + ) + if use_one_inst_qkv: + tmem_s_alloc = tmem_s0.get_tmem_requirements()[0] + # One-inst Keeps also transports stats through SMEM, so there is no + # TMEM stats allocation to alias with S. + assert cfg.keeps_stats_via_smem + assert cfg.tmem_stats_cols + cfg.tmem_p_cols <= cfg.tmem_s_cols + smem_p0.get_tmem_requirements()[0].offset = ( + tmem_s_alloc.offset + cfg.tmem_stats_cols + ) + + eager_init_resources = ( + [tmem_corr0] if use_one_inst_qkv else [tmem_corr0, tmem_corr1] + ) + if smem_kv_reuse_credit is not None: + # Initialize the persistent ring cursor under the same CTA-wide fence + # and barrier used by other manually managed SMEM control state. + eager_init_resources.append(smem_kv_reuse_credit) + if cfg.tile_size_kv == 256: + # KV256's TMEM P operands use one-way per-fragment ready barriers. + # Initialize them beside correction's manually managed SMEM state. + eager_init_resources.extend([smem_p0, smem_p1]) + + return ( + task_list, + resource_dependency_graph, + dma_consumer_release_labels, + smem_allocator, + tmem_allocator, + eager_init_resources, + ) + + +def _round_up_tmem_columns(num_columns: int) -> int: + """tcgen05_alloc requires a power-of-two column count in [32, 512].""" + return max(32, 1 << (num_columns - 1).bit_length()) + + +def _has_unmodeled_tmem_p_alias_protocol(cfg: FmhaDecodeConfig) -> bool: + """Whether exhaustive TS checking would report a known false P/S race. + + The staged D256 path selects one of two physical P/S stages at runtime. + Static KV256 instead orders streamed P fragments with private mbarriers and + reuses the matching TmemO-full barrier as the next-QK overwrite credit. + Those intra-work protocols are below TaskManager's resource transitions, + so its allocation-level checker cannot prove them. Persistent KV256 has + enough task-level ordering for the checker and remains covered. + """ + return cfg.uses_staged_one_inst_tmem_p or ( + cfg.streams_tmem_p_fragments and not cfg.use_persistent_scheduler + ) + + +def build_decode_task_manager( + cfg: FmhaDecodeConfig, + seq_len_kv: int = 2048, + batch_size: int = 8, + num_heads_kv: int = 8, + verbose: bool = True, + skip_validation: bool = False, + exhaustive_deadlock_race_check: bool = True, +) -> TaskManager: + """Build and validate the decode TS TaskManager (pure Python, no GPU). + + Parameters + ---------- + cfg : FmhaDecodeConfig + Pre-built configuration. Build it externally with + ``make_decode_config`` so callers can apply their own + overrides without monkey-patching the kernel module. + seq_len_kv : int + KV sequence length. + exhaustive_deadlock_race_check : bool + Run exhaustive interleaving validation where TaskManager can model the + complete synchronization protocol. Structural checks always run. + + Returns + ------- + TaskManager + Structurally validated task manager, exhaustively checked when the + profile does not use a private TMEM-P alias protocol. + """ + effective_seq_len_kv = _configure_static_sliding_window(cfg, seq_len_kv) + total_kv_tiles = _compute_total_kv_tiles(effective_seq_len_kv, cfg.tile_size_kv) + cfg.total_kv_tiles = total_kv_tiles + + ( + task_list, + resource_dependency_graph, + dma_consumer_release_labels, + smem_allocator, + tmem_allocator, + _eager_init_resources, + ) = _build_decode_gen_schedule( + cfg, + total_kv_tiles, + tile_sched_params=None, + num_heads_kv=Int32(num_heads_kv), + ) + + tm = TaskManager( + tasks=task_list, + resource_dependency_graph=resource_dependency_graph, + dma_consumer_release_labels=dma_consumer_release_labels, + smem_allocator=smem_allocator, + tmem_allocator=tmem_allocator, + verbose=verbose, + skip_validation=skip_validation, + exhaustive_deadlock_race_check=( + exhaustive_deadlock_race_check + and not _has_unmodeled_tmem_p_alias_protocol(cfg) + ), + ) + + return tm + + +# ===================================================================== +# GPU Kernel +# ===================================================================== + + +@cute.jit +def _run_decode_gen_active( + tma_desc_q: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k_atom: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v_atom: cutlass.GridConstant[cuda.TensorMap], + o_iter: cute.Pointer, + g_s_k: Int32, + g_h_k: Int32, + g_scale_s_log2_e: Float32, + g_output_scale: Float32, + g_seqlens_kv: cute.Pointer, + g_cu_seqlens_q: cute.Pointer, + g_page_idx_kv: cute.Pointer, + g_partial_o: cute.Pointer, + g_partial_stats: cute.Pointer, + g_split_kv_counter: cute.Pointer, + g_attention_sinks: cute.Pointer, + g_h_r: Int32, + q_group_idx: Int32, + h_k_idx: Int32, + b_idx: Int32, + q_token_offset: Int32, + seq_len_q: Int32, + active_splits_kv: Int32, + static_full_split_prefix: cutlass.Constexpr[bool], + tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + seq_len_kv: cutlass.Constexpr[int] = 2048, + use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, + use_native_paged_kv: cutlass.Constexpr[bool] = False, + use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, + g_block_tables: cute.Pointer | None = None, + g_block_table_capacity: Int32 | None = None, + g_block_table_row_stride: Int64 | None = None, + g_sparse_row_route_offsets: cute.Pointer | None = None, + g_sparse_row_route_counts: cute.Pointer | None = None, + g_sparse_route_metadata: cute.Pointer | None = None, +) -> None: + """Run the complete decode body for one runtime-valid Q tile. + + Builds resources/tasks via `_build_decode_gen_schedule` (shared with the + validation path), then owns the matched TMA prefetch, SMEM/TMEM setup, + TaskManager execution, and TMEM teardown lifecycle. Keeping that lifecycle + in a void JIT helper lets the kernel wrapper omit it entirely for an + overlaunched packed-Q tile without threading TaskManager state through a + dynamic branch. + """ + + WARP_SIZE = 32 + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + bias_static_sliding_kv_tma = ( + cfg.use_sliding_window_causal + and cfg.max_seq_len_q == 1 + and not cfg.use_paged_kv + and not cfg.use_split_kv + and not use_variable_seqlens_kv + ) + effective_seq_len_kv = _configure_static_sliding_window( + cfg, seq_len_kv, bias_static_sliding_kv_tma + ) + total_kv_tiles = _compute_total_kv_tiles(effective_seq_len_kv, cfg.tile_size_kv) + cfg.total_kv_tiles = total_kv_tiles + use_runtime_seqlens_kv = use_variable_seqlens_kv or ( + use_native_paged_kv and not use_static_native_seqlens_kv + ) + runtime_seqlens_kv = ( + g_seqlens_kv if cutlass.const_expr(use_runtime_seqlens_kv) else None + ) + runtime_max_seq_len_kv = ( + g_s_k + if cutlass.const_expr(use_runtime_seqlens_kv) + else Int32(cfg.static_seq_len_kv) + ) + use_clc_dynamic_scheduler = cfg.use_persistent_scheduler + + # Prefetch TMA + init_warp = 1 + if warp_idx == init_warp: + prims.prefetch_tensormap(tma_desc_q.get_ptr()) + prims.prefetch_tensormap(tma_desc_k.get_ptr()) + prims.prefetch_tensormap(tma_desc_v.get_ptr()) + if cutlass.const_expr( + cfg.use_block_sparse + and _block_sparse_kv_atom_size(cfg.kv_block_size) == 64 + and ( + cfg.tile_size_kv == 256 + or not _prepared_kv_routes_are_block_aligned( + cfg.kv_block_size, + cfg.tile_size_kv, + ) + ) + ): + # KV256 always issues semantic KV64 atoms. KV128 needs this second + # descriptor only for non-aligned coarse routes. + prims.prefetch_tensormap(tma_desc_k_atom.get_ptr()) + prims.prefetch_tensormap(tma_desc_v_atom.get_ptr()) + init_warp += 1 + + clc_response_ptr = None + if cutlass.const_expr(use_clc_dynamic_scheduler): + clc_response_ptr = cute.arch.alloc_smem( + cutlass.Int128, _PERSISTENT_SCHEDULE_TOKEN_STAGES + ) + + q_output_rows = g_h_r + if cutlass.const_expr(cfg.max_seq_len_q > 1): + q_output_rows = g_h_r * Int32(cfg.max_seq_len_q) + + ( + task_list, + dep_graph, + dma_consumer_release_labels, + smem_allocator, + tmem_allocator, + eager_init_resources, + ) = _build_decode_gen_schedule( + cfg, + total_kv_tiles, + scale_softmax_log2=g_scale_s_log2_e, + o_ptr=o_iter, + output_scale=g_output_scale, + partial_o_ptr=g_partial_o, + partial_stats_ptr=g_partial_stats, + split_kv_counter_ptr=g_split_kv_counter, + attention_sinks_ptr=g_attention_sinks, + seqlens_kv=runtime_seqlens_kv, + cu_seqlens_q=g_cu_seqlens_q, + max_seq_len_kv=runtime_max_seq_len_kv, + corr_max_seq_len_kv=seq_len_kv, + num_heads_kv=g_h_k, + h_r=q_output_rows, + tma_desc_q=tma_desc_q.get_ptr(), + tma_desc_k=tma_desc_k.get_ptr(), + tma_desc_v=tma_desc_v.get_ptr(), + tma_desc_k_atom=tma_desc_k_atom.get_ptr(), + tma_desc_v_atom=tma_desc_v_atom.get_ptr(), + page_idx_kv=g_page_idx_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + q_token_offset=q_token_offset, + seq_len_q=seq_len_q, + active_splits_kv=active_splits_kv, + static_full_split_prefix=static_full_split_prefix, + tile_sched_params=tile_sched_params, + clc_response_ptr=clc_response_ptr, + use_variable_seqlens_kv=use_variable_seqlens_kv, + use_native_paged_kv=use_native_paged_kv, + use_static_native_seqlens_kv=use_static_native_seqlens_kv, + block_tables=g_block_tables, + block_table_capacity=g_block_table_capacity, + block_table_row_stride=g_block_table_row_stride, + sparse_row_route_offsets=g_sparse_row_route_offsets, + sparse_row_route_counts=g_sparse_row_route_counts, + sparse_route_metadata=g_sparse_route_metadata, + ) + + smem_allocator.allocate() + + tmem_cols = tmem_allocator.total_tmem_columns + tmem_alloc_cols = _round_up_tmem_columns(tmem_cols) + tmem_ptr_alloc = smem_allocator.tmem_ptr_alloc + assert tmem_ptr_alloc is not None + tmem_ptr_i32 = smem_allocator.get(tmem_ptr_alloc) + if warp_idx == init_warp: + prims.tcgen05_alloc(tmem_ptr_i32, Int32(tmem_alloc_cols)) + prims.tcgen05_relinquish_alloc_permit() + init_warp += 1 + + task_manager = TaskManager( + tasks=task_list, + resource_dependency_graph=dep_graph, + dma_consumer_release_labels=dma_consumer_release_labels, + skip_validation=True, + verbose=False, + exhaustive_deadlock_race_check=not _has_unmodeled_tmem_p_alias_protocol(cfg), + smem_allocator=smem_allocator, + tmem_allocator=tmem_allocator, + ) + + task_manager.setup_resources_and_tasks() + resource_context = ResourceContext( + smem_base=smem_allocator.smem_base, + tmem_ptr_i32=tmem_ptr_i32, + ) + for resource in eager_init_resources: + # Materialize manually managed resource state before TS tasks start. + # This covers correction's cluster transaction barriers and KV256's + # per-fragment P-ready barriers. + resource.create_function_variables(resource_context) + # Ensure every CTA thread observes initialized mbarriers and SMEM resource + # bases before any task body can use them. + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + if cutlass.const_expr(cfg.supports_cluster_smem_reduction): + # Cluster-wide visibility point for the per-CTA transaction mbarriers. + # After this, peers may safely signal an owner CTA's mbarrier via mapa. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + task_manager.run() + + if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): + # Publish the dependent reducer only after this producer CTA has made + # its normalized partial O and log2-LSE stores visible. + thread_idx, _, _ = cute.arch.thread_idx() + if thread_idx == Int32(0): + prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) + + # Every peer async-stores into an owner CTA's distributed SMEM and charges + # the bytes to that owner's transaction mbarrier. Producer-only CTAs are + # never remote-SMEM targets, so they may retire after their store issues; + # each owner remains resident until its mbarrier completes and its local + # reduction consumes all delivered partials. A final cluster rendezvous + # would therefore only make completed producers wait for the owners. + + # ── TMEM cleanup ── + dealloc_barrier_id = 13 + if cutlass.const_expr(cfg.use_keeps_mma_ab): + # Keeps aliases score/stat TMEM columns across task warp groups. Wait + # for every task's producer tail before the correction owner warp + # releases the allocation; a correction-only barrier can race those + # tails and leave tcgen05.dealloc waiting on live TMEM users. + prims.barrier_cta_sync(dealloc_barrier_id) + if ( + warp_idx >= cfg.correction_warp_idx + and warp_idx < cfg.correction_warp_idx + cfg.correction_num_warps + ): + tidx, _, _ = cute.arch.thread_idx() + correction_thread_idx = tidx - cfg.correction_warp_idx * WARP_SIZE + if correction_thread_idx < WARP_SIZE: + tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) + prims.tcgen05_dealloc(tmem_ptr, Int32(tmem_alloc_cols)) + else: + if ( + warp_idx >= cfg.correction_warp_idx + and warp_idx < cfg.correction_warp_idx + cfg.correction_num_warps + ): + prims.barrier_cta_sync( + dealloc_barrier_id, + thread_count=cfg.correction_num_warps * WARP_SIZE, + ) + tidx, _, _ = cute.arch.thread_idx() + correction_thread_idx = tidx - cfg.correction_warp_idx * WARP_SIZE + if correction_thread_idx < WARP_SIZE: + tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) + prims.tcgen05_dealloc(tmem_ptr, Int32(tmem_alloc_cols)) + + +@cute.jit +def _run_decode_gen_inactive_cluster_rank() -> None: + """Join cluster initialization, then retire an inactive physical split rank.""" + # The physical cluster remains configured-max sized. Initialize a local + # zero-traffic barrier and join the same cluster visibility point as active + # ranks. Active peers never address this rank after runtime contraction. + inactive_mbarrier = cutlass.Array( + cutlass.Int64, + 1, + space=cutlass.AddressSpace.smem, + alignment=8, + ) + thread_idx, _, _ = cute.arch.thread_idx() + if thread_idx == Int32(0): + prims.mbarrier_init(inactive_mbarrier.data_ptr(), 1) + prims.fence_mbarrier_init() + prims.barrier_cta_sync(0) + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + + +@cute.jit +def _signal_padded_pdl_producer(cfg: cutlass.Constexpr[FmhaDecodeConfig]) -> None: + """Release the dependent reducer launch from a zero-work producer CTA.""" + if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): + thread_idx, _, _ = cute.arch.thread_idx() + if thread_idx == Int32(0): + prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) + + +@cute.jit +def _run_decode_gen_runtime_prefix( + tma_desc_q: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k_atom: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v_atom: cutlass.GridConstant[cuda.TensorMap], + o_iter: cute.Pointer, + g_s_k: Int32, + g_h_k: Int32, + g_scale_s_log2_e: Float32, + g_output_scale: Float32, + g_seqlens_kv: cute.Pointer, + g_cu_seqlens_q: cute.Pointer, + g_page_idx_kv: cute.Pointer, + g_partial_o: cute.Pointer, + g_partial_stats: cute.Pointer, + g_split_kv_counter: cute.Pointer, + g_attention_sinks: cute.Pointer, + g_h_r: Int32, + q_group_cta_idx: Int32, + q_group_idx: Int32, + h_k_idx: Int32, + b_idx: Int32, + q_token_offset: Int32, + seq_len_q: Int32, + q_token_base: Int32, + tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + seq_len_kv: cutlass.Constexpr[int], + use_variable_seqlens_kv: cutlass.Constexpr[bool], + use_native_paged_kv: cutlass.Constexpr[bool], + use_static_native_seqlens_kv: cutlass.Constexpr[bool], + g_block_tables: cute.Pointer | None, + g_block_table_capacity: Int32 | None, + g_block_table_row_stride: Int64 | None, + g_sparse_row_route_offsets: cute.Pointer | None, + g_sparse_row_route_counts: cute.Pointer | None, + g_sparse_route_metadata: cute.Pointer | None, +) -> None: + """Run the general runtime split-prefix producer or retire its suffix.""" + + split_is_active = cutlass.Boolean(True) + active_splits_kv = Int32(1) + if cutlass.const_expr(cfg.use_split_kv): + runtime_seq_len_kv = g_s_k + if cutlass.const_expr( + use_variable_seqlens_kv + or (use_native_paged_kv and not use_static_native_seqlens_kv) + ): + runtime_seq_len_kv = Int32(g_seqlens_kv[b_idx]) + active_splits_kv = _runtime_active_splits_kv( + cfg, + runtime_seq_len_kv, + seq_len_q, + q_token_base, + ) + if cutlass.const_expr(not cfg.use_separate_reduction_kernel): + # Preserve one neutral producer for fused empty-K semantics. + active_splits_kv = cute.math.max(active_splits_kv, Int32(1)) + split_idx = q_group_cta_idx % Int32(cfg.splits_kv) + split_is_active = split_idx < active_splits_kv + + if cutlass.const_expr(cfg.supports_cluster_smem_reduction): + if split_is_active: + _run_decode_gen_active( + tma_desc_q, + tma_desc_k, + tma_desc_v, + tma_desc_k_atom, + tma_desc_v_atom, + o_iter, + g_s_k, + g_h_k, + g_scale_s_log2_e, + g_output_scale, + g_seqlens_kv, + g_cu_seqlens_q, + g_page_idx_kv, + g_partial_o, + g_partial_stats, + g_split_kv_counter, + g_attention_sinks, + g_h_r, + q_group_idx, + h_k_idx, + b_idx, + q_token_offset, + seq_len_q, + active_splits_kv, + False, + tile_sched_params, + cfg, + seq_len_kv, + use_variable_seqlens_kv, + use_native_paged_kv, + use_static_native_seqlens_kv, + g_block_tables, + g_block_table_capacity, + g_block_table_row_stride, + g_sparse_row_route_offsets, + g_sparse_row_route_counts, + g_sparse_route_metadata, + ) + else: + _run_decode_gen_inactive_cluster_rank() + else: + if split_is_active: + _run_decode_gen_active( + tma_desc_q, + tma_desc_k, + tma_desc_v, + tma_desc_k_atom, + tma_desc_v_atom, + o_iter, + g_s_k, + g_h_k, + g_scale_s_log2_e, + g_output_scale, + g_seqlens_kv, + g_cu_seqlens_q, + g_page_idx_kv, + g_partial_o, + g_partial_stats, + g_split_kv_counter, + g_attention_sinks, + g_h_r, + q_group_idx, + h_k_idx, + b_idx, + q_token_offset, + seq_len_q, + active_splits_kv, + False, + tile_sched_params, + cfg, + seq_len_kv, + use_variable_seqlens_kv, + use_native_paged_kv, + use_static_native_seqlens_kv, + g_block_tables, + g_block_table_capacity, + g_block_table_row_stride, + g_sparse_row_route_offsets, + g_sparse_row_route_counts, + g_sparse_route_metadata, + ) + else: + _signal_padded_pdl_producer(cfg) + + +@cute.kernel +def decode_gen_kernel( + tma_desc_q: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k_atom: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v_atom: cutlass.GridConstant[cuda.TensorMap], + o_iter: cute.Pointer, + g_s_k: Int32, + g_h_k: Int32, + g_scale_s_log2_e: Float32, + g_output_scale: Float32, + g_seqlens_kv: cute.Pointer, + g_cu_seqlens_q: cute.Pointer, + g_page_idx_kv: cute.Pointer, + g_partial_o: cute.Pointer, + g_partial_stats: cute.Pointer, + g_split_kv_counter: cute.Pointer, + g_attention_sinks: cute.Pointer, + g_h_r: Int32, + tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + seq_len_kv: cutlass.Constexpr[int] = 2048, + use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, + use_native_paged_kv: cutlass.Constexpr[bool] = False, + use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, + g_block_tables: cute.Pointer | None = None, + g_block_table_capacity: Int32 | None = None, + g_block_table_row_stride: Int64 | None = None, + g_sparse_row_route_offsets: cute.Pointer | None = None, + g_sparse_row_route_counts: cute.Pointer | None = None, + g_sparse_route_metadata: cute.Pointer | None = None, + static_full_split_prefix: cutlass.Constexpr[bool] = False, +) -> None: + """Dispatch one static Q/split tile and drain padded launch slots safely.""" + q_group_cta_idx, h_k_idx, b_idx = cute.arch.block_idx() + q_group_idx = q_group_cta_idx + if cutlass.const_expr(cfg.use_split_kv): + # Grid coordinates and scratch linearization retain configured fanout. + q_group_idx = q_group_cta_idx // Int32(cfg.splits_kv) + + q_token_offset = Int32(0) + seq_len_q = Int32(cfg.max_seq_len_q) + q_token_base = Int32(0) + q_tile_is_active = cutlass.Boolean(True) + if cutlass.const_expr(not cfg.use_persistent_scheduler): + if cutlass.const_expr(cfg.use_variable_seqlens_q): + q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) + q_token_base = _q_group_token_base(cfg, q_group_idx) + q_tile_is_active = q_token_base < seq_len_q + + # Persistent block coordinates identify physical workers rather than + # logical tiles; their WorkQueue owns the equivalent Q predicate. Split-KV + # and persistent scheduling are mutually exclusive in supported profiles. + if q_tile_is_active: + if cutlass.const_expr(cfg.use_split_kv and static_full_split_prefix): + _run_decode_gen_active( + tma_desc_q, + tma_desc_k, + tma_desc_v, + tma_desc_k_atom, + tma_desc_v_atom, + o_iter, + g_s_k, + g_h_k, + g_scale_s_log2_e, + g_output_scale, + g_seqlens_kv, + g_cu_seqlens_q, + g_page_idx_kv, + g_partial_o, + g_partial_stats, + g_split_kv_counter, + g_attention_sinks, + g_h_r, + q_group_idx, + h_k_idx, + b_idx, + q_token_offset, + seq_len_q, + Int32(cfg.splits_kv), + True, + tile_sched_params, + cfg, + seq_len_kv, + use_variable_seqlens_kv, + use_native_paged_kv, + use_static_native_seqlens_kv, + g_block_tables, + g_block_table_capacity, + g_block_table_row_stride, + g_sparse_row_route_offsets, + g_sparse_row_route_counts, + g_sparse_route_metadata, + ) + else: + _run_decode_gen_runtime_prefix( + tma_desc_q, + tma_desc_k, + tma_desc_v, + tma_desc_k_atom, + tma_desc_v_atom, + o_iter, + g_s_k, + g_h_k, + g_scale_s_log2_e, + g_output_scale, + g_seqlens_kv, + g_cu_seqlens_q, + g_page_idx_kv, + g_partial_o, + g_partial_stats, + g_split_kv_counter, + g_attention_sinks, + g_h_r, + q_group_cta_idx, + q_group_idx, + h_k_idx, + b_idx, + q_token_offset, + seq_len_q, + q_token_base, + tile_sched_params, + cfg, + seq_len_kv, + use_variable_seqlens_kv, + use_native_paged_kv, + use_static_native_seqlens_kv, + g_block_tables, + g_block_table_capacity, + g_block_table_row_stride, + g_sparse_row_route_offsets, + g_sparse_row_route_counts, + g_sparse_route_metadata, + ) + else: + # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no + # producer state, unlike a valid Q tile whose K domain is empty. + _signal_padded_pdl_producer(cfg) + + +@cute.jit +def fmha_decode_launch( + problem_shape: tuple[Int32, Int32, Int32, Int32, Int32], + q_iter: cute.Pointer, + k_iter: cute.Pointer, + v_iter: cute.Pointer, + o_iter: cute.Pointer, + seqlens_kv_iter: cute.Pointer, + cu_seqlens_q_iter: cute.Pointer, + total_q_tokens: Int32, + page_idx_kv_iter: cute.Pointer, + partial_o_iter: cute.Pointer, + partial_stats_iter: cute.Pointer, + split_kv_counter_iter: cute.Pointer, + attention_sinks_iter: cute.Pointer, + scale_s: Float32, + output_scale: Float32, + kv_b_stride: Int32, + max_active_clusters: Int32, + stream: cuda_drv.CUstream, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + seq_len_kv: cutlass.Constexpr[int] = 2048, + use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, + use_native_paged_kv: cutlass.Constexpr[bool] = False, + block_tables_iter: cute.Pointer | None = None, + block_table_capacity: Int32 = 0, + block_table_row_stride: Int64 = 0, + num_physical_kv_pages: Int64 = 0, + k_page_stride: Int64 = 0, + v_page_stride: Int64 = 0, + static_full_split_prefix: cutlass.Constexpr[bool] = False, + use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, +) -> None: + """Standalone JIT launcher for FMHA decode TS.""" + log2_e = math.log2(math.e) + b, h_q, h_k, s_k, d = problem_shape + h_r = h_q // h_k + bias_static_sliding_kv_tma = ( + cfg.use_sliding_window_causal + and cfg.max_seq_len_q == 1 + and not cfg.use_paged_kv + and not cfg.use_split_kv + and not use_variable_seqlens_kv + ) + effective_seq_len_kv = _configure_static_sliding_window( + cfg, seq_len_kv, bias_static_sliding_kv_tma + ) + + q_seq = Int32(cfg.max_seq_len_q) + if cutlass.const_expr(cfg.use_paged_kv): + if cutlass.const_expr(use_native_paged_kv): + kv_shape = ( + d, + Int32(cfg.num_tokens_per_page), + h_k, + num_physical_kv_pages, + ) + k_layout = cute.make_layout( + kv_shape, + stride=( + 1, + d, + d * Int32(cfg.num_tokens_per_page), + k_page_stride, + ), + ) + v_layout = cute.make_layout( + kv_shape, + stride=( + 1, + d, + d * Int32(cfg.num_tokens_per_page), + v_page_stride, + ), + ) + k_tma = cute.make_tensor(k_iter, k_layout) + v_tma = cute.make_tensor(v_iter, v_layout) + else: + total_pages = b * Int32(cfg.max_num_pages_per_seq_kv) + kv_layout = cute.make_layout( + (d, Int32(cfg.num_tokens_per_page), h_k, total_pages), + stride=( + 1, + d, + d * Int32(cfg.num_tokens_per_page), + d * Int32(cfg.num_tokens_per_page) * h_k, + ), + ) + k_tma_iter = k_iter + v_tma_iter = v_iter + k_tma = cute.make_tensor(k_tma_iter, kv_layout) + v_tma = cute.make_tensor(v_tma_iter, kv_layout) + else: + kv_s_for_tma = s_k + k_tma_iter = k_iter + v_tma_iter = v_iter + if cutlass.const_expr(cfg.use_static_sliding_kv_tma_bias): + skipped_tokens = Int32( + _compute_static_num_skipped_kv_tiles(cfg, seq_len_kv) * cfg.tile_size_kv + ) + skipped_elems = skipped_tokens * d + k_tma_iter = k_iter + skipped_elems + v_tma_iter = v_iter + skipped_elems + kv_s_for_tma = Int32(effective_seq_len_kv) + kv_layout = cute.make_layout( + (d, kv_s_for_tma, h_k, b), stride=(1, d, d * s_k, kv_b_stride) + ) + k_tma = cute.make_tensor(k_tma_iter, kv_layout) + v_tma = cute.make_tensor(v_tma_iter, kv_layout) + + # Keep the TMA inner box at 128B when possible, but never exceed headDim. + # box_dim is expressed in elements of the source dtype, not bytes. + tma_box0_q = min(128 // cfg.q_dtype_bytes, cfg.headdim) + tma_box0_kv = min(128 // cfg.kv_dtype_bytes, cfg.headdim) + tma_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.use_fp8_qkv and cfg.headdim == 64): + tma_swizzle = cuda.TensorMapSwizzle.s64b + if cutlass.const_expr(cfg.tile_size_kv == 256): + # The 2x2 datapath consumes K in a (0, 2, 1, 3) KV64 permutation. + # A KV64 TensorMap atom lets the shared load resource place each + # semantic block directly in its physical slot for both paged and + # contiguous layouts. + tma_kv_tokens = min( + cfg.num_tokens_per_page if cfg.use_paged_kv else cfg.tile_size_kv, + 64, + ) + else: + tma_kv_tokens = ( + cfg.num_tokens_per_page + if cutlass.const_expr(cfg.use_paged_kv) + else cfg.tile_size_kv + ) + q_box_dims: tuple[object, ...] + if cutlass.const_expr(cfg.use_variable_seqlens_q): + # Packed Q is physically [sum_q_tokens, num_heads_q, head_dim]. Flatten + # the two head modes into Hq so the ragged token axis still fits in a + # rank-5 tensor map after the helper inserts its two synthetic modes. + q_tma = cute.make_tensor( + q_iter, + cute.make_layout( + (d, h_q, total_q_tokens), + stride=(1, d, h_q * d), + ), + ) + if cutlass.const_expr(cfg.groups_tokens_heads_q): + q_box_dims = ( + tma_box0_q, + cfg.heads_q_per_kv, + cfg.q_tokens_per_cta, + ) + q_groups = Int32( + (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta + ) + else: + q_box_dims = (tma_box0_q, cfg.tile_size_q, 1) + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + q_groups = Int32(cfg.max_seq_len_q) * head_ctas_per_token + tma_desc_q = create_tensor_map_ragged_from_tensor( + q_tma, + box_dims=q_box_dims, + ragged_dim=2, + stride_order=(0, 1, 2), + swizzle=tma_swizzle, + ) + else: + q_tma = cute.make_tensor( + q_iter, + cute.make_layout( + (d, h_r, h_k, q_seq, b), + stride=(1, d, h_r * d, h_q * d, q_seq * h_q * d), + ), + ) + if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): + q_box_dims = ( + tma_box0_q, + cfg.heads_q_per_kv, + 1, + cfg.q_tokens_per_cta, + 1, + ) + q_groups = Int32( + (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta + ) + else: + q_box_dims = (tma_box0_q, cfg.tile_size_q, 1, 1, 1) + q_groups = ( + (h_r + Int32(cfg.tile_size_q - 1)) // Int32(cfg.tile_size_q) + ) * q_seq + tma_desc_q = create_tensor_map_tiled_from_view( + q_tma, + box_dims=q_box_dims, + stride_order=(0, 1, 2, 3, 4), + swizzle=tma_swizzle, + ) + tma_desc_k = create_tensor_map_tiled_from_view( + k_tma, + box_dims=(tma_box0_kv, tma_kv_tokens, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle, + ) + tma_desc_v = create_tensor_map_tiled_from_view( + v_tma, + box_dims=(tma_box0_kv, tma_kv_tokens, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle, + ) + + grid_x = q_groups + if cutlass.const_expr(cfg.use_split_kv): + grid_x = q_groups * Int32(cfg.splits_kv) + + if cutlass.const_expr(cfg.use_persistent_scheduler): + tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + (grid_x, h_k, b), + (1, 1, 1), + ) + grid = tile_sched_params.get_grid_shape() + else: + tile_sched_params = None + grid = (grid_x, h_k, b) + # cluster distributed-SMEM reduction groups the splits_kv split CTAs + # of each sequence (contiguous in grid-X) into one cluster so they can write + # partials into each other's shared memory via prims.mapa. + cluster_shape = [1, 1, 1] + if cutlass.const_expr(cfg.use_cluster_smem_reduction): + cluster_shape = [cfg.splits_kv, 1, 1] + null_sparse_route_ptr = cute.make_ptr( + Int32, + 0, + mem_space=cutlass.AddressSpace.gmem, + ) + decode_gen_kernel( + tma_desc_q, + tma_desc_k, + tma_desc_v, + # Dense/paged profiles never inspect the 64-token descriptor slots. + tma_desc_k, + tma_desc_v, + o_iter, + s_k, + h_k, + Float32(scale_s * log2_e), + output_scale, + seqlens_kv_iter, + cu_seqlens_q_iter, + page_idx_kv_iter, + partial_o_iter, + partial_stats_iter, + split_kv_counter_iter, + attention_sinks_iter, + h_r, + tile_sched_params, + cfg, + seq_len_kv, + use_variable_seqlens_kv, + use_native_paged_kv, + use_static_native_seqlens_kv, + block_tables_iter, + block_table_capacity, + block_table_row_stride, + null_sparse_route_ptr, + null_sparse_route_ptr, + null_sparse_route_ptr, + static_full_split_prefix, + ).launch( + grid=grid, + block=[cfg.threads_per_cta, 1, 1], + cluster=cluster_shape, + stream=stream, + min_blocks_per_mp=_decode_min_blocks_per_mp(cfg, effective_seq_len_kv), + use_pdl=cfg.use_parallel_separate_reduction_pdl, + ) + + +@cute.jit +def fmha_block_sparse_launch( + problem_shape: tuple[Int32, Int32, Int32, Int32, Int32], + q_iter: cute.Pointer, + k_iter: cute.Pointer, + v_iter: cute.Pointer, + o_iter: cute.Pointer, + row_route_offsets_iter: cute.Pointer, + row_route_counts_iter: cute.Pointer, + route_metadata_iter: cute.Pointer, + scale_s: Float32, + stream: cuda_drv.CUstream, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + seq_len_kv: cutlass.Constexpr[int], + g_seqlens_kv: cute.Pointer | None = None, + use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, + num_physical_kv_pages: Int64 = 0, + k_page_stride: Int64 = 0, + v_page_stride: Int64 = 0, +) -> None: + """Launch attention over contiguous or paged prepared KV routes. + + A preceding prepare kernel has already resolved each BSR row into compact + logical atom origins, storage locators, validity flags, and optional token + words. Both layouts execute the same ``decode_gen_kernel`` schedule. + """ + if cutlass.const_expr(not cfg.use_block_sparse): + raise ValueError("fmha_block_sparse_launch requires cfg.use_block_sparse=True") + + log2_e = math.log2(math.e) + b, h_q, h_k, s_k, d = problem_shape + h_r = h_q // h_k + q_seq = Int32(cfg.max_seq_len_q) + q_strides, kv_strides = _block_sparse_bshd_tma_strides( + q_seq=q_seq, + h_q=h_q, + h_k=h_k, + s_k=s_k, + d=d, + ) + + # FP16/BF16 H128 uses a 64-element (128-byte) inner box. The sparse + # profile validator rejects other element widths and head dimensions. + tma_box0 = min(128 // cfg.kv_dtype_bytes, cfg.headdim) + tma_swizzle = cuda.TensorMapSwizzle.s128b + # Public tensors are contiguous BSHD. Q is factored into (Hr, Hkv) so the + # unchanged grouped-Q resource can address one KV head's grouped-Q tile. + q_desc = create_tensor_map_tiled( + global_address=q_iter.toint(), + dtype=cfg.q_dtype, + global_dims=(d, h_r, h_k, q_seq, b), + global_strides=q_strides, + box_dims=( + tma_box0, + cfg.heads_q_per_kv, + 1, + cfg.q_tokens_per_cta, + 1, + ), + swizzle=tma_swizzle, + ) + + kv_atom_size = _block_sparse_kv_atom_size(cfg.kv_block_size) + if cutlass.const_expr(cfg.use_paged_kv): + # Paged HND storage is addressed as (D, token-in-page, Hkv, page). + # Prepared routes already contain each atom's physical page ID, so no + # dense page table is passed to or staged by the attention kernel. + kv_shape = ( + d, + Int32(cfg.num_tokens_per_page), + h_k, + num_physical_kv_pages, + ) + k_layout = cute.make_layout( + kv_shape, + stride=( + 1, + d, + d * Int32(cfg.num_tokens_per_page), + k_page_stride, + ), + ) + v_layout = cute.make_layout( + kv_shape, + stride=( + 1, + d, + d * Int32(cfg.num_tokens_per_page), + v_page_stride, + ), + ) + k_desc_atom = create_tensor_map_tiled_from_view( + cute.make_tensor(k_iter, k_layout), + box_dims=(tma_box0, kv_atom_size, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle, + ) + v_desc_atom = create_tensor_map_tiled_from_view( + cute.make_tensor(v_iter, v_layout), + box_dims=(tma_box0, kv_atom_size, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle, + ) + k_desc_primary = k_desc_atom + v_desc_primary = v_desc_atom + else: + # Contiguous sparse coordinates retain the logical (D, S, H, B) + # order and the established primary/atom descriptor split. + primary_kv_box_size = 2 * kv_atom_size if kv_atom_size == 64 else kv_atom_size + kv_dims = (d, s_k, h_k, b) + k_desc_primary = create_tensor_map_tiled( + global_address=k_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=kv_dims, + global_strides=kv_strides, + box_dims=(tma_box0, primary_kv_box_size, 1, 1), + swizzle=tma_swizzle, + ) + v_desc_primary = create_tensor_map_tiled( + global_address=v_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=kv_dims, + global_strides=kv_strides, + box_dims=(tma_box0, primary_kv_box_size, 1, 1), + swizzle=tma_swizzle, + ) + k_desc_atom = k_desc_primary + v_desc_atom = v_desc_primary + if cutlass.const_expr( + kv_atom_size == 64 + and ( + cfg.tile_size_kv == 256 + or not _prepared_kv_routes_are_block_aligned( + cfg.kv_block_size, + cfg.tile_size_kv, + ) + ) + ): + # KV256 always stages four semantic KV64 atoms. KV128 needs this + # map only when a route may join unrelated BSR entries. + k_desc_atom = create_tensor_map_tiled( + global_address=k_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=kv_dims, + global_strides=kv_strides, + box_dims=(tma_box0, kv_atom_size, 1, 1), + swizzle=tma_swizzle, + ) + v_desc_atom = create_tensor_map_tiled( + global_address=v_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=kv_dims, + global_strides=kv_strides, + box_dims=(tma_box0, kv_atom_size, 1, 1), + swizzle=tma_swizzle, + ) + + q_groups = Int32( + (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta + ) + if cutlass.const_expr(cfg.use_persistent_scheduler): + tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + (q_groups, h_k, b), + (1, 1, 1), + ) + grid = tile_sched_params.get_grid_shape() + else: + tile_sched_params = None + grid = (q_groups, h_k, b) + + null_i32_ptr = cute.make_ptr( + Int32, + 0, + mem_space=cutlass.AddressSpace.gmem, + ) + null_f32_ptr = cute.make_ptr( + Float32, + 0, + mem_space=cutlass.AddressSpace.gmem, + ) + seqlens_kv_iter = ( + g_seqlens_kv if cutlass.const_expr(use_variable_seqlens_kv) else null_i32_ptr + ) + decode_gen_kernel( + q_desc, + k_desc_primary, + v_desc_primary, + k_desc_atom, + v_desc_atom, + o_iter, + s_k, + h_k, + Float32(scale_s * log2_e), + Float32(1.0), + seqlens_kv_iter, + null_i32_ptr, + null_i32_ptr, + o_iter, + null_f32_ptr, + null_i32_ptr, + null_f32_ptr, + h_r, + tile_sched_params, + cfg, + seq_len_kv, + use_variable_seqlens_kv, + False, # use_native_paged_kv + False, # use_static_native_seqlens_kv + null_i32_ptr, # g_block_tables + Int32(0), # g_block_table_capacity + Int64(0), # g_block_table_row_stride + row_route_offsets_iter, + row_route_counts_iter, + route_metadata_iter, + False, # static_full_split_prefix + ).launch( + grid=grid, + block=[cfg.threads_per_cta, 1, 1], + cluster=[1, 1, 1], + stream=stream, + # Reuse dense decode's entry-occupancy contract, including the long + # KV256 profiles that execute dynamic register reallocation. + min_blocks_per_mp=_decode_min_blocks_per_mp(cfg, seq_len_kv), + use_pdl=cfg.use_parallel_separate_reduction_pdl, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py new file mode 100644 index 000000000000..60360f90de17 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py @@ -0,0 +1,155 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resource definitions for the FMHA decode TS kernel. + +-------------------------------- + +SMEM resources +~~~~~~~~~~~~~~ +- SmemQResource : Q tile in SMEM, q_stages-deep TMA pipeline. + Producer (LoadTask): TMA Q GMEM -> SMEM (one or two + head-dim halves for 16-bit Q). Consumer (MmaTask): + builds a tcgen05 SMEM descriptor per stage for the + BMM1 (QK) A operand. Q is loaded once per work + tile in HEAD and reused across all BMM1 calls. + +- SmemPageOffsetsKvResource : SMEM-cached paged-KV page table entries. + Producer (dedicated prefetch warp): coalesced load + of a 32-page window from ``page_idx_kv`` into SMEM + (one stage per HEAD/LOOP K0/K1/V0/V1 cadence). + Consumer (TMA load warp): reads the ``pages_per_tile`` + slice for the current tile when issuing page-sized + TMA copies. Multiple consecutive tiles whose page + IDs share a 32-page window reuse the same stage. + +- SmemBlockSparseKvMetadataResource : Pipeline-free prepared route metadata + retained from one K load through the matching V. + +- SmemBlockSparseSoftmaxMetadataResource : Staged prepared route/token metadata + copied to Softmax task-local registers before the + corresponding pipeline stage is released. + +- SmemKvResource : Shared SMEM ring for K and V tiles. K and V + alternate in one allocation/pipeline; consumer + descriptors target the same SMEM but may use + different leading-byte offsets for K vs. V. + Producers (LoadTask): TMA K0/K1/V0/V1 GMEM -> SMEM + (paged or contiguous, depending on cfg). Consumers + (MmaTask): build tcgen05 K descriptors (BMM1 B + operand) and V descriptors (BMM2 B operand). + +- SmemKvTileResource : Dedicated SMEM tile used by split-head-dimension + profiles for one K or V producer instance. + +- SmemPResource : P operand for BMM2. The validated one-instance + staged-D256 Keeps profile places P in two TMEM views + aliased with the matching S stages; other profiles + use an SMEM tile. Producer + (Softmax): converts S in registers to P, writes the + selected operand layout, and publishes per-lane + local sums back through TmemS. Consumer (MmaTask): + publishes the matching TMEM address or SMEM + descriptor used by BMM2. + +TMEM resources +~~~~~~~~~~~~~~ +- TmemSResource : TMEM score buffer (BMM1 accumulator / softmax + input). Producer (MmaTask): QK MMA -> S in TMEM. + Consumer (Softmax): loads S to registers, maintains + running row max/sum, applies optional causal / + sliding-window / attention-sink masks, and feeds the + P producer. Also owns the SMEM softmax scratch + buffer used for the cross-warp atomic-max reduction. + +- TmemOResource : TMEM O accumulator for BMM2, o_stages deep. + Producer (MmaTask): P x V MMA -> O. Consumer + (Correction): tracks which O stage is ready + (``o_stage_idx`` plus tail stage indices) so the + in-place rescale path can find the correct columns. + +- TmemSoftmaxLocalResource : TMEM-local softmax stats exchanged with the + correction warps. Producer (Softmax): writes the + per-loop ``old_max``/``new_max``/``sum`` arrays and + the tail-visible per-instance copies. Consumer + (Correction): loads the stats to drive O rescaling + (LOOP) and final normalization (TAIL). + +- TmemSoftmaxOrderResource : Barrier-only ordering resource for softmax-stat + publication and correction consumption. + +- TmemStatsDoneResource : Barrier-only lifetime credit for TMEM columns + shared by S and local stats. MMA acquires it before + overwriting S; Correction returns it after loading + the matching stats into registers. + +- TmemSoftmaxGlobalResource : FP8 sum-correction helper. Producer-only + resource that, after P quantization, reapplies the + running-max correction to the running denominator + (using the TmemS local-sum array) and publishes the + corrected sums back through TmemS. Inactive when + non-FP8 Q/K/V. + +- TmemCorrResource : Correction and output resource. LOOP stages + rescale an in-flight O tile when the running max + changes; TAIL stages combine the two BMM2 instances, + normalize by the final denominator, and either + store the final O tile or write partial O for a + split-KV reduction. + +All resource classes derive from ``DecodeGenResourceBase`` (a thin +``MemoryResource`` subclass that marks ``consumer_vars`` / ``producer_vars`` +as Constexpr so the @cute.jit tracer does not traverse them during +dynamic-if serialization). +""" + +from .helpers_common import DecodeGenResourceBase +from .smem_resources import ( + SmemKvTileResource, + SmemKvResource, + SmemPageOffsetsKvResource, + SmemQResource, +) +from .smem_block_sparse_metadata import ( + SmemBlockSparseKvMetadataResource, + SmemBlockSparseSoftmaxMetadataResource, +) +from .tmem_corr import TmemCorrResource +from .tmem_o import TmemOResource +from .smem_p import SmemPResource +from .tmem_s import TmemSResource +from .tmem_softmax_stats import ( + TmemStatsDoneResource, + TmemSoftmaxGlobalResource, + TmemSoftmaxLocalResource, + TmemSoftmaxOrderResource, +) + +__all__ = [ + "DecodeGenResourceBase", + "SmemKvResource", + "SmemKvTileResource", + "SmemPageOffsetsKvResource", + "SmemPResource", + "SmemQResource", + "SmemBlockSparseKvMetadataResource", + "SmemBlockSparseSoftmaxMetadataResource", + "TmemCorrResource", + "TmemOResource", + "TmemSResource", + "TmemStatsDoneResource", + "TmemSoftmaxGlobalResource", + "TmemSoftmaxLocalResource", + "TmemSoftmaxOrderResource", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py new file mode 100644 index 000000000000..e3190f09097c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py @@ -0,0 +1,738 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Common helpers shared across FMHA decode TS resource files. + +Holds the small primitives, type aliases, task-cache offsets, config-driven +shape/dtype/swizzle helpers, and ``DecodeGenResourceBase`` — anything that +multiple resource classes (and the other ``_helpers_*`` modules) need. +""" + +from functools import partial +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float16, Float32, Int32, Int64, Uint32 +from cutlass.experimental import primitives as prims + +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + TaskLocalVariable, +) + +from ..fmha_decode_config import FmhaDecodeConfig + +Constexpr = cutlass.Constexpr +NEG_FLT_MAX = -3.4028235e38 +fadd2 = partial(cute.arch.add_packed_f32x2, ftz=False, rnd="rn") +fmul2 = partial(cute.arch.mul_packed_f32x2, ftz=False, rnd="rn") +ffma2 = partial(cute.arch.fma_packed_f32x2, ftz=False, rnd="rn") + +TaskCache = tuple[ + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, + Int32, +] +DescriptorValue = prims.Tcgen05SmemDesc | cutlass.Int64 +ResourceVarValue = ( + Int32 | Float32 | Uint32 | cutlass.Int64 | cutlass.Array | DescriptorValue +) +ResourceVars = dict[str, ResourceVarValue] + +# Offsets into DecodeGenTask.make_task_cache(). Keeping these symbolic makes +# resource code explicit about which task-local lane or address value it needs. +_TASK_CACHE_TMEM_BASE_OFFSET = 0 +_TASK_CACHE_WARP_GRP_THREAD_IDX = 1 +_TASK_CACHE_WARP_IDX = 2 +_TASK_CACHE_LANE_IDX = 3 +_TASK_CACHE_SEQ_LEN_KV = 4 +_TASK_CACHE_KV_REQUEST_BEGIN = 5 +_TASK_CACHE_KV_PAGE_IDX_UB = 6 +_TASK_CACHE_KV_RAW_TILE_BASE = 7 +_TASK_CACHE_KV_VALID_TILE_END = 8 +_TASK_CACHE_KV_WINDOW_START = 9 +# Block-sparse rows share the two generic KV-span words with paged decode. +# Semantic aliases keep sparse consumers independent of page-table naming +# without extending the stable ten-word task-cache ABI. +_TASK_CACHE_SPARSE_ROUTE_BEGIN = _TASK_CACHE_KV_REQUEST_BEGIN +_TASK_CACHE_SPARSE_ROUTE_COUNT = _TASK_CACHE_KV_PAGE_IDX_UB + + +@cute.jit +def _sparse_task_cache_route_begin(task_cache: TaskCache) -> Int32: + """Load the first prepared-route ordinal for one sparse row.""" + + return Int32(task_cache[_TASK_CACHE_SPARSE_ROUTE_BEGIN]) + + +@cute.jit +def _sparse_task_cache_route_count(task_cache: TaskCache) -> Int32: + """Load the live prepared-route count for one sparse row.""" + + return Int32(task_cache[_TASK_CACHE_SPARSE_ROUTE_COUNT]) + + +@cute.jit +def _warp_broadcast_i32(value: Int32, source_lane: Constexpr[int]) -> Int32: + """Broadcast one source-lane scalar as a warp-uniform Int32 value.""" + + return cute.arch.make_warp_uniform( + Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=value, + offset=source_lane, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + ) + + +def _mma_kind_for_qkv(cfg: FmhaDecodeConfig) -> prims.Tcgen05MMAKind: + """Select the tcgen05 MMA opcode family used for Q/K/V operands.""" + return prims.Tcgen05MMAKind.F8F6F4 if cfg.use_fp8_qkv else prims.Tcgen05MMAKind.F16 + + +def _mma_k_step(cfg: FmhaDecodeConfig) -> int: + """Return the K dimension advanced by one tcgen05 MMA instruction.""" + return 32 if cfg.use_fp8_qkv else 16 + + +@cute.jit +def _freeze_smem_descriptor(desc): + """Copy a SMEM descriptor through a register before MMA integer offsets.""" + return cute.arch.inline_ptx( + "mov.b64 {$w0}, {$r0};", + write_only_types=[Int64], + read_only_args=[desc], + ) + + +@cute.jit +def _wait_for_mbarrier_phase(barrier, phase: Int32) -> None: + """Wait for one parity of a CTA-local reusable mbarrier.""" + + while not prims.mbarrier_try_wait_parity( + barrier, + phase, + time_limit=10_000_000, + ): + pass + + +def _softmax_scale_pair_width(num_scale_groups: int, scale_base: int) -> int: + """Return the number of live lanes in one packed two-group operation.""" + return min(2, num_scale_groups - scale_base) + + +def _shape_tuple(shape: int | tuple[int, ...]) -> tuple[int, ...]: + """Normalize placeholder shapes to a tuple form accepted by cutlass.Array.""" + if isinstance(shape, tuple): + return shape + return (shape,) + + +def _placeholder_smem_array( + dtype: type, shape: int | tuple[int, ...] = 1 +) -> cutlass.Array | None: + """Create a fake SMEM array when resource construction needs a placeholder.""" + try: + return cutlass.Array( + cutlass.Int64(0), + dtype=dtype, + shape=_shape_tuple(shape), + addrspace=3, + ) + except (RuntimeError, ValueError): + return None + + +def _placeholder_local_array( + dtype: type, shape: int | tuple[int, ...] = 1, alignment: int | None = None +) -> cutlass.Array | None: + """Create a fake register-space array when tracing needs a placeholder.""" + try: + if alignment is None: + return cutlass.Array(dtype, shape, space=cutlass.AddressSpace.rmem) + return cutlass.Array( + dtype, shape, space=cutlass.AddressSpace.rmem, alignment=alignment + ) + except (RuntimeError, ValueError): + return None + + +@cute.jit +def _keeps_q64_row_idx(warp_grp_thread_idx: Int32) -> Int32: + """Map a correction lane to its q-row for tileSizeQ <= 64 keeps paths.""" + return (warp_grp_thread_idx >> Int32(5)) * Int32(16) + ( + warp_grp_thread_idx & Int32(0xF) + ) + + +@cute.jit +def _keeps_q64_col_base(lane_idx: Int32, half_cols: int) -> Int32: + """Return the head-dim column base for one q64 keeps lane group.""" + return (lane_idx >> Int32(4)) * Int32(half_cols) + + +@cute.jit +def _keeps_row_idx(cfg: Constexpr[FmhaDecodeConfig], warp_grp_thread_idx: Int32): + """Map a correction lane to the logical output row it owns.""" + if cutlass.const_expr(cfg.tile_size_kv == 256): + # KV256's 2x2 datapath exposes two spatial KV128 partials for each + # logical Q row. Threads [0, 64) and [64, 128) therefore share rows. + return warp_grp_thread_idx & Int32(63) + if cutlass.const_expr(cfg.tile_size_q == 128): + return warp_grp_thread_idx + return _keeps_q64_row_idx(warp_grp_thread_idx) + + +@cute.jit +def _keeps_col_base( + cfg: Constexpr[FmhaDecodeConfig], lane_idx: Int32, half_cols: int +) -> Int32: + """Return the lane's keepsMmaAb output-column base.""" + if cutlass.const_expr(cfg.tile_size_kv == 256): + return Int32(0) + if cutlass.const_expr(cfg.tile_size_q == 128): + return Int32(0) + return _keeps_q64_col_base(lane_idx, half_cols) + + +@cute.jit +def _keeps_score_col( + cfg: Constexpr[FmhaDecodeConfig], + warp_grp_thread_idx: Int32, + reg_idx: Constexpr[int], + col_base: Int32, +) -> Int32: + """Return the semantic KV column represented by one Keeps score register.""" + if cutlass.const_expr(cfg.tile_size_kv == 256): + # A physical thread owns four K32 fragments. The spatial half selects + # alternating KV64 blocks; the temporal fragment selects the low/high + # K32 sub-block inside that semantic KV64 block. + spatial = warp_grp_thread_idx >> Int32(6) + fragment = reg_idx // 32 + semantic_block = Int32(2 * (fragment // 2)) + spatial + return ( + semantic_block * Int32(64) + + Int32((fragment % 2) * 32) + + Int32(reg_idx % 32) + ) + return col_base + Int32(reg_idx) + + +@cute.jit +def _keeps_tcgen05_ld( + cfg: Constexpr[FmhaDecodeConfig], + tmem_addr, + *, + num: Constexpr[int], + offset: Constexpr[int], +): + """Load keepsMmaAb TMEM fragments using the tileSizeQ-specific shape.""" + if cutlass.const_expr(cfg.tile_size_kv == 256 or cfg.tile_size_q == 128): + return prims.tcgen05_ld( + "32x32b", + tmem_addr, + num=num, + ) + # The 16x32bx2 variant has a required half-split offset operand. Route + # through the public primitive wrapper so Python constants are materialized + # as MLIR values before reaching the low-level operation. + return prims.tcgen05_ld( + "16x32bx2", + tmem_addr, + num=num, + offset=offset, + ) + + +@cute.jit +def _keeps_tcgen05_st( + cfg: Constexpr[FmhaDecodeConfig], + tmem_addr, + val, + *, + offset: Constexpr[int], +) -> None: + """Store keepsMmaAb TMEM fragments using the tileSizeQ-specific shape.""" + if cutlass.const_expr(cfg.tile_size_kv == 256 or cfg.tile_size_q == 128): + prims.tcgen05_st( + "32x32b", + tmem_addr, + val, + ) + else: + prims.tcgen05_st( + "16x32bx2", + tmem_addr, + val, + offset=offset, + ) + + +@cute.jit +def _pack_float2_to_fp16(v0: Float32, v1: Float32) -> Int32: + """Pack two FP32 values into one FP16x2 register.""" + return cutlass.Vector.from_elements((v0, v1), Float32).to(Float16).bitcast(Int32)[0] + + +@cute.jit +def _pack_float2_to_bf16(v0: Float32, v1: Float32) -> Int32: + """Pack two FP32 values into one BF16x2 register.""" + return ( + cutlass.Vector.from_elements((v0, v1), Float32).to(BFloat16).bitcast(Int32)[0] + ) + + +def _qkv_smem_swizzle(cfg: FmhaDecodeConfig) -> prims.Tcgen05SmemSwizzle: + """Select the tcgen05 SMEM swizzle for staged Q/K/V tiles.""" + if cfg.use_fp8_qkv and cfg.headdim == 64: + return prims.Tcgen05SmemSwizzle.SWIZZLE_64B + return prims.Tcgen05SmemSwizzle.SWIZZLE_128B + + +def _major_k_stride_bytes(dtype_bytes: int, headdim: int) -> int: + """Return the descriptor K-major stride in bytes for one swizzle block.""" + # The descriptor swizzle shape depends on head dim and operand type, not + # on the number of Q rows in the tile. + num_smem_cols = 128 // dtype_bytes + rows_per_smem_row = max(1, num_smem_cols // headdim) + if rows_per_smem_row == 1: + rows_per_swizzle_blk = 8 + elif rows_per_smem_row == 2: + rows_per_swizzle_blk = 4 + elif rows_per_smem_row == 4: + rows_per_swizzle_blk = 2 + else: + rows_per_swizzle_blk = 1 + return 128 * rows_per_swizzle_blk + + +@cute.jit +def _fp8_log2_quant_scale() -> Float32: + """Return log2 scaling used by FP8 probability quantization.""" + return Float32(8.8073549) + + +def _neg_max_f32() -> Float32: + """Return the negative sentinel used for running softmax maxima.""" + return Float32(NEG_FLT_MAX) + + +def _softmax_tile_idx( + cfg: FmhaDecodeConfig, stage_info: StageInfo, inst_id: int +) -> Int32: + """Tile index consumed by the softmax-side MMA loop (inst_id ∈ {0, 1}).""" + return stage_info.loop_offset * Int32(cfg.num_insts_kv) + Int32(inst_id) + + +@cute.jit +def _named_barrier_arrive( + number_of_threads: Constexpr[int], barrier_id: Constexpr[int] +) -> None: + """Arrive at a named barrier from the configured warp subset.""" + # The non-aligned primitive is required here: ordered softmax uses four + # participating warps rather than a CTA-wide converged barrier. + prims.barrier_cta_arrive(barrier_id, number_of_threads) + + +@cute.jit +def _named_barrier_sync( + number_of_threads: Constexpr[int], barrier_id: Constexpr[int] +) -> None: + """Wait at a named barrier from the configured warp subset.""" + prims.barrier_cta_sync(barrier_id, thread_count=number_of_threads) + + +def _is_last_loop_iteration(stage_info: StageInfo) -> cutlass.Boolean: + """Return whether the current schedule loop iteration is the final one.""" + return stage_info.loop_offset + Int32(1) == stage_info.loop_end + + +def _clamp_valid_tile_idx(cfg: FmhaDecodeConfig, tile_idx: Int32) -> Int32: + """Clamp a static K/V tile index to the last valid tile.""" + return cute.math.min(tile_idx, Int32(cfg.total_kv_tiles - 1)) + + +@cute.jit +def _decode_gen_task_cache(stage_info: StageInfo) -> TaskCache: + """Return the task cache or a zero-filled placeholder cache.""" + if cutlass.const_expr(stage_info.task_cache is None): + zero = Int32(0) + return ( + zero, + zero, + zero, + zero, + zero, + zero, + zero, + zero, + zero, + zero, + ) + return stage_info.task_cache + + +@cute.jit +def _logical_head_batch( + stage_info: StageInfo, fallback_h_k_idx: Int32, fallback_b_idx: Int32 +) -> tuple[Int32, Int32]: + """Resolve logical KV head and batch from work tile or static launch.""" + if cutlass.const_expr(stage_info.work_tile is not None): + tile_idx = stage_info.work_tile.tile_idx + return Int32(tile_idx[1]), Int32(tile_idx[2]) + return fallback_h_k_idx, fallback_b_idx + + +@cute.jit +def _logical_q_group_idx( + cfg: Constexpr[FmhaDecodeConfig], + stage_info: StageInfo, + fallback_q_group_idx: Int32, +) -> Int32: + """Resolve the q-group from the persistent work tile or static launch.""" + if cutlass.const_expr(cfg.has_single_q_cta): + # A split coordinate may still vary in grid X, but every physical CTA + # maps to logical Q group zero. State this explicitly because the + # release compiler does not infer the range from the launch geometry. + return Int32(0) + if cutlass.const_expr(stage_info.work_tile is not None): + q_group_cta_idx = Int32(stage_info.work_tile.tile_idx[0]) + if cutlass.const_expr(cfg.use_split_kv): + return q_group_cta_idx // Int32(cfg.splits_kv) + return q_group_cta_idx + return fallback_q_group_idx + + +@cute.jit +def _q_tile_output_row_base( + cfg: Constexpr[FmhaDecodeConfig], q_group_idx: Int32 +) -> Int32: + """Return the first packed output row owned by a logical Q CTA.""" + if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): + return q_group_idx * Int32(cfg.q_tma_rows_per_cta) + if cutlass.const_expr(cfg.max_seq_len_q > 1): + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + token_idx = q_group_idx // head_ctas_per_token + head_cta_idx = q_group_idx - token_idx * head_ctas_per_token + return token_idx * Int32(cfg.heads_q_per_kv) + head_cta_idx * Int32( + cfg.tile_size_q + ) + return q_group_idx * Int32(cfg.tile_size_q) + + +@cute.jit +def _q_seq_bounds( + cfg: Constexpr[FmhaDecodeConfig], + cu_seqlens_q: cute.Pointer | None, + batch_idx: Int32, +) -> tuple[Int32, Int32]: + """Return packed Q token offset/length, or fixed-SQ neutral bounds.""" + if cutlass.const_expr(cfg.use_variable_seqlens_q): + q_begin = Int32(cu_seqlens_q[batch_idx]) + q_end = Int32(cu_seqlens_q[batch_idx + Int32(1)]) + return q_begin, q_end - q_begin + return Int32(0), Int32(cfg.max_seq_len_q) + + +@cute.jit +def _q_group_token_base(cfg: Constexpr[FmhaDecodeConfig], q_group_idx: Int32) -> Int32: + """Return the first Q token owned by a logical Q CTA.""" + if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): + return q_group_idx * Int32(cfg.q_tokens_per_cta) + if cutlass.const_expr(cfg.max_seq_len_q > 1): + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + return q_group_idx // head_ctas_per_token + return Int32(0) + + +@cute.jit +def _q_tile_valid_rows_for_seq( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + q_group_idx: Int32, + seq_len_q: Int32, +) -> Int32: + """Return runtime-valid MMA rows in one logical Q CTA. + + ``h_r`` is the total packed output-row count supplied to correction. Grouped + profiles advance by complete tokens, so structural padding and the final + partial token group are excluded independently from that global bound. + """ + if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): + token_base = q_group_idx * Int32(cfg.q_tokens_per_cta) + remaining_tokens = cute.math.max(seq_len_q - token_base, Int32(0)) + valid_tokens = cute.math.min(remaining_tokens, Int32(cfg.q_tokens_per_cta)) + return valid_tokens * Int32(cfg.heads_q_per_kv) + if cutlass.const_expr(cfg.max_seq_len_q > 1): + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + token_idx = q_group_idx // head_ctas_per_token + head_cta_idx = q_group_idx - token_idx * head_ctas_per_token + head_base = head_cta_idx * Int32(cfg.tile_size_q) + return cute.math.min( + cute.math.max(Int32(cfg.heads_q_per_kv) - head_base, Int32(0)), + Int32(cfg.tile_size_q), + ) + return cute.math.min( + cute.math.max( + h_r - q_group_idx * Int32(cfg.tile_size_q), + Int32(0), + ), + Int32(cfg.tile_size_q), + ) + + +@cute.jit +def _q_logical_output_row_token_and_local_head( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + logical_output_row_idx: Int32, +) -> tuple[Int32, Int32]: + """Decompose one batch-local scratch row into token and local Q head. + + Split-KV scratch stores logical rows densely as ``[token, local_head]``. + Packed final output needs that pair to replace the scratch buffer's padded + batch stride with the public cumulative-token offset. + """ + heads_q_per_kv = h_r + if cutlass.const_expr(cfg.heads_q_per_kv != 0): + heads_q_per_kv = Int32(cfg.heads_q_per_kv) + heads_q_per_kv_fdd = cute.fast_divmod_create_divisor(heads_q_per_kv) + return divmod(logical_output_row_idx, heads_q_per_kv_fdd) + + +@cute.jit +def _q_logical_output_row_is_valid_for_seq( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + logical_output_row_idx: Int32, + seq_len_q: Int32, +) -> cutlass.Boolean: + """Return whether a batch-local split scratch row owns packed output.""" + if cutlass.const_expr(cfg.use_variable_seqlens_q): + return logical_output_row_idx < seq_len_q * Int32(cfg.heads_q_per_kv) + return logical_output_row_idx < h_r + + +@cute.jit +def _q_physical_output_row_from_token_and_local_head( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + num_heads_kv: Int32, + batch_idx: Int32, + kv_head_idx: Int32, + q_token_idx: Int32, + local_head_idx: Int32, + q_token_offset: Int32, +) -> Int32: + """Map one token/local-head pair to the selected public O ABI.""" + if cutlass.const_expr(cfg.use_variable_seqlens_q): + heads_q_per_kv = Int32(cfg.heads_q_per_kv) + num_heads_q = num_heads_kv * heads_q_per_kv + global_head_idx = kv_head_idx * heads_q_per_kv + local_head_idx + return (q_token_offset + q_token_idx) * num_heads_q + global_head_idx + if cutlass.const_expr(cfg.max_seq_len_q > 1): + heads_q_per_kv = Int32(cfg.heads_q_per_kv) + num_heads_q = num_heads_kv * heads_q_per_kv + global_head_idx = kv_head_idx * heads_q_per_kv + local_head_idx + return ( + batch_idx * Int32(cfg.max_seq_len_q) + q_token_idx + ) * num_heads_q + global_head_idx + return (batch_idx * num_heads_kv + kv_head_idx) * h_r + local_head_idx + + +@cute.jit +def _q_physical_output_row_from_logical( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + num_heads_kv: Int32, + batch_idx: Int32, + kv_head_idx: Int32, + logical_output_row_idx: Int32, + q_token_offset: Int32, +) -> Int32: + """Map a batch-local split scratch row to the physical output tensor.""" + q_token_idx, local_head_idx = _q_logical_output_row_token_and_local_head( + cfg, h_r, logical_output_row_idx + ) + return _q_physical_output_row_from_token_and_local_head( + cfg, + h_r, + num_heads_kv, + batch_idx, + kv_head_idx, + q_token_idx, + local_head_idx, + q_token_offset, + ) + + +@cute.jit +def _q_physical_output_row( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + num_heads_kv: Int32, + batch_idx: Int32, + kv_head_idx: Int32, + q_group_idx: Int32, + tile_row_idx: Int32, + q_token_offset: Int32, +) -> Int32: + """Map a CTA-local Q row to the public output tensor's flat row. + + Fixed output is token-major ``[B, SQ, Hq, D]`` and packed variable-Q output + is token-major ``[sumQ, Hq, D]``. Resolve the common token/local-head pair + once so direct, split-GMEM, separate, and cluster final stores share the same + ABI. + """ + q_token_idx, local_head_idx = _q_row_token_and_local_head( + cfg, h_r, q_group_idx, tile_row_idx + ) + return _q_physical_output_row_from_token_and_local_head( + cfg, + h_r, + num_heads_kv, + batch_idx, + kv_head_idx, + q_token_idx, + local_head_idx, + q_token_offset, + ) + + +@cute.jit +def _q_row_is_valid_for_seq( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + q_group_idx: Int32, + tile_row_idx: Int32, + seq_len_q: Int32, +) -> cutlass.Boolean: + """Return whether a CTA-local Q row is valid for a runtime Q length.""" + return tile_row_idx < _q_tile_valid_rows_for_seq(cfg, h_r, q_group_idx, seq_len_q) + + +@cute.jit +def _q_row_token_and_local_head( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + q_group_idx: Int32, + tile_row_idx: Int32, +) -> tuple[Int32, Int32]: + """Map a valid CTA-local Q row to its token and KV-local Q head.""" + if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): + heads_q_per_kv_value = cfg.heads_q_per_kv + heads_q_per_kv = Int32(heads_q_per_kv_value) + if cutlass.const_expr( + heads_q_per_kv_value > 0 + and (heads_q_per_kv_value & (heads_q_per_kv_value - 1)) == 0 + ): + # Ratio-32 grouped decode lowers directly to shift/mask. Keep a + # FastDivmod fallback for future non-power-of-two grouped profiles. + shift = heads_q_per_kv_value.bit_length() - 1 + token_offset = tile_row_idx >> Int32(shift) + local_head = tile_row_idx & Int32(heads_q_per_kv_value - 1) + else: + heads_q_per_kv_fdd = cute.fast_divmod_create_divisor(heads_q_per_kv) + token_offset, local_head = divmod(tile_row_idx, heads_q_per_kv_fdd) + return ( + q_group_idx * Int32(cfg.q_tokens_per_cta) + token_offset, + local_head, + ) + if cutlass.const_expr(cfg.max_seq_len_q > 1): + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + token_idx = q_group_idx // head_ctas_per_token + head_cta_idx = q_group_idx - token_idx * head_ctas_per_token + return ( + token_idx, + head_cta_idx * Int32(cfg.tile_size_q) + tile_row_idx, + ) + return ( + Int32(0), + q_group_idx * Int32(cfg.tile_size_q) + tile_row_idx, + ) + + +@cute.jit +def _attention_sink_head_stride(cfg: Constexpr[FmhaDecodeConfig], h_r: Int32) -> Int32: + """Return the per-token local-head stride for attention-sink indexing.""" + if cutlass.const_expr(cfg.heads_q_per_kv != 0): + return Int32(cfg.heads_q_per_kv) + return h_r + + +@cute.jit +def _local_head_from_q_output_row( + cfg: Constexpr[FmhaDecodeConfig], + h_r: Int32, + output_row_idx: Int32, +) -> Int32: + """Recover the local Q head from a packed output row.""" + _, local_head_idx = _q_logical_output_row_token_and_local_head( + cfg, h_r, output_row_idx + ) + return local_head_idx + + +class DecodeGenResourceBase(MemoryResource): + """Base for decode-gen resources. + + Captured schedules let the framework manage variable lifecycle per work + tile. + + consumer_vars / producer_vars are marked Constexpr so that the @cute.jit + tracer's tree_flatten does NOT traverse them during dynamic-if + serialization. The framework accesses these dicts via + object.__getattribute__ which bypasses both the Constexpr filter and + the __getattribute__ guard. + """ + + consumer_vars: Constexpr[dict] = None + producer_vars: Constexpr[dict] = None + _task_local_specs: ClassVar[tuple[tuple, ...]] = () + + def __post_init__(self) -> None: + """Materialize task-local variables and placeholder resource state.""" + for name, dtype, default, docs in self._task_local_specs: + object.__setattr__( + self, + name, + TaskLocalVariable(dtype=dtype, default=default, docs=docs), + ) + self._init_placeholder_state() + + def _init_placeholder_state(self) -> None: + """Hook for subclasses to install placeholder arrays before tracing.""" + return diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_kv_tile_idx.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_kv_tile_idx.py new file mode 100644 index 000000000000..538a360badbb --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_kv_tile_idx.py @@ -0,0 +1,339 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""KV tile-index resolution helpers for FMHA decode TS resources. + +Handles seq-len lookup, sliding-window prefix skipping, and +split-KV / static-vs-runtime tile-index math used by ``SmemPageOffsetsKvResource``, +``SmemKvResource``, and the softmax / correction consumers. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + +from cutlass.experimental.task_scheduling.resources import StageInfo + +from ...mask import kv_tile_is_fully_visible +from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig +from .helpers_common import ( + _TASK_CACHE_SEQ_LEN_KV, + _decode_gen_task_cache, + _logical_head_batch, +) + + +@cute.jit +def _load_runtime_seq_len_kv( + seq_lens_kv: cute.Pointer | None, + max_seq_len_kv: Int32 | int, + stage_info: StageInfo, + fallback_h_k_idx: Int32, + fallback_b_idx: Int32, +) -> Int32: + """Load runtime KV length for the logical batch, or return static max.""" + # Persistent scheduling carries logical (head, batch) in the work tile; + # non-persistent kernels use the launch-time fallback coordinates. + _, logical_b_idx = _logical_head_batch(stage_info, fallback_h_k_idx, fallback_b_idx) + if cutlass.const_expr(seq_lens_kv is None): + return Int32(max_seq_len_kv) + if cutlass.const_expr(stage_info.task_cache is not None): + return Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_SEQ_LEN_KV]) + return Int32(seq_lens_kv[logical_b_idx]) + + +@cute.jit +def _runtime_total_kv_tiles( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return runtime KV tiles spanning this CTA's causal/window union.""" + seq_len_kv = _runtime_effective_seq_len_kv(cfg, seq_len_kv, seq_len_q, q_token_base) + return cute.ceil_div(seq_len_kv, cfg.tile_size_kv) + + +@cute.jit +def _runtime_configured_local_kv_tiles( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return the instruction-aligned local span for configured split capacity.""" + total_kv_tiles = _runtime_total_kv_tiles( + cfg, + seq_len_kv, + seq_len_q, + q_token_base, + ) + num_insts_kv = Int32(cfg.num_insts_kv) + configured_splits = Int32(cfg.splits_kv if cfg.use_split_kv else 1) + tiles_per_group = configured_splits * num_insts_kv + num_groups = (total_kv_tiles + tiles_per_group - Int32(1)) // tiles_per_group + return cute.math.max(num_groups * num_insts_kv, num_insts_kv) + + +@cute.jit +def _sliding_window_start_idx( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return the earliest sliding-window token needed by one logical Q CTA.""" + if cutlass.const_expr(not cfg.use_sliding_window_causal): + return Int32(0) + return cute.math.max( + seq_len_kv + - seq_len_q + + q_token_base + + Int32(1) + - Int32(cfg.attention_window_size), + Int32(0), + ) + + +@cute.jit +def _num_skipped_kv_tiles( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return full leading KV tiles skipped by the runtime sliding window.""" + return _sliding_window_start_idx(cfg, seq_len_kv, seq_len_q, q_token_base) // Int32( + cfg.tile_size_kv + ) + + +@cute.jit +def _runtime_effective_seq_len_kv( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return the CTA-visible K span after causal/window trimming. + + The left endpoint is rounded down to a complete K tile so the partial + boundary remains available for row-exact masking. The right endpoint is + the last valid Q row owned by this CTA, which is the union bound needed by + grouped multi-token Q tiles. + """ + skipped_tokens = _num_skipped_kv_tiles( + cfg, seq_len_kv, seq_len_q, q_token_base + ) * Int32(cfg.tile_size_kv) + visible_k_end = seq_len_kv + if cutlass.const_expr(cfg.mask_type == CAUSAL): + q_token_end = cute.math.min( + q_token_base + Int32(cfg.q_tokens_per_cta), + seq_len_q, + ) + visible_k_end = cute.math.min( + cute.math.max( + seq_len_kv - seq_len_q + q_token_end, + Int32(0), + ), + seq_len_kv, + ) + return cute.math.max(visible_k_end - skipped_tokens, Int32(0)) + + +@cute.jit +def _kv_tile_is_fully_unmasked_for_q_group( + cfg: FmhaDecodeConfig, + tile_offset_k: Int32, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, + tile_has_valid_scores: cutlass.Boolean, +) -> cutlass.Boolean: + """Return whether a KV tile needs no score mask for any active Q row. + + Grouped causal rows expose an intersection bounded on the right by the + earliest Q token. A sliding window also bounds it on the left by the + latest active Q token. Tiles outside that intersection retain the exact + per-row boundary path. + """ + visible_k_begin = Int32(0) + visible_k_end = seq_len_kv + if cutlass.const_expr(cfg.mask_type == CAUSAL): + visible_k_end = cute.math.min( + cute.math.max( + seq_len_kv - seq_len_q + q_token_base + Int32(1), + Int32(0), + ), + seq_len_kv, + ) + if cutlass.const_expr(cfg.use_sliding_window_causal): + q_token_end = cute.math.min( + q_token_base + Int32(cfg.q_tokens_per_cta), + seq_len_q, + ) + latest_causal_end = cute.math.min( + cute.math.max( + seq_len_kv - seq_len_q + q_token_end, + Int32(0), + ), + seq_len_kv, + ) + visible_k_begin = cute.math.max( + latest_causal_end - Int32(cfg.attention_window_size), + Int32(0), + ) + return tile_has_valid_scores and kv_tile_is_fully_visible( + tile_offset_k, + Int32(cfg.tile_size_kv), + visible_k_begin, + visible_k_end, + ) + + +@cute.jit +def _runtime_active_splits_kv( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return the useful split prefix for one runtime Q/KV work item.""" + if cutlass.const_expr(not cfg.use_split_kv): + return Int32(1) + total_kv_tiles = _runtime_total_kv_tiles( + cfg, + seq_len_kv, + seq_len_q, + q_token_base, + ) + local_kv_tiles = _runtime_configured_local_kv_tiles( + cfg, + seq_len_kv, + seq_len_q, + q_token_base, + ) + # Split ranges are instruction-group aligned. Only ranges intersecting the + # valid K domain participate; remaining configured grid slots are padding. + return (total_kv_tiles + local_kv_tiles - Int32(1)) // local_kv_tiles + + +@cute.jit +def _runtime_execution_splits_kv( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return a nonzero producer/reduction fanout for a runtime-valid Q tile.""" + return cute.math.max( + _runtime_active_splits_kv(cfg, seq_len_kv, seq_len_q, q_token_base), + Int32(1), + ) + + +@cute.jit +def _runtime_last_valid_tile_idx( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return the last valid runtime KV tile index.""" + return cute.math.max( + _runtime_total_kv_tiles(cfg, seq_len_kv, seq_len_q, q_token_base) - Int32(1), + Int32(0), + ) + + +@cute.jit +def _runtime_clamp_valid_tile_idx( + cfg: FmhaDecodeConfig, + tile_idx: Int32, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Clamp a runtime KV tile index to the last valid tile.""" + return cute.math.min( + tile_idx, + _runtime_last_valid_tile_idx(cfg, seq_len_kv, seq_len_q, q_token_base), + ) + + +@cute.jit +def _runtime_last_valid_page_idx(cfg: FmhaDecodeConfig, seq_len_kv: Int32) -> Int32: + """Return the last valid page index for a runtime KV length.""" + num_pages = cute.ceil_div(seq_len_kv, cfg.num_tokens_per_page) + return cute.math.max(num_pages - Int32(1), Int32(0)) + + +@cute.jit +def _static_split_kv_global_tile_idx( + cfg: FmhaDecodeConfig, stage_info: StageInfo, local_tile_idx: Int32 +) -> Int32: + """Map a static local tile index to a global split-KV tile index.""" + if cutlass.const_expr(not cfg.use_split_kv): + return local_tile_idx + return ( + _logical_cta_kv_idx(cfg, stage_info) * Int32(cfg.static_local_kv_tiles) + + local_tile_idx + ) + + +@cute.jit +def _logical_cta_kv_idx(cfg: FmhaDecodeConfig, stage_info: StageInfo) -> Int32: + """Resolve the split index from the Q-group-major launch coordinate.""" + if cutlass.const_expr(stage_info.work_tile is not None): + q_group_cta_idx = Int32(stage_info.work_tile.tile_idx[0]) + else: + q_group_cta_idx, _, _ = cute.arch.block_idx() + if cutlass.const_expr(cfg.use_split_kv): + return q_group_cta_idx % Int32(cfg.splits_kv) + return q_group_cta_idx + + +@cute.jit +def _runtime_local_kv_tiles( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Return local KV tiles assigned to one split CTA at runtime.""" + return _runtime_configured_local_kv_tiles( + cfg, + seq_len_kv, + seq_len_q, + q_token_base, + ) + + +@cute.jit +def _runtime_split_kv_global_tile_idx( + cfg: FmhaDecodeConfig, + stage_info: StageInfo, + local_tile_idx: Int32, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, +) -> Int32: + """Map a runtime local tile index to a global split-KV tile index.""" + if cutlass.const_expr(not cfg.use_split_kv): + return local_tile_idx + return ( + _logical_cta_kv_idx(cfg, stage_info) + * _runtime_local_kv_tiles(cfg, seq_len_kv, seq_len_q, q_token_base) + + local_tile_idx + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_output.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_output.py new file mode 100644 index 000000000000..db6bd6b8de8c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_output.py @@ -0,0 +1,459 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SMEM/GMEM output helpers for FMHA decode TS resources. + +Holds the transposed FP8 STSM stores, the SMEM→GMEM 16-byte vector copy, +the 16-bit O-reorg offset math, the P-STSM offset math, and partial-O +load helpers used by ``SmemPResource`` and ``TmemCorrResource``. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float16, Float32, Int32, Int64 +from cutlass.experimental import primitives as prims + +from ..fmha_decode_config import FmhaDecodeConfig +from .helpers_common import ( + Constexpr, + _q_physical_output_row, + _q_row_is_valid_for_seq, +) + + +@cute.jit +def _keeps_p_smem_block_offset_bytes( + cfg: Constexpr[FmhaDecodeConfig], row_idx: Int32, col_idx: Int32 +) -> Int32: + """JIT form of the Keeps P SWIZZLE_128B vector-block address.""" + dtype_bytes = cfg.q_dtype_bytes + chunk_cols = 128 // dtype_bytes + chunk_idx = col_idx // Int32(chunk_cols) + col_in_chunk = col_idx - chunk_idx * Int32(chunk_cols) + byte_col = col_in_chunk * Int32(dtype_bytes) + return ( + chunk_idx * Int32(cfg.tile_size_q * 128) + + row_idx * Int32(128) + + (((byte_col >> Int32(4)) ^ (row_idx & Int32(0x7))) << Int32(4)) + ) + + +@cute.jit +def _fp8_stsm_smem_dst( + smem_base_i32: cutlass.Array, + warp_grp_thread_idx: Int32, + num_trans_rows: int, + num_trans_cols: int, + stsm_idx: int, +) -> cute.Pointer: + """Return the swizzled SMEM destination for one FP8 transposed STSM.""" + # Compute the swizzled SMEM address used by transposed 8-bit STSM. The + # mapping writes lane fragments in the same layout the later vector GMEM + # copy expects. + num_rows = Int32(num_trans_rows) + num_bytes_per_row = Int32(num_trans_cols) + num_rows_per_128b = Int32(128) // num_bytes_per_row + num_segs_per_warp_per_row = num_bytes_per_row // Int32(16 * 4) + num_stsm_per_row = max(8 // 32, 1) + num_mtx_per_col = num_rows // Int32(8) + warp_idx = warp_grp_thread_idx >> Int32(5) + lane_idx = warp_grp_thread_idx & Int32(0x1F) + thr_row_idx = lane_idx & Int32(0x7) + mtx_idx = lane_idx >> Int32(3) + mtx_row_idx = mtx_idx % num_mtx_per_col + mtx_col_idx = mtx_idx // num_mtx_per_col + + stsm_row_idx = Int32(stsm_idx % num_stsm_per_row) + stsm_col_idx = Int32(stsm_idx // num_stsm_per_row) + xor_mask = thr_row_idx // num_rows_per_128b + seg_col_idx = ( + warp_idx * num_segs_per_warp_per_row + mtx_col_idx + stsm_col_idx + ) ^ xor_mask + smem_offset = ( + mtx_row_idx * Int32(8) + thr_row_idx + stsm_row_idx * Int32(32) + ) * num_bytes_per_row + seg_col_idx * Int32(16) + return (smem_base_i32.subview((smem_offset >> Int32(2)))).data_ptr() + + +@cute.jit +def _store_transposed_smem8b_x1( + smem_base_i32: cutlass.Array, + reg0: Int32, + warp_grp_thread_idx: Int32, + num_trans_rows: int, + num_trans_cols: int, + stsm_idx: int = 0, +) -> None: + """Store one packed 8-bit register with transposed stmatrix.""" + smem_dst = _fp8_stsm_smem_dst( + smem_base_i32, + warp_grp_thread_idx, + num_trans_rows, + num_trans_cols, + stsm_idx, + ) + prims.inline_ptx_hl( + "stmatrix.sync.aligned.m16n8.x1.trans.shared.b8 [{$r0}], {{$r1}};", + read_only_args=[smem_dst, reg0], + ) + + +@cute.jit +def _store_transposed_smem8b_x2( + smem_base_i32: cutlass.Array, + reg0: Int32, + reg1: Int32, + warp_grp_thread_idx: Int32, + num_trans_rows: int, + num_trans_cols: int, + stsm_idx: int = 0, +) -> None: + """Store two packed 8-bit registers with transposed stmatrix.""" + smem_dst = _fp8_stsm_smem_dst( + smem_base_i32, + warp_grp_thread_idx, + num_trans_rows, + num_trans_cols, + stsm_idx, + ) + prims.inline_ptx_hl( + "stmatrix.sync.aligned.m16n8.x2.trans.shared.b8 [{$r0}], {{$r1}, {$r2}};", + read_only_args=[smem_dst, reg0, reg1], + ) + + +@cute.jit +def _store_transposed_smem8b_x4( + smem_base_i32: cutlass.Array, + reg0: Int32, + reg1: Int32, + reg2: Int32, + reg3: Int32, + warp_grp_thread_idx: Int32, + num_trans_rows: int, + num_trans_cols: int, + stsm_idx: int = 0, +) -> None: + """Store four packed 8-bit registers with transposed stmatrix.""" + smem_dst = _fp8_stsm_smem_dst( + smem_base_i32, + warp_grp_thread_idx, + num_trans_rows, + num_trans_cols, + stsm_idx, + ) + prims.inline_ptx_hl( + "stmatrix.sync.aligned.m16n8.x4.trans.shared.b8 [{$r0}], {{$r1}, {$r2}, {$r3}, {$r4}};", + read_only_args=[smem_dst, reg0, reg1, reg2, reg3], + ) + + +@cute.jit +def _store_transposed_smem8b( + smem_base_i32: cutlass.Array, + regs: cutlass.Array, + warp_grp_thread_idx: Int32, + num_trans_rows: int, + num_trans_cols: int, + num_regs: int, +) -> None: + """Dispatch FP8 transposed stores for the register count owned by a lane.""" + # STSM's swizzled row is at most 128 bytes. Wider outputs are laid out as + # consecutive 128-byte head-dimension bands, each containing every Q row. + # This keeps the existing D64/D128 mapping and makes D256 independent of + # TileQ instead of accidentally treating the second band as extra Q rows. + assert num_trans_rows >= 8 and num_trans_rows % 8 == 0 + assert num_trans_cols > 0 and ( + num_trans_cols in (64, 128) or num_trans_cols % 128 == 0 + ) + num_col_bands = max((num_trans_cols + 127) // 128, 1) + band_cols = min(num_trans_cols, 128) + assert num_regs % num_col_bands == 0 + regs_per_band = num_regs // num_col_bands + assert regs_per_band in (1, 2, 4, 8) + for band_idx in cutlass.range_constexpr(num_col_bands): + band_smem_base = smem_base_i32.subview( + band_idx * num_trans_rows * band_cols // 4 + ) + reg_base = band_idx * regs_per_band + if cutlass.const_expr(regs_per_band == 1): + _store_transposed_smem8b_x1( + band_smem_base, + regs[reg_base], + warp_grp_thread_idx, + num_trans_rows, + band_cols, + ) + elif cutlass.const_expr(regs_per_band == 2): + _store_transposed_smem8b_x2( + band_smem_base, + regs[reg_base], + regs[reg_base + 1], + warp_grp_thread_idx, + num_trans_rows, + band_cols, + ) + else: + _store_transposed_smem8b_x4( + band_smem_base, + regs[reg_base], + regs[reg_base + 1], + regs[reg_base + 2], + regs[reg_base + 3], + warp_grp_thread_idx, + num_trans_rows, + band_cols, + ) + if cutlass.const_expr(regs_per_band > 4): + _store_transposed_smem8b_x4( + band_smem_base, + regs[reg_base + 4], + regs[reg_base + 5], + regs[reg_base + 6], + regs[reg_base + 7], + warp_grp_thread_idx, + num_trans_rows, + band_cols, + 1, + ) + + +def _fp8_smem_load_xor_bytes_for_shape( + headdim: int, smem_row_idx: int | Int32 +) -> int | Int32: + """Return the byte XOR that reverses the FP8 STSM row swizzle. + + A 128-byte swizzle atom packs multiple D64 rows. The store-side XOR is + derived from the row inside each eight-row matrix, not the global packed + row number; including matrix-row bits would flip bit 6 and exchange the two + D64 rows sharing one 128-byte atom. + """ + num_bytes_per_smem_row = min(headdim, 128) + num_packed_smem_rows = 128 // num_bytes_per_smem_row + return ((smem_row_idx % Int32(8)) // Int32(num_packed_smem_rows)) * Int32(16) + + +@cute.jit +def _copy_transposed_smem8b_to_gmem( + smem_base_i32: cutlass.Array, + o_ptr: cute.Pointer, + cfg: Constexpr[FmhaDecodeConfig], + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_q_group_idx: Int32, + h_r: Int32, + num_heads_kv: Int32, + seq_len_q: Int32, + q_token_offset: Int32, + warp_grp_thread_idx: Int32, + full_tile_rows: Constexpr[bool] = False, +) -> None: + """Copy a transposed FP8 SMEM tile to the logical GMEM output layout.""" + # Reload the transposed 8-bit SMEM tile as contiguous 16-byte vectors and + # store the vectors into the logical output row. + headdim = cfg.headdim + tile_size_q = cfg.tile_size_q + num_bytes_per_smem_row = Int32(headdim if headdim <= 128 else 128) + num_copy_segments = max((tile_size_q * headdim + 2047) // 2048, 1) + for copy_segment_idx in cutlass.range_constexpr(num_copy_segments): + base_offset = warp_grp_thread_idx * Int32(16) + Int32(copy_segment_idx * 2048) + smem_row_idx = base_offset // num_bytes_per_smem_row + load_smem_offset = base_offset ^ _fp8_smem_load_xor_bytes_for_shape( + headdim, smem_row_idx + ) + dst_row_idx = smem_row_idx + dst_col_offset = base_offset % num_bytes_per_smem_row + if headdim > 128: + dst_row_idx = smem_row_idx % Int32(tile_size_q) + dst_col_offset = ( + smem_row_idx // Int32(tile_size_q) + ) * num_bytes_per_smem_row + (base_offset % num_bytes_per_smem_row) + physical_dst_row_idx = _q_physical_output_row( + cfg, + h_r, + num_heads_kv, + logical_b_idx, + logical_h_k_idx, + logical_q_group_idx, + dst_row_idx, + q_token_offset, + ) + valid_output_row = _q_row_is_valid_for_seq( + cfg, + h_r, + logical_q_group_idx, + dst_row_idx, + seq_len_q, + ) + if cutlass.const_expr(full_tile_rows): + smem_src_full = ( + smem_base_i32.subview((load_smem_offset >> Int32(2))) + ).data_ptr() + # The flattened row is intentionally Int32, but the byte address + # may exceed 2 GiB for a valid packed-Q output. Widen before the + # row-stride product so it cannot wrap in signed 32-bit arithmetic. + dst_row_base_full = Int64(physical_dst_row_idx) * Int64(headdim) + dst_ptr_full = cutlass.inttoptr( + o_ptr.toint() + dst_row_base_full + Int64(dst_col_offset), + mem_space=1, + dtype=Int32, + ) + dst_ptr_full.store(smem_src_full.load(count=4, alignment=16), alignment=16) + else: + if valid_output_row: + smem_src_guarded = ( + smem_base_i32.subview((load_smem_offset >> Int32(2))) + ).data_ptr() + dst_row_base_guarded = Int64(physical_dst_row_idx) * Int64(headdim) + dst_ptr_guarded = cutlass.inttoptr( + o_ptr.toint() + dst_row_base_guarded + Int64(dst_col_offset), + mem_space=1, + dtype=Int32, + ) + dst_ptr_guarded.store( + smem_src_guarded.load(count=4, alignment=16), alignment=16 + ) + + +@cute.jit +def _transposed_smem128x16b_stsm_offset_bytes( + num_rows: Constexpr[int], + local_warp_idx: Int32, + lane_idx: Int32, + stsm_group_idx: int = 0, +) -> Int32: + """Compute the SMEM byte offset for one transposed stmatrix store.""" + # Store transposed 128x16b rows. Once there are multiple TMEM load repeats, STSM + # groups advance both row and column sub-tiles. + num_tmem_load_reps = max(num_rows // 8, 1) + if cutlass.const_expr(num_tmem_load_reps == 1): + num_mtx_per_row_per_stsm = 4 + num_mtx_per_col_per_stsm = 1 + else: + num_mtx_per_row_per_stsm = 2 + num_mtx_per_col_per_stsm = 2 + num_stsm_per_col = max(num_tmem_load_reps // num_mtx_per_col_per_stsm, 1) + stsm_row_group_idx = stsm_group_idx // num_stsm_per_col + stsm_col_group_idx = stsm_group_idx % num_stsm_per_col + + slice_idx = local_warp_idx // Int32(2) + warp_idx_in_slice = local_warp_idx % Int32(2) + mtx_idx = lane_idx // Int32(8) + mtx_row_idx = mtx_idx // Int32(num_mtx_per_row_per_stsm) + thr_row_idx = lane_idx % Int32(8) + mtx_col_idx = ( + warp_idx_in_slice * Int32(4) + + (mtx_idx % Int32(num_mtx_per_row_per_stsm)) + + Int32(stsm_row_group_idx * num_mtx_per_row_per_stsm) + ) + return ( + slice_idx * Int32(num_rows * 128) + + ( + mtx_row_idx * Int32(8) + + thr_row_idx + + Int32(stsm_col_group_idx * num_mtx_per_col_per_stsm * 8) + ) + * Int32(128) + + ((mtx_col_idx ^ thr_row_idx) * Int32(16)) + ) + + +@cute.jit +def _fp16_o_reorg_offsets( + cfg: FmhaDecodeConfig, + warp_grp_thread_idx: Int32, + local_warp_idx: Int32, + lane_idx: Int32, + stsm_group_idx: int = 0, + copy_segment_idx: int = 0, +) -> tuple[Int32, Int32, Int32, Int32]: + """Return the SMEM reorg offsets for a 16-bit O/partial-O row.""" + + o_stage_dtype_bytes = cfg.o_dtype_bytes + if cutlass.const_expr(cfg.use_split_kv and cfg.use_fp8_output): + # Split-KV partial O is staged in 16-bit form even when final O is FP8. + o_stage_dtype_bytes = 2 + base_offset = (warp_grp_thread_idx << Int32(4)) + Int32(copy_segment_idx * 2048) + smem_row_idx = base_offset >> Int32(7) + slice_idx = local_warp_idx >> Int32(1) + warp_idx_in_slice = local_warp_idx & Int32(1) + thr_row_idx = lane_idx & Int32(0x7) + if cutlass.const_expr(cfg.headdim * 2 > 128): + if cutlass.const_expr(cfg.tile_size_q >= 16): + stsm_per_head_dim_stage = max(cfg.tile_size_q // 8, 1) + head_dim_stage_idx = stsm_group_idx // stsm_per_head_dim_stage + stage_stsm_group_idx = stsm_group_idx % stsm_per_head_dim_stage + smem_offset_bytes = Int32( + head_dim_stage_idx + * cfg.tile_size_q + * cfg.head_dim_kv_stage + * o_stage_dtype_bytes + ) + _transposed_smem128x16b_stsm_offset_bytes( + cfg.tile_size_q, + local_warp_idx, + lane_idx, + stage_stsm_group_idx, + ) + else: + mtx_col_idx_256b = (warp_idx_in_slice << Int32(2)) + ( + (lane_idx >> Int32(3)) & Int32(0x3) + ) + smem_offset_bytes = ( + slice_idx * Int32(8 * 128) + + Int32(stsm_group_idx * 16 * 128) + + thr_row_idx * Int32(128) + + ((mtx_col_idx_256b ^ thr_row_idx) * Int32(16)) + ) + load_smem_offset = base_offset ^ ((smem_row_idx & Int32(0x7)) << Int32(4)) + dst_row_idx = smem_row_idx % Int32(cfg.tile_size_q) + dst_col_offset = (smem_row_idx // Int32(cfg.tile_size_q)) * Int32(128) + ( + base_offset & Int32(0x7F) + ) + else: + mtx_idx = lane_idx >> Int32(3) + mtx_row_idx = mtx_idx >> Int32(1) + mtx_col_idx = mtx_idx & Int32(1) + seg_col_idx = ((local_warp_idx << Int32(1)) + mtx_col_idx) ^ thr_row_idx + smem_row = mtx_row_idx * Int32(8) + thr_row_idx + Int32(stsm_group_idx * 16) + smem_offset_bytes = smem_row * Int32(128) + seg_col_idx * Int32(16) + load_smem_offset = base_offset ^ ((smem_row_idx & Int32(0x7)) << Int32(4)) + dst_row_idx = smem_row_idx + dst_col_offset = base_offset & Int32(0x7F) + return smem_offset_bytes, load_smem_offset, dst_row_idx, dst_col_offset + + +@cute.jit +def _p_stsm_smem_offset_bytes( + local_warp_idx: Int32, + lane_idx: Int32, + stsm_group_idx: int = 0, + tile_size_q: int = 16, +) -> Int32: + """Return the SMEM byte offset used when storing P with stmatrix.""" + return _transposed_smem128x16b_stsm_offset_bytes( + tile_size_q, local_warp_idx, lane_idx, stsm_group_idx + ) + + +@cute.jit +def _load_partial_o_vec8_as_f32( + regs_i32: cutlass.Array, use_bf16_partial: Constexpr[bool] +) -> cutlass.Array: + """Convert one packed split-KV partial-O vector to FP32 values.""" + regs_vec = cutlass.Vector.from_elements( + (regs_i32[0], regs_i32[1], regs_i32[2], regs_i32[3]), Int32 + ) + if cutlass.const_expr(use_bf16_partial): + return regs_vec.bitcast(BFloat16).to(Float32) + return regs_vec.bitcast(Float16).to(Float32) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py new file mode 100644 index 000000000000..4585213cd80b --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py @@ -0,0 +1,461 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Softmax / P-compute helpers for FMHA decode TS resources. + +Used by ``TmemSResource`` (softmax reduction, atomic-max scratch), +``SmemPResource`` (S→P conversion and quantization), ``TmemCorrResource`` +(attention-sink normalization), and the softmax-stats resources. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Uint32 + +from cutlass.experimental import primitives as prims + +from ..fmha_decode_config import FmhaDecodeConfig +from .helpers_common import ( + Constexpr, + fadd2, + ffma2, + fmul2, + _fp8_log2_quant_scale, + _neg_max_f32, + _q_row_token_and_local_head, +) + + +@cute.jit +def _pack_float4_to_fp8_e4m3( + v0: Float32, v1: Float32, v2: Float32, v3: Float32 +) -> Int32: + """Pack four FP32 values into one FP8 E4M3x4 register.""" + return _pack_float4_to_fp8_e4m3_inline(v0, v1, v2, v3) + + +@cute.jit +def _pack_float4_to_fp8_e4m3_inline( + v0: Float32, + v1: Float32, + v2: Float32, + v3: Float32, + *, + loc=None, + ip=None, +) -> Int32: + """Pack FP8x4 in one public inline-PTX block.""" + return cute.arch.inline_ptx( + "{\n" + " .reg .b16 lo;\n" + " .reg .b16 hi;\n" + " cvt.rn.satfinite.e4m3x2.f32 lo, {$r1}, {$r0};\n" + " cvt.rn.satfinite.e4m3x2.f32 hi, {$r3}, {$r2};\n" + " mov.b32 {$w0}, {lo, hi};\n" + "}", + write_only_types=[Int32], + read_only_args=[v0, v1, v2, v3], + loc=loc, + ip=ip, + ) + + +@cute.jit +def _compute_fp8_p_regs_and_local_sums( + scale_softmax_log2: Float32, + new_max_0: Float32, + new_max_1: Float32, + s0: Float32, + s1: Float32, + s2: Float32, + s3: Float32, + s4: Float32, + s5: Float32, + s6: Float32, + s7: Float32, +) -> tuple[Int32, Int32, Float32, Float32]: + """Compute masked FP8 P registers and local softmax sums.""" + # Safe path: masked tiles can produce NEG_FLT_MAX as new_max. Treat that + # as zero for the exponent offset so invalid rows generate zero P instead + # of NaNs while local sums remain zero. + safe_new_max_0 = new_max_0 + safe_new_max_1 = new_max_1 + if safe_new_max_0 == _neg_max_f32(): + safe_new_max_0 = Float32(0.0) + if safe_new_max_1 == _neg_max_f32(): + safe_new_max_1 = Float32(0.0) + neg_scaled_max_pair = ffma2( + (safe_new_max_0, safe_new_max_1), + (-scale_softmax_log2, -scale_softmax_log2), + (_fp8_log2_quant_scale(), _fp8_log2_quant_scale()), + ) + + # Scale S by log2(e) * softmax_scale and include the FP8 quantization + # shift. The packed f32x2 operations keep paired scale groups aligned. + scaled_pair_01 = ffma2( + (s0, s1), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + scaled_pair_23 = ffma2( + (s2, s3), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob0 = cute.math.exp2(scaled_pair_01[0], fastmath=True) + prob1 = cute.math.exp2(scaled_pair_01[1], fastmath=True) + scaled_pair_45 = ffma2( + (s4, s5), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob2 = cute.math.exp2(scaled_pair_23[0], fastmath=True) + prob3 = cute.math.exp2(scaled_pair_23[1], fastmath=True) + # Accumulate the local softmax sums while packing P for the BMM2 operand. + local_sum_pair_01 = fadd2((prob0, prob1), (prob2, prob3)) + packed_p_0 = _pack_float4_to_fp8_e4m3(prob0, prob1, prob2, prob3) + scaled_pair_67 = ffma2( + (s6, s7), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob4 = cute.math.exp2(scaled_pair_45[0], fastmath=True) + prob5 = cute.math.exp2(scaled_pair_45[1], fastmath=True) + prob6 = cute.math.exp2(scaled_pair_67[0], fastmath=True) + prob7 = cute.math.exp2(scaled_pair_67[1], fastmath=True) + local_sum_pair_45 = fadd2((prob4, prob5), (prob6, prob7)) + packed_p_1 = _pack_float4_to_fp8_e4m3(prob4, prob5, prob6, prob7) + local_sum_pair = fadd2(local_sum_pair_01, local_sum_pair_45) + + return packed_p_0, packed_p_1, local_sum_pair[0], local_sum_pair[1] + + +@cute.jit +def _compute_fp8_p_regs_and_local_sums_dense( + scale_softmax_log2: Float32, + new_max_0: Float32, + new_max_1: Float32, + s0: Float32, + s1: Float32, + s2: Float32, + s3: Float32, + s4: Float32, + s5: Float32, + s6: Float32, + s7: Float32, +) -> tuple[Int32, Int32, Float32, Float32]: + """Compute dense FP8 P registers and local softmax sums.""" + # Dense path: all S entries are valid, so no NEG_FLT_MAX guard is needed. + neg_scaled_max_pair = ffma2( + (new_max_0, new_max_1), + (-scale_softmax_log2, -scale_softmax_log2), + (_fp8_log2_quant_scale(), _fp8_log2_quant_scale()), + ) + + scaled_pair_01 = ffma2( + (s0, s1), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + scaled_pair_23 = ffma2( + (s2, s3), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob0 = cute.math.exp2(scaled_pair_01[0], fastmath=True) + prob1 = cute.math.exp2(scaled_pair_01[1], fastmath=True) + scaled_pair_45 = ffma2( + (s4, s5), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob2 = cute.math.exp2(scaled_pair_23[0], fastmath=True) + prob3 = cute.math.exp2(scaled_pair_23[1], fastmath=True) + local_sum_pair_01 = fadd2((prob0, prob1), (prob2, prob3)) + packed_p_0 = _pack_float4_to_fp8_e4m3(prob0, prob1, prob2, prob3) + scaled_pair_67 = ffma2( + (s6, s7), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob4 = cute.math.exp2(scaled_pair_45[0], fastmath=True) + prob5 = cute.math.exp2(scaled_pair_45[1], fastmath=True) + prob6 = cute.math.exp2(scaled_pair_67[0], fastmath=True) + prob7 = cute.math.exp2(scaled_pair_67[1], fastmath=True) + local_sum_pair_45 = fadd2((prob4, prob5), (prob6, prob7)) + packed_p_1 = _pack_float4_to_fp8_e4m3(prob4, prob5, prob6, prob7) + local_sum_pair = fadd2(local_sum_pair_01, local_sum_pair_45) + + return packed_p_0, packed_p_1, local_sum_pair[0], local_sum_pair[1] + + +@cute.jit +def _compute_p_values_and_local_sums_dense( + scale_softmax_log2: Float32, + new_max_0: Float32, + new_max_1: Float32, + s0: Float32, + s1: Float32, + s2: Float32, + s3: Float32, + s4: Float32, + s5: Float32, + s6: Float32, + s7: Float32, +) -> tuple[ + Float32, + Float32, + Float32, + Float32, + Float32, + Float32, + Float32, + Float32, + Float32, + Float32, +]: + """Compute dense 16-bit P values and paired local softmax sums.""" + # Dense 16-bit path: compute eight P values and the two local sums without + # per-row validity checks. + neg_scaled_max_pair = fmul2( + (new_max_0, new_max_1), + (-scale_softmax_log2, -scale_softmax_log2), + ) + scaled_pair_01 = ffma2( + (s0, s1), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + scaled_pair_23 = ffma2( + (s2, s3), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob0 = cute.math.exp2(scaled_pair_01[0], fastmath=True) + prob1 = cute.math.exp2(scaled_pair_01[1], fastmath=True) + scaled_pair_45 = ffma2( + (s4, s5), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob2 = cute.math.exp2(scaled_pair_23[0], fastmath=True) + prob3 = cute.math.exp2(scaled_pair_23[1], fastmath=True) + scaled_pair_67 = ffma2( + (s6, s7), + (scale_softmax_log2, scale_softmax_log2), + neg_scaled_max_pair, + ) + prob4 = cute.math.exp2(scaled_pair_45[0], fastmath=True) + prob5 = cute.math.exp2(scaled_pair_45[1], fastmath=True) + prob6 = cute.math.exp2(scaled_pair_67[0], fastmath=True) + prob7 = cute.math.exp2(scaled_pair_67[1], fastmath=True) + local_sum_pair_02 = fadd2((prob0, prob1), (prob2, prob3)) + local_sum_pair_46 = fadd2((prob4, prob5), (prob6, prob7)) + local_sum_pair = fadd2(local_sum_pair_02, local_sum_pair_46) + return ( + prob0, + prob1, + prob2, + prob3, + prob4, + prob5, + prob6, + prob7, + local_sum_pair[0], + local_sum_pair[1], + ) + + +@cute.jit +def _float_to_u32_for_atomic_max(val: Float32) -> Uint32: + """Encode a float so unsigned atomic max preserves float ordering.""" + # Encode signed floats so unsigned atomic max has the same ordering. + bits = prims.mov_b32(val, target_type=Int32) + mask = (bits >> Int32(31)) | Int32(0x80000000) + encoded = bits ^ mask + return prims.mov_b32(encoded, target_type=Uint32) + + +@cute.jit +def _u32_to_float_for_atomic_max(val: Uint32) -> Float32: + """Decode the unsigned atomic-max representation back to float.""" + # Decode the monotonic unsigned representation back to float. + encoded = prims.mov_b32(val, target_type=Int32) + mask = (~(encoded >> Int32(31))) | Int32(0x80000000) + bits = encoded ^ mask + return prims.mov_b32(bits, target_type=Float32) + + +@cute.jit +def _smem_atomic_max_u32(ptr: cute.Pointer, val: Uint32) -> None: + """Atomically update a CTA-scope SMEM max encoded as unsigned int.""" + # Softmax row max is reduced through CTA SMEM using unsigned atomics over + # the encoded float representation. CTA scope is sufficient because all + # participating softmax warps are inside one CTA. + prims.atomicrmw( + prims.AtomicOp.MAX, + ptr, + val, + syncscope=prims.MemScope.CTA, + space=prims.SharedSpace.shared_cta, + ) + + +@cute.jit +def _wspro_reduce_max4( + val0: Float32, + val1: Float32, + val2: Float32, + val3: Float32, + local_row_idx: Int32, +) -> Float32: + """Reduce four independent maxima across four strided warp rows.""" + # The four column groups are interleaved every four lanes. The conditional + # swaps transpose independent scale groups into row ownership before each + # full-warp butterfly; every lane executes all three shuffle operations. + left01 = val0 + right01 = val1 + left23 = val2 + right23 = val3 + if (local_row_idx & Int32(1)) == Int32(0): + tmp = left01 + left01 = right01 + right01 = tmp + tmp = left23 + left23 = right23 + right23 = tmp + left01 = Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=left01, + offset=4, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + left23 = Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=left23, + offset=4, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + reduced01 = cute.math.max(left01, right01, ftz=True) + reduced23 = cute.math.max(left23, right23, ftz=True) + + if local_row_idx < Int32(2): + tmp = reduced01 + reduced01 = reduced23 + reduced23 = tmp + reduced01 = Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=reduced01, + offset=8, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + return cute.math.max(reduced01, reduced23, ftz=True) + + +@cute.jit +def _init_softmax_scratch_u32( + scratch: cutlass.Array, warp_grp_thread_idx: Int32, num_entries: int +) -> None: + """Initialize encoded softmax scratch maxima across the warp group.""" + encoded = _float_to_u32_for_atomic_max(_neg_max_f32()) + # Initialize scratch with a warp-group-thread-index-stepped loop instead + # of a simple `thread < 8` predicate. This avoids the widened store pattern + # that the TS helper currently lowers into. + for scratch_idx in cutlass.range( + warp_grp_thread_idx, Int32(num_entries), Int32(128), unroll=1 + ): + scratch[scratch_idx] = encoded + + +@cute.jit +def _attention_sink_for_local_head( + cfg: FmhaDecodeConfig, + attention_sinks_ptr: cute.Pointer | None, + scale_softmax_log2: Float32, + max_val: Float32, + logical_h_k_idx: Int32, + h_r: Int32, + num_heads_kv: Int32, + local_head_idx: Int32, +) -> Float32: + """Return the attention-sink denominator contribution for one local head.""" + # Attention sinks add a synthetic key/value entry to the normalization + # denominator. Return zero when the feature is disabled so call sites can + # use the same reduction flow. + if cutlass.const_expr(not cfg.use_attention_sinks): + return Float32(0.0) + head_idx = cute.math.min( + logical_h_k_idx * h_r + local_head_idx, + h_r * num_heads_kv - Int32(1), + ) + sink_ptr = cutlass.inttoptr( + attention_sinks_ptr.toint() + cutlass.Int64(head_idx * Int32(4)), + mem_space=1, + dtype=Float32, + ) + sink_val = sink_ptr.load(count=1, alignment=4)[0] + sink_exp = cute.math.exp2( + sink_val * Float32(1.4426950408889634) - max_val * scale_softmax_log2, + fastmath=True, + ) + if cutlass.const_expr(cfg.use_fp8_qkv): + sink_exp = sink_exp * Float32(448.0) + return sink_exp + + +@cute.jit +def _attention_sink_for_scale_idx( + cfg: FmhaDecodeConfig, + attention_sinks_ptr: cute.Pointer | None, + scale_softmax_log2: Float32, + max_val: Float32, + logical_h_k_idx: Int32, + h_r: Int32, + num_heads_kv: Int32, + logical_q_group_idx: Int32, + col_group_idx: Int32, + scale_idx: Constexpr[int], +) -> Float32: + """Map a softmax scale group to its attention-sink contribution.""" + tile_row_idx = ( + Int32((scale_idx // 2) * 8) + col_group_idx * Int32(2) + Int32(scale_idx % 2) + ) + _, local_head_idx = _q_row_token_and_local_head( + cfg, h_r, logical_q_group_idx, tile_row_idx + ) + head_stride = h_r + if cutlass.const_expr(cfg.heads_q_per_kv > 0): + # Multi-token profiles flatten token/head rows differently, but + # attention sinks remain one value per physical Q head. Resolve the + # CTA row through the shared Q geometry and stride by the true ratio. + head_stride = Int32(cfg.heads_q_per_kv) + return _attention_sink_for_local_head( + cfg, + attention_sinks_ptr, + scale_softmax_log2, + max_val, + logical_h_k_idx, + head_stride, + num_heads_kv, + local_head_idx, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py new file mode 100644 index 000000000000..3254de0e79e7 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py @@ -0,0 +1,1077 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared-memory metadata resources for prepared K/V and Softmax work. + +Each KV instruction owns two independent resources. During HEAD, the load task +loads a prepared route, stores it for the matching K/V pair, issues K, then +stages a copy for Softmax. During LOOP, V consumes the previous K/V metadata +before the load task overwrites it with the next route and issues K. Softmax +waits for its staged copy, moves the seven-slot task payload to registers, and +releases the stage before masking. +""" + +from dataclasses import dataclass +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Uint32 +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...._block_sparse.prepared import ( + _PREPARED_ROUTE_IS_FULL_FLAG, + _BlockSparseRouteLayout, +) +from ...placeholder_helpers import _placeholder_smem_array +from ...stage import FmhaStage +from ..fmha_decode_config import FmhaDecodeConfig +from .helpers_common import ( + _TASK_CACHE_SEQ_LEN_KV, + _TASK_CACHE_WARP_IDX, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + Constexpr, + DecodeGenResourceBase, + ResourceVars, + _decode_gen_task_cache, + _keeps_col_base, + _sparse_task_cache_route_begin, + _sparse_task_cache_route_count, + _warp_broadcast_i32, +) + + +# Keeps staging uses the low four bits for structural KV64 validity. Bit 4 +# carries the conservative prepared summary that token masking can be skipped; +# structural, tail, and causal masking remain independent. +_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG = 1 << 4 + +# B8 SWAP origins are eight-token aligned, so bit 0 is free while the route is +# in Softmax's private staging payload. Reusing it avoids adding a word to every +# pipeline stage merely to forward prepare's route-full summary. +_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK = ~_PREPARED_ROUTE_IS_FULL_FLAG + + +@cute.jit +def _paged_sparse_kv_load_coordinate( + logical_origin: Int32, + physical_page_id: Int32, + atom_is_valid: cutlass.Boolean, + page_size: Constexpr[int], +) -> tuple[Int32, Int32]: + """Return the in-page token and physical-page TMA coordinates. + + Invalid atoms still issue their normal TMA transaction so the pipeline + barrier observes a fixed transaction count. Mapping them to the first + token just beyond page zero makes the token coordinate OOB while keeping + the page coordinate itself valid. + """ + + token_in_page = Int32(page_size) + page_id = Int32(0) + if atom_is_valid: + token_in_page = logical_origin % Int32(page_size) + page_id = physical_page_id + return token_in_page, page_id + + +def _swaps_forwards_packed_route_full(cfg: FmhaDecodeConfig) -> bool: + """Whether SWAP forwards prepare's route-full bit in a staged origin. + + Q8/B8 would otherwise need four origin checks, so it forwards the prepared + summary. Larger Q tiles keep straight-line per-score predicates: forwarding + the summary lengthens their hot path more than the skipped checks save. + """ + + return ( + cfg.tile_size_q == 8 + and cfg.kv_block_size == 8 + and not cfg.use_kv_valid_bits + and not cfg.uses_uniform_causal_mask + and not cfg.uses_per_row_causal_mask + ) + + +def _kv_retained_route_words( + route_layout: _BlockSparseRouteLayout, +) -> int: + """Return the aligned SMEM words retained from K issue through V. + + Contiguous routes retain their existing load-origin payload. Paged routes + retain parallel logical-origin and physical-page-ID arrays so every atom + has an independent storage locator; invalid entries use ``(-1, -1)``. + """ + + payload_words = route_layout.logical_origins_per_route + if route_layout.is_paged: + payload_words *= 2 + elif route_layout.logical_origins_per_route == 2: + payload_words += 1 + return ((payload_words + 3) // 4) * 4 + + +@dataclass(frozen=True) +class _BlockSparseSoftmaxStagingLayout: + """Layout of the staged cross-warp Softmax metadata payload. + + Keeps retains all route origins, a flags word, alignment padding, and the + optional K32 token words. KV256 consumers then select the four words owned + by their spatial half. SWAP stores execution-ordered origins followed by + optional logical K32 token words, one for each consumer warp. Its + noncausal Q8/B8 profile without token bits packs route-full into the + otherwise-zero low bit of each warp's first aligned origin. + """ + + # Logical-origin scalars staged for one complete KV route. + num_origin_words: int + # SWAP origins consumed by one Softmax warp; Keeps retains all origins. + origins_per_warp: int + route_flags_word_offset: int | None + token_words_word_offset: int | None + stage_stride_words: int + total_words: int + + @property + def size_bytes(self) -> int: + """Return the 16-byte-aligned staged allocation size.""" + + return self.total_words * 4 + + @staticmethod + def create( + *, + use_keeps_mma_ab: bool, + route_layout: _BlockSparseRouteLayout, + num_stages: int, + ) -> "_BlockSparseSoftmaxStagingLayout": + """Build a stage-count-dependent layout for Softmax metadata.""" + + kv_route_size = route_layout.kv_route_size + atom_size = route_layout.atom_size + has_token_bits = route_layout.has_token_bits + if use_keeps_mma_ab: + num_origin_words = kv_route_size // atom_size + assert num_origin_words <= 4, ( + "Keeps softmax staging supports at most four route origins" + ) + origins_per_warp = num_origin_words + route_flags_word_offset = num_origin_words + aligned_payload_words = ((route_flags_word_offset + 1 + 3) // 4) * 4 + token_words_word_offset = aligned_payload_words if has_token_bits else None + token_words = kv_route_size // 32 if has_token_bits else 0 + stage_stride_words = ((aligned_payload_words + token_words + 3) // 4) * 4 + else: + softmax_atom_size = min(atom_size, 32) + num_origin_words = kv_route_size // softmax_atom_size + origins_per_warp = 32 // softmax_atom_size + route_flags_word_offset = None + token_words_word_offset = num_origin_words if has_token_bits else None + stage_stride_words = num_origin_words + ( + kv_route_size // 32 if has_token_bits else 0 + ) + return _BlockSparseSoftmaxStagingLayout( + num_origin_words=num_origin_words, + origins_per_warp=origins_per_warp, + route_flags_word_offset=route_flags_word_offset, + token_words_word_offset=token_words_word_offset, + stage_stride_words=stage_stride_words, + total_words=num_stages * stage_stride_words, + ) + + +@dataclass(kw_only=True) +class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): + """Pipeline-free route metadata retained from one K issue through V. + + ``route_metadata`` points at the first prepared GMEM record. Resolution + returns logical origins to masking consumers. The private SMEM copy keeps + contiguous load origins or paged ``(logical origin, physical page ID)`` + pairs through the matching V issue. Invalid atoms are retained as safe + storage-specific OOB coordinates. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("resolved_origin0_slot", Int32, Int32(0), "First logical origin."), + ("resolved_origin1_slot", Int32, Int32(0), "Second logical origin."), + ( + "resolved_atom_validity_slot", + Int32, + Int32(0), + "Fine lane validity or coarse two-fragment validity mask.", + ), + ( + "route_record_word_offset_slot", + Int32, + Int32(-1), + "Metadata-relative record offset, or -1 for a dummy route.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + inst_id: Constexpr[int] = 0 + route_metadata: cute.Pointer | None = None + route_layout: Constexpr[_BlockSparseRouteLayout | None] = None + tma_oob_origin: Int32 = None + _retained_route_words: Constexpr[int] = 0 + _alloc: Constexpr[SmemAllocation | None] = None + _smem_words: cutlass.Array = None + resolved_origin0_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + resolved_origin1_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + resolved_atom_validity_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + route_record_word_offset_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def __post_init__(self) -> None: + """Derive the retained K/V payload from the prepared route layout.""" + + assert self.route_layout is not None + assert self.route_layout.is_paged == self.cfg.use_paged_kv + self._retained_route_words = _kv_retained_route_words(self.route_layout) + super().__post_init__() + + def _init_placeholder_state(self) -> None: + """Create shape-correct K/V metadata SMEM for task-graph tracing.""" + + self._smem_words = _placeholder_smem_array(Int32, self._retained_route_words) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate one aligned instruction-local K/V metadata slot.""" + + if self._alloc is None: + self._alloc = SmemAllocation( + name=self.name, + size_bytes=self._retained_route_words * 4, + alignment=16, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """K/V route metadata uses SMEM only.""" + + return [] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind the K/V metadata allocation on the load warp.""" + + if cutlass.const_expr(context is not None and context.smem_base is not None): + self._smem_words = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=Int32, + shape=(self._retained_route_words,), + addrspace=3, + ) + return {} + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Bind producer-local K/V metadata before the first resolution.""" + + self._create_initial_task_locals(stage_info.context) + + @cute.jit + def _prepared_route_logical_origin( + self, + route_record_word_offset: Int32, + atom_idx: Int32, + ) -> Int32: + """Load one logical KV-token origin from a prepared GMEM record.""" + + assert self.route_metadata is not None + return Int32(self.route_metadata[route_record_word_offset + atom_idx]) + + @cute.jit + def _prepared_route_physical_page_id_if_valid( + self, + route_record_word_offset: Int32, + atom_idx: Int32, + ) -> Int32: + """Load a page ID only when the prepared-record offset is valid.""" + + assert self.route_metadata is not None + assert self.route_layout.is_paged + physical_page_id = Int32(-1) + if route_record_word_offset >= Int32(0): + physical_page_id = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.physical_page_ids_word_offset) + + atom_idx + ] + ) + return physical_page_id + + @consumer_work( + returns=( + resolved_origin0_slot, + resolved_origin1_slot, + resolved_atom_validity_slot, + route_record_word_offset_slot, + ) + ) + @cute.jit + def resolve_route( + self, stage_info: StageInfo, *, section: Constexpr[FmhaStage] + ) -> tuple[Int32, Int32, Int32, Int32]: + """Load this resource instance's real or dummy prepared KV route.""" + + assert self.route_metadata is not None + task_cache = _decode_gen_task_cache(stage_info) + row_route_begin = _sparse_task_cache_route_begin(task_cache) + route_count = _sparse_task_cache_route_count(task_cache) + # HEAD publishes one route per instruction. LOOP starts after those + # two publications, hence the one-based loop offset below. Keeping the + # constexpr branch local lets the task scheduler specialize each work + # clone together with its HEAD/LOOP section. + if cutlass.const_expr(section == FmhaStage.Head): + route_idx = Int32(self.inst_id) + else: + route_idx = (stage_info.loop_offset + Int32(1)) * Int32( + self.cfg.num_insts_kv + ) + Int32(self.inst_id) + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + route_record_word_offset = Int32(-1) + if route_idx < route_count: + route_record_word_offset = (row_route_begin + route_idx) * Int32( + self.route_layout.route_metadata_stride_words + ) + route_record_word_offset = cute.arch.make_warp_uniform(route_record_word_offset) + + num_logical_origins = self.route_layout.logical_origins_per_route + uses_two_fragment_route = num_logical_origins == 2 + # Keep the validity load adjacent to the lane-distributed origins. + # For the maximal 32-origin layout, lane 31 can retain its origin and + # load the independent validity scalar before the warp broadcasts it. + valid_mask_lane = min(num_logical_origins, 31) + logical_origin = Int32(-1) + atom_valid_mask = Int32(0) + route_record_is_valid = route_record_word_offset >= Int32(0) + + if route_record_is_valid: + if lane_idx < Int32(num_logical_origins): + logical_origin = self._prepared_route_logical_origin( + route_record_word_offset, + lane_idx, + ) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if lane_idx == Int32(valid_mask_lane): + atom_valid_mask = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.atom_valid_mask_word_offset) + ] + ) + + if cutlass.const_expr(uses_two_fragment_route): + origin0 = _warp_broadcast_i32(logical_origin, 0) + origin1 = _warp_broadcast_i32(logical_origin, 1) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + # The validity word shares the prepared record's cache line + # with fields consumed shortly afterward by Softmax. + atom_valid_mask = _warp_broadcast_i32(atom_valid_mask, valid_mask_lane) + else: + # Unmasked metadata needs no later fields. Recover validity + # from the invalid-origin sentinel and avoid the dead load. + origin_is_valid = cutlass.Boolean( + lane_idx < Int32(2) and logical_origin >= Int32(0) + ) + atom_valid_mask = Int32(cute.arch.vote_ballot_sync(origin_is_valid)) + return origin0, origin1, atom_valid_mask, route_record_word_offset + + # Wider routes stay lane-distributed: each active lane carries only + # its origin and validity through the existing three-scalar K/V ABI. + valid = cutlass.Boolean(False) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + atom_valid_mask = _warp_broadcast_i32(atom_valid_mask, valid_mask_lane) + if lane_idx < Int32(num_logical_origins): + valid = (atom_valid_mask & (Int32(1) << lane_idx)) != Int32(0) + else: + if lane_idx < Int32(num_logical_origins): + valid = logical_origin >= Int32(0) + return logical_origin, Int32(0), Int32(valid), route_record_word_offset + + @producer_work + @cute.jit + def store_route( + self, + stage_info: StageInfo, + *, + resolved_origin0: Int32, + resolved_origin1: Int32, + resolved_atom_validity: Int32, + route_record_word_offset: Int32, + ) -> None: + """Retain storage coordinates for the matching K/V load pair.""" + + del stage_info + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + num_origins = self.route_layout.logical_origins_per_route + if cutlass.const_expr(self.route_layout.is_paged): + if lane_idx < Int32(num_origins): + logical_origin = Int32(resolved_origin0) + atom_is_valid = resolved_atom_validity != Int32(0) + if cutlass.const_expr(num_origins == 2): + if lane_idx == Int32(0): + logical_origin = resolved_origin0 + else: + logical_origin = resolved_origin1 + atom_is_valid = ( + resolved_atom_validity & (Int32(1) << lane_idx) + ) != Int32(0) + physical_page_id = Int32(-1) + if atom_is_valid: + physical_page_id = self._prepared_route_physical_page_id_if_valid( + route_record_word_offset, + lane_idx, + ) + else: + logical_origin = Int32(-1) + self._smem_words[lane_idx] = logical_origin + self._smem_words[ + Int32(self.route_layout.physical_page_ids_word_offset) + lane_idx + ] = physical_page_id + elif cutlass.const_expr(num_origins == 2): + if lane_idx == Int32(0): + self._smem_words[Int32(0)] = resolved_origin0 + self._smem_words[Int32(1)] = resolved_origin1 + self._smem_words[ + Int32(self.route_layout.atom_valid_mask_word_offset) + ] = resolved_atom_validity + else: + if lane_idx < Int32(self.route_layout.logical_origins_per_route): + load_origin = Int32(resolved_origin0) + if resolved_atom_validity == Int32(0): + # Fine-route K and V both consume this retained value. + # Materialize their TensorMap OOB coordinate once here + # instead of rechecking the invalid-origin sentinel for + # every atom copy in both producer passes. + load_origin = Int32(self.tma_oob_origin) + self._smem_words[lane_idx] = load_origin + # K consumes this slot immediately, while V consumes it at the start + # of the next cadence. Both execute in this warp, so a warp fence is + # sufficient; no cross-warp mbarrier belongs here. + cute.arch.sync_warp() + + @cute.jit + def route_tma_coordinate( + self, + atom_idx: Int32, + logical_b_idx: Int32, + ) -> tuple[Int32, Int32]: + """Load one retained sparse atom and return its TensorMap coordinates.""" + + logical_origin = Int32(self._smem_words[atom_idx]) + if cutlass.const_expr(not self.route_layout.is_paged): + # Invalid contiguous atoms were normalized to the TensorMap OOB + # token coordinate when the route was retained, so their storage + # coordinate remains the live batch index. + return logical_origin, logical_b_idx + + physical_page_id = Int32( + self._smem_words[ + Int32(self.route_layout.physical_page_ids_word_offset) + atom_idx + ] + ) + atom_is_valid = cutlass.Boolean( + logical_origin >= Int32(0) and physical_page_id >= Int32(0) + ) + return _paged_sparse_kv_load_coordinate( + logical_origin, + physical_page_id, + atom_is_valid, + self.route_layout.paged_page_size, + ) + + @cute.jit + def route_atom_valid_mask(self) -> Int32: + """Load the retained route's two-fragment validity mask.""" + + assert not self.route_layout.is_paged + assert self.route_layout.logical_origins_per_route == 2 + return Int32( + self._smem_words[Int32(self.route_layout.atom_valid_mask_word_offset)] + ) + + +@dataclass(kw_only=True) +class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): + """Staged route and token metadata consumed by one Softmax group. + + ``inst_id`` identifies which of the two Softmax pipelines owns this + resource. Route resolution belongs to the paired K/V resource, so the + producer passes the resolved payload explicitly instead of recomputing it. + For Keeps, every route token word moves through SMEM without a + data-dependent branch; each consumer receives at most four words through + the stable task-local ABI. A runtime route-full bit can skip per-score token + predicates while leaving structural masking independent. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("softmax_origin0_slot", Int32, Int32(0), "Loaded first logical origin."), + ("softmax_origin1_slot", Int32, Int32(0), "Loaded second logical origin."), + ( + "softmax_route_flags_slot", + Int32, + Int32(0), + "Keeps route flags or SWAP's third logical origin.", + ), + ( + "softmax_token_word0_slot", + Uint32, + Uint32(0xFFFFFFFF), + "Keeps token word 0 or SWAP's fourth origin as unsigned bits.", + ), + ( + "softmax_token_word1_slot", + Uint32, + Uint32(0xFFFFFFFF), + "Keeps token word 1 or SWAP's packed logical K32 token word.", + ), + ( + "softmax_token_word2_slot", + Uint32, + Uint32(0xFFFFFFFF), + "Loaded third Keeps token word or SWAP's B8 route-full summary.", + ), + ( + "softmax_token_word3_slot", + Uint32, + Uint32(0xFFFFFFFF), + "Loaded fourth Keeps token-validity word; unused by SWAP.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + inst_id: Constexpr[int] = 0 + route_metadata: cute.Pointer | None = None + staging_layout: Constexpr[_BlockSparseSoftmaxStagingLayout | None] = None + route_layout: Constexpr[_BlockSparseRouteLayout | None] = None + _alloc: Constexpr[SmemAllocation | None] = None + _smem_words: cutlass.Array = None + softmax_origin0_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + softmax_origin1_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + softmax_route_flags_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + softmax_token_word0_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + softmax_token_word1_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + softmax_token_word2_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + softmax_token_word3_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def __post_init__(self) -> None: + """Derive the staged metadata layout.""" + + assert self.route_layout is not None + assert self.route_layout.is_paged == self.cfg.use_paged_kv + self.staging_layout = _BlockSparseSoftmaxStagingLayout.create( + use_keeps_mma_ab=self.cfg.use_keeps_mma_ab, + route_layout=self.route_layout, + num_stages=self.pipeline_config.num_stages, + ) + super().__post_init__() + + def _init_placeholder_state(self) -> None: + """Create shape-correct staged SMEM for task-graph tracing.""" + + self._smem_words = _placeholder_smem_array( + Int32, self.staging_layout.total_words + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate one metadata payload per configured pipeline stage.""" + + if self._alloc is None: + self._alloc = SmemAllocation( + name=self.name, + size_bytes=self.staging_layout.size_bytes, + alignment=16, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Softmax route metadata uses SMEM and registers only.""" + + return [] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind the staged allocation on producer and consumer tasks.""" + + if cutlass.const_expr(context is not None and context.smem_base is not None): + self._smem_words = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=Int32, + shape=(self.staging_layout.total_words,), + addrspace=3, + ) + return {} + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Bind producer-side metadata storage before the first route.""" + + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_read_state(self, stage_info: StageInfo) -> None: + """Bind consumer-side metadata storage before the first wait.""" + + self._create_initial_task_locals(stage_info.context) + + @cute.jit + def _producer_stage_base(self, stage_info: StageInfo) -> Int32: + """Return the producer stage selected by the task scheduler.""" + + return stage_info.stage_idx * Int32(self.staging_layout.stage_stride_words) + + @cute.jit + def _consumer_stage_base(self) -> Int32: + """Return the consumer stage selected by the latest wait.""" + + return self.consumer_work_stage * Int32(self.staging_layout.stage_stride_words) + + @cute.jit + def _store_route_swaps( + self, + stage_info: StageInfo, + resolved_origin0: Int32, + resolved_origin1: Int32, + resolved_atom_validity: Int32, + route_record_word_offset: Int32, + ) -> None: + """Stage SWAP origins and optional logical-K32 token metadata. + + The noncausal Q8/B8 profile without token bits also packs prepare's + route-full summary into bit 0 of each warp's first aligned origin. + """ + + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + stage_base = self._producer_stage_base(stage_info) + task_cache = _decode_gen_task_cache(stage_info) + seq_len_kv = Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]) + route_record_is_valid = route_record_word_offset >= Int32(0) + + packed_route_full = Int32(0) + if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): + if lane_idx == Int32(0) and route_record_is_valid: + packed_route_full = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.route_flags_word_offset) + ] + ) & Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + packed_route_full = _warp_broadcast_i32(packed_route_full, 0) + + softmax_origin = Int32(-1) + if cutlass.const_expr(self.cfg.kv_block_size < 64): + if lane_idx < Int32(self.staging_layout.num_origin_words): + softmax_origin = Int32(resolved_origin0) + if resolved_atom_validity == Int32(0): + softmax_origin = Int32(-1) + if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): + # Replicate route-full in each K32 slice's first origin; + # B8 alignment leaves bit 0 free for the summary. + if lane_idx % Int32(self.staging_layout.origins_per_warp) == Int32( + 0 + ): + softmax_origin = ( + softmax_origin & Int32(_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK) + ) | packed_route_full + self._smem_words[stage_base + lane_idx] = softmax_origin + else: + # SWAP with a coarse KV atom expands the two resolved KV64 + # fragments into the four logical K32 origins consumed by its + # four softmax warps. + if lane_idx < Int32(4): + fragment_idx = lane_idx >> Int32(1) + softmax_origin = Int32(resolved_origin0) + valid = (resolved_atom_validity & Int32(1)) != Int32(0) + if fragment_idx == Int32(1): + softmax_origin = Int32(resolved_origin1) + valid = (resolved_atom_validity & Int32(2)) != Int32(0) + softmax_origin = softmax_origin + (lane_idx & Int32(1)) * Int32(32) + if not valid or softmax_origin >= seq_len_kv: + softmax_origin = Int32(-1) + self._smem_words[stage_base + lane_idx] = softmax_origin + + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + assert self.route_metadata is not None + assert self.route_layout.token_words_word_offset is not None + assert self.staging_layout.token_words_word_offset is not None + if lane_idx < Int32(self.route_layout.token_words_per_route): + logical_word = Uint32(0) + if route_record_is_valid: + logical_word = Uint32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.token_words_word_offset) + + lane_idx + ] + ) + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + lane_idx + ] = Int32(logical_word) + cute.arch.sync_warp() + + @cute.jit + def _store_route_keeps( + self, + stage_info: StageInfo, + resolved_origin0: Int32, + resolved_origin1: Int32, + resolved_atom_validity: Int32, + route_record_word_offset: Int32, + ) -> None: + """Stage a Keeps route and its already prepared token metadata. + + Origins and token words are lane-distributed, so KV128 and KV256 use + the same producer shape. The consumer later selects the two KV64 + origins and K32 words owned by its KV256 spatial half. + """ + + assert self.staging_layout.route_flags_word_offset is not None + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + route_record_is_valid = route_record_word_offset >= Int32(0) + num_origins = self.route_layout.logical_origins_per_route + route_flags = Int32(resolved_atom_validity) + if cutlass.const_expr(num_origins > 2): + route_flags = Int32( + cute.arch.vote_ballot_sync( + lane_idx < Int32(num_origins) and resolved_atom_validity != Int32(0) + ) + ) + token_word = Uint32(0) + route_token_mask_is_full = cutlass.Boolean(False) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + assert self.route_metadata is not None + assert self.route_layout.token_words_word_offset is not None + gmem_route_flags = Int32(0) + if lane_idx == Int32(0) and route_record_is_valid: + gmem_route_flags = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.route_flags_word_offset) + ] + ) + gmem_route_flags = _warp_broadcast_i32(gmem_route_flags, 0) + # Prepared bit 0 summarizes the whole route. Staged low bits are + # already fragment validity, so remap the summary above them. + route_token_mask_is_full = cutlass.Boolean( + (gmem_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) != Int32(0) + ) + if ( + lane_idx < Int32(self.route_layout.token_words_per_route) + and route_record_is_valid + ): + token_word = Uint32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.token_words_word_offset) + + lane_idx + ] + ) + + stage_base = self._producer_stage_base(stage_info) + if cutlass.const_expr(num_origins == 2): + if lane_idx == Int32(0): + self._smem_words[stage_base] = Int32(resolved_origin0) + self._smem_words[stage_base + Int32(1)] = Int32(resolved_origin1) + else: + if lane_idx < Int32(num_origins): + self._smem_words[stage_base + lane_idx] = Int32(resolved_origin0) + if lane_idx == Int32(0): + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + route_flags = route_flags | ( + Int32(route_token_mask_is_full) + * Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) + ) + self._smem_words[ + stage_base + Int32(self.staging_layout.route_flags_word_offset) + ] = route_flags + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + assert self.staging_layout.token_words_word_offset is not None + if lane_idx < Int32(self.route_layout.token_words_per_route): + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + lane_idx + ] = Int32(token_word) + cute.arch.sync_warp() + + @producer_work + @cute.jit + def store_route( + self, + stage_info: StageInfo, + *, + resolved_origin0: Int32, + resolved_origin1: Int32, + resolved_atom_validity: Int32, + route_record_word_offset: Int32, + ) -> None: + """Store one resolved route and its optional token words in a stage.""" + + if cutlass.const_expr(self.cfg.use_keeps_mma_ab): + self._store_route_keeps( + stage_info, + resolved_origin0, + resolved_origin1, + resolved_atom_validity, + route_record_word_offset, + ) + else: + self._store_route_swaps( + stage_info, + resolved_origin0, + resolved_origin1, + resolved_atom_validity, + route_record_word_offset, + ) + + @cute.jit + def _load_route_swaps_values( + self, stage_info: StageInfo + ) -> tuple[Int32, Int32, Int32, Uint32, Uint32, Uint32]: + """Load one SWAP warp's logical KV origins and optional token mask. + + ``origin0..3`` are logical KV-token atom bases assigned to + this Softmax warp's logical K32 slice; unused or invalid origins are + negative. To preserve the shared seven-slot task ABI, origin 2/3 + subsequently travel through the shared route-flags/token-word-0 slots. + Token-word 1 carries the logical K32 mask, token-word 2 carries + route-full, and token-word 3 is unused. + """ + + stage_base = self._consumer_stage_base() + local_warp_idx = Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_IDX]) + warp_origin_base = stage_base + local_warp_idx * Int32( + self.staging_layout.origins_per_warp + ) + origin0 = Int32(self._smem_words[warp_origin_base]) + origin1 = Int32(-1) + origin2 = Int32(-1) + origin3 = Int32(-1) + if cutlass.const_expr(self.staging_layout.origins_per_warp >= 2): + origin1 = Int32(self._smem_words[warp_origin_base + Int32(1)]) + if cutlass.const_expr(self.staging_layout.origins_per_warp >= 3): + origin2 = Int32(self._smem_words[warp_origin_base + Int32(2)]) + if cutlass.const_expr(self.staging_layout.origins_per_warp >= 4): + origin3 = Int32(self._smem_words[warp_origin_base + Int32(3)]) + + token_word = Uint32(0xFFFFFFFF) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + assert self.staging_layout.token_words_word_offset is not None + token_word = Uint32( + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + local_warp_idx + ] + ) + route_flags = Uint32(0) + if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): + route_flags = Uint32(origin0 & Int32(1)) + origin0 = origin0 & Int32(_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK) + return ( + origin0, + origin1, + origin2, + origin3.bitcast(Uint32), + token_word, + route_flags, + ) + + @consumer_work( + returns=( + softmax_origin0_slot, + softmax_origin1_slot, + softmax_route_flags_slot, + softmax_token_word0_slot, + softmax_token_word1_slot, + softmax_token_word2_slot, + softmax_token_word3_slot, + ) + ) + @cute.jit + def load_route( + self, stage_info: StageInfo + ) -> tuple[Int32, Int32, Int32, Uint32, Uint32, Uint32, Uint32]: + """Copy the waited stage to task-local registers before release.""" + + if cutlass.const_expr(not self.cfg.use_keeps_mma_ab): + # Reuse the original seven-slot task ABI: Task7 interprets the + # middle fields as origin2, origin3 bits, and the logical K32 mask. + origin0, origin1, origin2, origin3_bits, token_word, route_flags = ( + self._load_route_swaps_values(stage_info) + ) + return ( + origin0, + origin1, + origin2, + origin3_bits, + token_word, + route_flags, + Uint32(0xFFFFFFFF), + ) + + assert self.staging_layout.route_flags_word_offset is not None + stage_base = self._consumer_stage_base() + stored_route_flags = Int32( + self._smem_words[ + stage_base + Int32(self.staging_layout.route_flags_word_offset) + ] + ) + origin0_idx = Int32(0) + origin1_idx = Int32(1) + route_flags = stored_route_flags + if cutlass.const_expr(self.route_layout.kv_route_size == 256): + # KV256 threads [0, 64) and [64, 128) own alternating KV64 + # origins: (0, 2) and (1, 3), respectively. Remap their validity + # bits into the existing two-origin task-local ABI. + assert self.route_layout.atom_size == 64 + warp_grp_thread_idx = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] + ) + spatial = warp_grp_thread_idx >> Int32(6) + origin0_idx = spatial + origin1_idx = spatial + Int32(2) + valid0 = (stored_route_flags >> origin0_idx) & Int32(1) + valid1 = (stored_route_flags >> origin1_idx) & Int32(1) + route_flags = valid0 | (valid1 << Int32(1)) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + route_flags = route_flags | ( + stored_route_flags & Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) + ) + origin0 = Int32(self._smem_words[stage_base + origin0_idx]) + origin1 = Int32(self._smem_words[stage_base + origin1_idx]) + token_word0 = Uint32(0xFFFFFFFF) + token_word1 = Uint32(0xFFFFFFFF) + token_word2 = Uint32(0xFFFFFFFF) + token_word3 = Uint32(0xFFFFFFFF) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + assert self.staging_layout.token_words_word_offset is not None + if cutlass.const_expr(self.route_layout.kv_route_size == 256): + token_base = Int32(self.staging_layout.token_words_word_offset) + word0_idx = origin0_idx * Int32(2) + word1_idx = origin1_idx * Int32(2) + token_word0 = Uint32( + self._smem_words[stage_base + token_base + word0_idx] + ) + token_word1 = Uint32( + self._smem_words[stage_base + token_base + word0_idx + Int32(1)] + ) + token_word2 = Uint32( + self._smem_words[stage_base + token_base + word1_idx] + ) + token_word3 = Uint32( + self._smem_words[stage_base + token_base + word1_idx + Int32(1)] + ) + elif cutlass.const_expr(self.cfg.tile_size_q == 64): + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + local_word_base = _keeps_col_base( + self.cfg, + lane_idx, + self.cfg.num_s_regs_per_thread, + ) >> Int32(5) + token_word0 = Uint32( + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + local_word_base + ] + ) + token_word1 = Uint32( + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + local_word_base + + Int32(1) + ] + ) + else: + token_word0 = Uint32( + self._smem_words[ + stage_base + Int32(self.staging_layout.token_words_word_offset) + ] + ) + token_word1 = Uint32( + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset + 1) + ] + ) + token_word2 = Uint32( + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset + 2) + ] + ) + token_word3 = Uint32( + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset + 3) + ] + ) + return ( + origin0, + origin1, + route_flags, + token_word0, + token_word1, + token_word2, + token_word3, + ) + + +__all__ = [ + "SmemBlockSparseKvMetadataResource", + "SmemBlockSparseSoftmaxMetadataResource", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py new file mode 100644 index 000000000000..bb27c3c553f0 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py @@ -0,0 +1,1209 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``SmemPResource`` — P operand staging for BMM2. + +Producer (Softmax): converts S in registers to P and publishes per-lane local +sums back through ``TmemSResource``. Keeps Q64/Q128 overlays P on consumed S +columns in TMEM; Swaps retains the SMEM operand layout. Consumer (MmaTask) +publishes the corresponding TMEM address or SMEM descriptor. +""" + +from dataclasses import dataclass +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 +from cutlass.experimental import primitives as prims + +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + WorkAttr, + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ..fmha_decode_config import FmhaDecodeConfig +from ...placeholder_helpers import _placeholder_smem_array +from .helpers_common import ( + Constexpr, + DecodeGenResourceBase, + ResourceVars, + fadd2, + ffma2, + fmul2, + _TASK_CACHE_LANE_IDX, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _TASK_CACHE_WARP_IDX, + _decode_gen_task_cache, + _fp8_log2_quant_scale, + _is_last_loop_iteration, + _keeps_col_base, + _keeps_row_idx, + _keeps_tcgen05_st, + _named_barrier_arrive, + _neg_max_f32, + _pack_float2_to_bf16, + _pack_float2_to_fp16, + _wait_for_mbarrier_phase, +) +from .helpers_output import ( + _keeps_p_smem_block_offset_bytes, + _p_stsm_smem_offset_bytes, + _store_transposed_smem8b, + _store_transposed_smem8b_x2, + _store_transposed_smem8b_x4, +) +from .helpers_softmax import ( + _compute_fp8_p_regs_and_local_sums, + _compute_fp8_p_regs_and_local_sums_dense, + _compute_p_values_and_local_sums_dense, + _pack_float4_to_fp8_e4m3, + _pack_float4_to_fp8_e4m3_inline, +) +from .tmem_s import TmemSResource + + +@dataclass(kw_only=True) +class SmemPResource(DecodeGenResourceBase): + """P operand resource consumed by BMM2. + + Softmax producers convert S to P, store it in the profile's TMEM or SMEM + layout, and publish local sums back to TmemS. Most profiles use the generic + full/empty P pipeline. KV256 instead publishes four independently ready + K32 TMEM fragments; BMM2 consumes those fragments in order, while the + matching TmemO full barrier prevents the next QK from overwriting aliased P. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "p_desc_0_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "P descriptor for VP MMA call 0.", + ), + ( + "p_desc_1_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "P descriptor for VP MMA call 1.", + ), + ( + "p_tmem_addr_0_slot", + Int32, + Int32(0), + "TMEM P address for VP MMA call 0.", + ), + ( + "p_tmem_addr_1_slot", + Int32, + Int32(0), + "TMEM P address for VP MMA call 1.", + ), + ) + inst_id: Constexpr[int] = 0 + cfg: Constexpr[FmhaDecodeConfig] = None + scale_softmax_log2: Float32 = None + use_variable_seqlens_kv: Constexpr[bool] = False + tmem_s_ref: Constexpr[TmemSResource] = None + tmem_o_ref: Constexpr[object] = None + _alloc: Constexpr[SmemAllocation | None] = None + _fragment_ready_alloc: Constexpr[SmemAllocation | None] = None + _tmem_alloc: Constexpr[TmemAllocation | None] = None + _tmem_base_addr: Int32 = None + _smem_base_p: cutlass.Array = None + _smem_base_p_i32: cutlass.Array = None + _fragment_ready: cutlass.Array = None + p_desc_0_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + p_desc_1_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + p_tmem_addr_0_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + p_tmem_addr_1_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def _init_placeholder_state(self) -> None: + """Create placeholder P storage state.""" + self._tmem_base_addr = Int32(0) + self._smem_base_p = _placeholder_smem_array( + self.cfg.q_dtype, + self.cfg.smem_p_tile_bytes // self.cfg.q_dtype_bytes, + ) + self._smem_base_p_i32 = _placeholder_smem_array( + Int32, self.cfg.smem_p_tile_bytes // 4 + ) + self._fragment_ready = _placeholder_smem_array( + Int64, self.cfg.num_softmax_score_fragments + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate P storage or the KV256 fragment-ready barriers.""" + if self.cfg.streams_tmem_p_fragments: + if self._fragment_ready_alloc is None: + self._fragment_ready_alloc = SmemAllocation( + name=f"{self.name}_fragmentReady", + size_bytes=self.cfg.num_softmax_score_fragments * 8, + alignment=16, + ) + return [self._fragment_ready_alloc] + if self.cfg.uses_tmem_p: + return [] + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.smem_p_tile_bytes, + alignment=self.cfg.stensor_align, + ) + return [self._alloc] + + @cute.jit + def _bind_fragment_ready(self, context: ResourceContext | None = None) -> None: + """Bind the one-way KV256 P-ready barriers from the SMEM context.""" + if cutlass.const_expr( + self.cfg.streams_tmem_p_fragments + and context is not None + and context.smem_base is not None + and self._fragment_ready_alloc is not None + ): + self._fragment_ready = cutlass.Array( + context.smem_base.data_ptr() + self._fragment_ready_alloc.offset, + dtype=Int64, + shape=(self.cfg.num_softmax_score_fragments,), + addrspace=3, + ) + + @cute.jit + def create_function_variables( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind and initialize KV256's per-fragment ready barriers.""" + self._bind_fragment_ready(context) + if cutlass.const_expr(self.cfg.streams_tmem_p_fragments): + tidx, _, _ = cute.arch.thread_idx() + producer_warps = ( + self.cfg.softmax0_num_warps + if self.inst_id == 0 + else self.cfg.softmax1_num_warps + ) + if tidx == Int32(0): + for fragment_idx in cutlass.range_constexpr( + self.cfg.num_softmax_score_fragments + ): + prims.mbarrier_init( + self._fragment_ready.data_ptr() + fragment_idx, + producer_warps, + ) + return {} + + @cute.jit + def initialize_runtime_state_internal( + self, + context: ResourceContext | None = None, + captured_schedule: bool = False, + ) -> None: + """Initialize generic resource state and bind fragment barriers.""" + super().initialize_runtime_state_internal(context, captured_schedule) + self._bind_fragment_ready(context) + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Allocate the TMEM columns occupied by the Keeps P operand.""" + if not self.cfg.uses_tmem_p: + return [] + if self._tmem_alloc is None: + self._tmem_alloc = TmemAllocation( + name=f"{self.name}_tmem", + num_columns=self.cfg.tmem_p_cols_per_inst, + ) + return [self._tmem_alloc] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind P storage and initialize operand task locals.""" + self._bind_fragment_ready(context) + if cutlass.const_expr( + not self.cfg.uses_tmem_p + and context is not None + and context.smem_base is not None + ): + # P is materialized in SMEM because BMM2 consumes it as a tcgen05 + # SMEM operand. + smem_base_ptr = context.smem_base.data_ptr() + self._alloc.offset + self._smem_base_p = cutlass.Array( + smem_base_ptr, + dtype=self.cfg.q_dtype, + shape=(self.cfg.smem_p_tile_bytes // self.cfg.q_dtype_bytes,), + addrspace=3, + ) + self._smem_base_p_i32 = cutlass.Array( + smem_base_ptr, + dtype=Int32, + shape=(self.cfg.smem_p_tile_bytes // 4,), + addrspace=3, + ) + if cutlass.const_expr( + self.cfg.uses_tmem_p + and context is not None + and context.tmem_ptr_i32 is not None + ): + self._tmem_base_addr = context.tmem_ptr_i32.load() + return { + "p_desc_0": cutlass.Int64(0), + "p_desc_1": cutlass.Int64(0), + "p_tmem_addr_0": Int32(0), + "p_tmem_addr_1": Int32(0), + } + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Provide P operand slots for one work tile.""" + _ = context + return { + "p_desc_0": cutlass.Int64(0), + "p_desc_1": cutlass.Int64(0), + "p_tmem_addr_0": Int32(0), + "p_tmem_addr_1": Int32(0), + } + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_compute_state(self, stage_info: StageInfo) -> None: + """Initialize producer-side P registers and local sums.""" + # ProdAuxWork: bind the SMEM P tile and reset producer-local P/sum + # state before the softmax producer starts writing this work tile. + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize consumer-side P operand placeholders.""" + # ConsAuxWork: mirror the P storage binding on the BMM2 side so operand + # work can publish a valid descriptor or TMEM address for this tile. + self._create_initial_task_locals(stage_info.context) + + @producer_work + @cute.jit + def compute_p_fragment( + self, + stage_info: StageInfo, + *, + fragment_idx: Constexpr[int], + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> None: + """Convert one KV256 K32 score fragment and publish its TMEM P slice.""" + cfg = self.cfg + assert cfg.streams_tmem_p_fragments + assert not cfg.use_fp8_qkv and cfg.uses_two_inst_tmem_p + assert cfg.softmax_score_fragment_regs == 32 + + new_max = new_max_arr[0] + safe_new_max = new_max + if safe_new_max == _neg_max_f32(): + safe_new_max = Float32(0.0) + minus_max_scale = Float32(-self.scale_softmax_log2 * safe_new_max) + + # Eight independent chains keep the denominator update off one long + # dependency chain. Reuse s_arr for probabilities so only one K32 score + # fragment remains live while P is packed. + sum_chains = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) + for chain_idx in cutlass.range_constexpr(8): + sum_chains[chain_idx] = Float32(0.0) + for pair_idx in cutlass.range_constexpr(16): + value_idx = pair_idx * 2 + p0, p1 = cute.arch.fma_packed_f32x2( + (Float32(s_arr[value_idx]), Float32(s_arr[value_idx + 1])), + (self.scale_softmax_log2, self.scale_softmax_log2), + (minus_max_scale, minus_max_scale), + ) + p0 = Float32(cute.math.exp2(p0, fastmath=True)) + p1 = Float32(cute.math.exp2(p1, fastmath=True)) + s_arr[value_idx] = p0 + s_arr[value_idx + 1] = p1 + chain_idx = (pair_idx & 3) * 2 + sum_chains[chain_idx], sum_chains[chain_idx + 1] = ( + cute.arch.add_packed_f32x2( + (sum_chains[chain_idx], sum_chains[chain_idx + 1]), + (p0, p1), + ) + ) + + # Collapse the eight reduction chains before packing P and publishing + # its barrier. This keeps only one sum scalar live across STTM instead + # of overlapping the full reduction state with packed P and addresses. + sum01 = cute.arch.add_packed_f32x2( + (sum_chains[0], sum_chains[1]), + (sum_chains[2], sum_chains[3]), + ) + sum23 = cute.arch.add_packed_f32x2( + (sum_chains[4], sum_chains[5]), + (sum_chains[6], sum_chains[7]), + ) + total_pair = cute.arch.add_packed_f32x2(sum01, sum23) + local_sum = Float32(total_pair[0] + total_pair[1]) + + packed_p = ( + s_arr.data_ptr().load(count=32, alignment=4).to(cfg.q_dtype).bitcast(Int32) + ) + + fragment_cols = cfg.softmax_score_fragment_regs // 2 + p_tmem_addr = ( + self._tmem_base_addr + + Int32(self._tmem_alloc.offset) + + Int32(fragment_idx * fragment_cols) + ) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(p_tmem_addr, Int32), + packed_p, + offset=cfg.tmem_p_cols_per_inst, + ) + # This lowers to the warp-collective tcgen05.wait::st. The explicit + # proxy fence then makes every lane's completed STTM visible through + # the lane-0 mbarrier publication consumed by the MMA warp. + cute.arch.fence_view_async_tmem_store() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + + # KV256 aliases P with the score tile that produced it. Each softmax + # warp publishes its own rows after the TMEM store drains; BMM2 waits + # for all producer warps before consuming the fragment. + tidx, _, _ = cute.arch.thread_idx() + if (tidx & Int32(31)) == Int32(0): + prims.mbarrier_arrive(self._fragment_ready.data_ptr() + Int32(fragment_idx)) + + if cutlass.const_expr(fragment_idx != 0): + local_sum += self.tmem_s_ref.load_p_local_sum(0) + self.tmem_s_ref.store_p_local_sum(0, local_sum) + + @cute.jit + def _compute_keeps_p( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> None: + """Materialize one non-KV256 row-major Keeps probability tile. + + TQ128 gives each warp-group thread a complete 128-column row. TQ64 + gives paired lanes the low/high 64-column halves of one row. Each lane + writes disjoint packed blocks into the TMEM or SMEM layout consumed by + BMM2. + """ + cfg = self.cfg + # KV256 uses compute_p_fragment so only one K32 score fragment is live. + assert not cfg.streams_tmem_p_fragments + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + num_s_regs = cfg.num_s_regs_per_thread + vector_elements = cfg.keeps_p_smem_vector_elements + num_vector_blocks = num_s_regs // vector_elements + row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) + col_base = _keeps_col_base(cfg, lane_idx, num_s_regs) + p_tmem_stage_base = Int32(0) + if cutlass.const_expr(cfg.uses_tmem_p): + p_tmem_stage_base = ( + self._tmem_base_addr + + Int32(self._tmem_alloc.offset) + + stage_info.stage_idx * cfg.tmem_s_cols + ) + + new_max = new_max_arr[0] + safe_new_max = new_max + if safe_new_max == _neg_max_f32(): + safe_new_max = Float32(0.0) + neg_scaled_max = -self.scale_softmax_log2 * safe_new_max + if cutlass.const_expr(cfg.use_fp8_qkv): + neg_scaled_max += _fp8_log2_quant_scale() + + # Preserve four independent modulo-4 sum chains as two packed pairs, + # without keeping a second 16-value P array live beside the S row. + local_sum_pair_01 = (Float32(0.0), Float32(0.0)) + local_sum_pair_23 = (Float32(0.0), Float32(0.0)) + + # Each vector block is exactly 16 bytes after conversion. Compute and + # pack adjacent pairs directly into their final register payload. + packed_p_regs = cfg.num_packed_p_regs if cfg.uses_two_inst_tmem_p else 4 + packed_p = cutlass.Array(Int32, packed_p_regs, space=cutlass.AddressSpace.rmem) + for block_idx in cutlass.range_constexpr(num_vector_blocks): + s_base = block_idx * vector_elements + packed_base = block_idx * 4 if cfg.uses_two_inst_tmem_p else 0 + if cutlass.const_expr(cfg.use_fp8_qkv): + for packed_idx in cutlass.range_constexpr(4): + val_base = packed_idx * 4 + scaled_pair_01 = ffma2( + ( + s_arr[s_base + val_base], + s_arr[s_base + val_base + 1], + ), + (self.scale_softmax_log2, self.scale_softmax_log2), + (neg_scaled_max, neg_scaled_max), + ) + p_pair_01 = ( + cute.math.exp2(scaled_pair_01[0], fastmath=True), + cute.math.exp2(scaled_pair_01[1], fastmath=True), + ) + local_sum_pair_01 = fadd2(local_sum_pair_01, p_pair_01) + scaled_pair_23 = ffma2( + ( + s_arr[s_base + val_base + 2], + s_arr[s_base + val_base + 3], + ), + (self.scale_softmax_log2, self.scale_softmax_log2), + (neg_scaled_max, neg_scaled_max), + ) + p_pair_23 = ( + cute.math.exp2(scaled_pair_23[0], fastmath=True), + cute.math.exp2(scaled_pair_23[1], fastmath=True), + ) + local_sum_pair_23 = fadd2(local_sum_pair_23, p_pair_23) + packed_p[packed_base + packed_idx] = ( + _pack_float4_to_fp8_e4m3_inline( + p_pair_01[0], + p_pair_01[1], + p_pair_23[0], + p_pair_23[1], + ) + ) + else: + for packed_idx in cutlass.range_constexpr(4): + val_base = packed_idx * 2 + scaled_pair = ffma2( + ( + s_arr[s_base + val_base], + s_arr[s_base + val_base + 1], + ), + (self.scale_softmax_log2, self.scale_softmax_log2), + (neg_scaled_max, neg_scaled_max), + ) + p_pair = ( + cute.math.exp2(scaled_pair[0], fastmath=True), + cute.math.exp2(scaled_pair[1], fastmath=True), + ) + if cutlass.const_expr(packed_idx % 2 == 0): + local_sum_pair_01 = fadd2(local_sum_pair_01, p_pair) + else: + local_sum_pair_23 = fadd2(local_sum_pair_23, p_pair) + if cutlass.const_expr(cfg.use_bf16_qkv): + packed_p[packed_base + packed_idx] = _pack_float2_to_bf16( + p_pair[0], p_pair[1] + ) + else: + packed_p[packed_base + packed_idx] = _pack_float2_to_fp16( + p_pair[0], p_pair[1] + ) + + if cutlass.const_expr(cfg.uses_two_inst_tmem_p): + # Retain the complete packed row and publish it once below. + pass + elif cutlass.const_expr(cfg.uses_tmem_p): + # Each register packs two 16-bit P values. The q64 TMEM store shape + # maps paired half-warps onto the low/high 32-column halves of + # the 64-column UInt32 P tile. + p_tmem_addr = p_tmem_stage_base + Int32(block_idx * 4) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(p_tmem_addr, Int32), + packed_p.data_ptr().load(count=4, alignment=4), + offset=cfg.num_packed_p_regs, + ) + else: + logical_col = col_base + Int32(s_base) + smem_offset_bytes = _keeps_p_smem_block_offset_bytes( + cfg, row_idx, logical_col + ) + smem_dst = self._smem_base_p_i32.subview( + smem_offset_bytes >> Int32(2) + ).data_ptr() + smem_dst.store( + packed_p.data_ptr().load(count=4, alignment=4), alignment=16 + ) + if cutlass.const_expr(cfg.uses_two_inst_tmem_p): + # FP8 publishes a complete row with one x16/x32 STTM. FP16/BF16 + # uses x16 slices to limit Softmax register pressure. This is the + # complete-row Q128/KV128 path; KV256 publishes K32 fragments. + assert cfg.num_packed_p_regs in (16, 32, 64) + regs_per_store = cfg.num_packed_p_regs if cfg.use_fp8_qkv else 16 + assert cfg.num_packed_p_regs % regs_per_store == 0 + for store_idx in cutlass.range_constexpr( + cfg.num_packed_p_regs // regs_per_store + ): + packed_offset = store_idx * regs_per_store + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr( + p_tmem_stage_base + Int32(packed_offset), Int32 + ), + (packed_p.data_ptr() + packed_offset).load( + count=regs_per_store, alignment=4 + ), + # Separate the paired Softmax destinations by one packed + # row (the half-row split for x16/x32 TMEM layouts). + offset=cfg.num_packed_p_regs, + ) + if cutlass.const_expr(cfg.ordered_softmax_early_release): + # Hand the baton over as soon as this group's TMEM store has + # issued: the partner's exp2/pack/TMEM store touch only its own + # registers and TMEM region, so it need not wait for this + # store to drain, the async fence, or the pipeline commit. + # The exp2 phases stay serialized (shared MUFU), but the + # store-drain + commit tail overlaps the partner's wakeup. + _named_barrier_arrive( + cfg.resolved_softmax_order_barrier_threads, + barrier_id=cfg.softmax_order_barrier_id + 1 - self.inst_id, + ) + local_sum_pair0 = fadd2(local_sum_pair_01, local_sum_pair_23) + local_sum = local_sum_pair0[0] + local_sum_pair0[1] + self.tmem_s_ref.store_p_local_sum(0, local_sum) + + # Publish the selected memory view before the task-level P pipeline + # exposes this stage to BMM2. + if cutlass.const_expr(cfg.uses_tmem_p): + cute.arch.fence_view_async_tmem_store() + if cutlass.const_expr(cfg.uses_staged_one_inst_tmem_p): + # Synchronize the D256 producer warp group after its TMEM stores + # are visible. + prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) + else: + # Each producer thread orders its own SMEM P stores with an + # async-proxy fence before its own AsyncUmma producer-commit + # mbarrier arrive. The full barrier counts all 128 softmax + # threads, so the commit itself is the warp-group visibility + # point and no extra named barrier is needed here. + cute.arch.fence_view_async_shared() + + @producer_work + @cute.jit + def compute_p( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> None: + """Compute P from S, stage its BMM2 operand, and publish local sums.""" + cfg = self.cfg + if cutlass.const_expr(cfg.use_keeps_mma_ab): + self._compute_keeps_p( + stage_info, + new_max_arr=new_max_arr, + s_arr=s_arr, + ) + return + # ProdWork: transform the softmax S registers into the P operand layout + # expected by BMM2, while recording the per-scale local sums consumed by + # the denominator update. + # Decode the scheduler cache once so every store path uses the same + # warp/lane ownership for SMEM offsets and STSM swizzles. + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + if cutlass.const_expr(cfg.tile_size_q == 32 and cfg.use_fp8_qkv): + # Tile-Q=32 FP8 fast path: compute E4M3 P registers in the + # same order consumed by the STSM helper, while also capturing + # one local denominator sum per softmax scale group. + packed_p = cutlass.Array( + Int32, cfg.num_packed_p_regs, space=cutlass.AddressSpace.rmem + ) + local_sum = cutlass.Array( + Float32, + cfg.num_softmax_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + if cutlass.const_expr( + not self.use_variable_seqlens_kv + and cfg.total_kv_tiles > 0 + and (cfg.total_kv_tiles % cfg.num_insts_kv) == 0 + and cfg.q_tiles_are_full + ): + # Dense full tiles have no tail/window masking. Use the + # straight-line helper to keep this path compact. + for scale_pair_idx in cutlass.range_constexpr(4): + scale_base = scale_pair_idx * 2 + q32_s_base = scale_pair_idx * 4 + p_lo, p_hi, sum_lo, sum_hi = ( + _compute_fp8_p_regs_and_local_sums_dense( + self.scale_softmax_log2, + new_max_arr[scale_base], + new_max_arr[scale_base + 1], + s_arr[q32_s_base], + s_arr[q32_s_base + 1], + s_arr[q32_s_base + 2], + s_arr[q32_s_base + 3], + s_arr[q32_s_base + 16], + s_arr[q32_s_base + 17], + s_arr[q32_s_base + 18], + s_arr[q32_s_base + 19], + ) + ) + packed_p[scale_pair_idx] = p_lo + packed_p[scale_pair_idx + 4] = p_hi + local_sum[scale_base] = sum_lo + local_sum[scale_base + 1] = sum_hi + else: + # General path handles masked S values from variable + # seqlens, odd tail waves, sliding window, or split-KV. + for scale_pair_idx in cutlass.range_constexpr(4): + scale_base = scale_pair_idx * 2 + q32_s_base = scale_pair_idx * 4 + p_lo, p_hi, sum_lo, sum_hi = _compute_fp8_p_regs_and_local_sums( + self.scale_softmax_log2, + new_max_arr[scale_base], + new_max_arr[scale_base + 1], + s_arr[q32_s_base], + s_arr[q32_s_base + 1], + s_arr[q32_s_base + 2], + s_arr[q32_s_base + 3], + s_arr[q32_s_base + 16], + s_arr[q32_s_base + 17], + s_arr[q32_s_base + 18], + s_arr[q32_s_base + 19], + ) + packed_p[scale_pair_idx] = p_lo + packed_p[scale_pair_idx + 4] = p_hi + local_sum[scale_base] = sum_lo + local_sum[scale_base + 1] = sum_hi + # Publish the local denominator contribution before committing P so + # TmemS can update the online-softmax sum after P materialization. + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + self.tmem_s_ref.store_p_local_sum(scale_idx, local_sum[scale_idx]) + # Store the low and high K halves separately. The x4 helper writes + # transposed E4M3 bytes into the SMEM layout expected by BMM2. + _store_transposed_smem8b_x4( + self._smem_base_p_i32, + packed_p[0], + packed_p[1], + packed_p[2], + packed_p[3], + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.tile_size_kv, + ) + _store_transposed_smem8b_x4( + self._smem_base_p_i32, + packed_p[4], + packed_p[5], + packed_p[6], + packed_p[7], + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.tile_size_kv, + 1, + ) + # The P producer writes SMEM directly with STSM helpers. Fence + # and synchronize the warpgroup before the TS pipeline commits + # the stage to the BMM2 consumer. + cute.arch.fence_view_async_shared() + prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) + return + + if cutlass.const_expr(cfg.tile_size_q == 16 and cfg.use_fp8_qkv): + # Tile-Q=16 FP8 fast path: each helper call handles two softmax + # scale groups and returns the low/high K halves already packed for + # STSM. This avoids keeping all 16 FP32 P values live and packing + # them in a separate pass. + packed_p = cutlass.Array( + Int32, cfg.num_packed_p_regs, space=cutlass.AddressSpace.rmem + ) + local_sum = cutlass.Array( + Float32, + cfg.num_softmax_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + for scale_pair_idx in cutlass.range_constexpr(2): + scale_base = scale_pair_idx * 2 + s_base = scale_pair_idx * 4 + if cutlass.const_expr( + not self.use_variable_seqlens_kv + and cfg.total_kv_tiles > 0 + and (cfg.total_kv_tiles % cfg.num_insts_kv) == 0 + and cfg.q_tiles_are_full + ): + p_lo, p_hi, sum_lo, sum_hi = ( + _compute_fp8_p_regs_and_local_sums_dense( + self.scale_softmax_log2, + new_max_arr[scale_base], + new_max_arr[scale_base + 1], + s_arr[s_base], + s_arr[s_base + 1], + s_arr[s_base + 2], + s_arr[s_base + 3], + s_arr[s_base + 8], + s_arr[s_base + 9], + s_arr[s_base + 10], + s_arr[s_base + 11], + ) + ) + else: + p_lo, p_hi, sum_lo, sum_hi = _compute_fp8_p_regs_and_local_sums( + self.scale_softmax_log2, + new_max_arr[scale_base], + new_max_arr[scale_base + 1], + s_arr[s_base], + s_arr[s_base + 1], + s_arr[s_base + 2], + s_arr[s_base + 3], + s_arr[s_base + 8], + s_arr[s_base + 9], + s_arr[s_base + 10], + s_arr[s_base + 11], + ) + packed_p[scale_pair_idx] = p_lo + packed_p[scale_pair_idx + 2] = p_hi + local_sum[scale_base] = sum_lo + local_sum[scale_base + 1] = sum_hi + + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + self.tmem_s_ref.store_p_local_sum(scale_idx, local_sum[scale_idx]) + _store_transposed_smem8b( + self._smem_base_p_i32, + packed_p.data_ptr().load(count=cfg.num_packed_p_regs, alignment=4), + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.tile_size_kv, + cfg.num_packed_p_regs, + ) + cute.arch.fence_view_async_shared() + prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) + return + + if cutlass.const_expr(cfg.tile_size_q in (16, 32)): + # Generic tile-Q 16/32 path: compute P scalars for each + # softmax scale group, accumulate local sums, then pack/store + # to the BMM2 SMEM layout. + q_repeats = max(cfg.tile_size_q // 8, 1) + num_s_regs = cfg.num_s_regs_per_thread + num_scale_groups = cfg.num_softmax_scale_groups + p_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) + local_sums = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + for idx in cutlass.range_constexpr(num_s_regs): + p_vals[idx] = Float32(0.0) + for idx in cutlass.range_constexpr(num_scale_groups): + local_sums[idx] = Float32(0.0) + + for scale_idx in cutlass.range_constexpr(num_scale_groups): + # Convert each softmax scale group from S to P. Masked rows have + # new_max == -inf and keep their initialized zero P/local_sum. + new_max = new_max_arr[scale_idx] + safe_new_max = new_max + if safe_new_max == _neg_max_f32(): + safe_new_max = Float32(0.0) + neg_scaled_max = -self.scale_softmax_log2 * safe_new_max + if cutlass.const_expr(cfg.use_fp8_qkv): + neg_scaled_max += _fp8_log2_quant_scale() + if new_max != _neg_max_f32(): + repeat_idx = scale_idx // 2 + pair_idx = scale_idx % 2 + generic_s_base = repeat_idx * 4 + pair_idx + for k_pair_idx in cutlass.range_constexpr(4): + if cutlass.const_expr(k_pair_idx < 2): + s_idx = generic_s_base + k_pair_idx * 2 + else: + s_idx = ( + generic_s_base + q_repeats * 4 + (k_pair_idx - 2) * 2 + ) + p_val = cute.math.exp2( + s_arr[s_idx] * self.scale_softmax_log2 + neg_scaled_max, + fastmath=True, + ) + p_vals[s_idx] = p_val + local_sums[scale_idx] += p_val + # Hand off denominator contributions through TmemS. P remains a pure + # MMA operand in SMEM; sums are not reloaded from the P tile. + for scale_idx in cutlass.range_constexpr(num_scale_groups): + self.tmem_s_ref.store_p_local_sum(scale_idx, local_sums[scale_idx]) + + if cutlass.const_expr(cfg.use_fp8_qkv): + # FP8 P is packed four values per register and stored with + # transposed 8-bit helpers so BMM2 sees the tcgen05 layout. + packed_p = cutlass.Array( + Int32, cfg.num_packed_p_regs, space=cutlass.AddressSpace.rmem + ) + for packed_idx in cutlass.range_constexpr(cfg.num_packed_p_regs): + val_base = packed_idx * 4 + packed_p[packed_idx] = _pack_float4_to_fp8_e4m3( + p_vals[val_base], + p_vals[val_base + 1], + p_vals[val_base + 2], + p_vals[val_base + 3], + ) + _store_transposed_smem8b( + self._smem_base_p_i32, + packed_p.data_ptr().load(count=cfg.num_packed_p_regs, alignment=4), + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.tile_size_kv, + cfg.num_packed_p_regs, + ) + # Inline byte stores need an explicit CTA barrier before the + # pipeline stage can be observed by the BMM2 consumer. + cute.arch.fence_view_async_shared() + prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) + return + + # FP16/BF16 P uses stmatrix stores. Each stmatrix group writes + # one 8x8 fragment to the swizzled SMEM tile consumed by UMMA. + regs_p = cutlass.Array( + Int32, + cfg.num_packed_p_regs, + space=cutlass.AddressSpace.rmem, + ) + for pair_idx in cutlass.range_constexpr(cfg.num_packed_p_regs): + val_base = pair_idx * 2 + if cutlass.const_expr(cfg.use_bf16_qkv): + regs_p[pair_idx] = _pack_float2_to_bf16( + p_vals[val_base], p_vals[val_base + 1] + ) + else: + regs_p[pair_idx] = _pack_float2_to_fp16( + p_vals[val_base], p_vals[val_base + 1] + ) + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + for stsm_group_idx in cutlass.range_constexpr(cfg.num_packed_p_regs // 4): + smem_offset_bytes = _p_stsm_smem_offset_bytes( + warp_idx, lane_idx, stsm_group_idx, cfg.tile_size_q + ) + smem_dst = ( + self._smem_base_p_i32.subview((smem_offset_bytes >> 2)) + ).data_ptr() + prims.stmatrix( + smem_dst, + (regs_p.data_ptr() + stsm_group_idx * 4).load(count=4, alignment=4), + prims.MMALayout.COL, + shape=prims.StoreShape.M8N8, + ) + cute.arch.fence_view_async_shared() + return + + if cutlass.const_expr(cfg.use_fp8_qkv): + # Tile-Q=8 FP8 path: compute packed P and local sums directly + # from the eight S registers owned by this lane. + packed_p = cutlass.Array( + Int32, cfg.num_packed_p_regs, space=cutlass.AddressSpace.rmem + ) + local_sum = cutlass.Array( + Float32, + cfg.num_softmax_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + if cutlass.const_expr( + not self.use_variable_seqlens_kv + and not cfg.uses_runtime_q_kv_union + and not cfg.use_split_kv + and cfg.has_odd_kv_tail + and self.inst_id == 1 + ): + # In the static nonsplit profile the final inst1 wave is only + # structural padding. Publish an exact zero contribution so + # the paired instance cannot perturb final normalization. + if _is_last_loop_iteration(stage_info): + packed_p[0] = Int32(0) + packed_p[1] = Int32(0) + local_sum[0] = Float32(0.0) + local_sum[1] = Float32(0.0) + else: + packed_p[0], packed_p[1], local_sum[0], local_sum[1] = ( + _compute_fp8_p_regs_and_local_sums( + self.scale_softmax_log2, + new_max_arr[0], + new_max_arr[1], + s_arr[0], + s_arr[1], + s_arr[2], + s_arr[3], + s_arr[4], + s_arr[5], + s_arr[6], + s_arr[7], + ) + ) + elif cutlass.const_expr( + not self.use_variable_seqlens_kv + and cfg.total_kv_tiles > 0 + and (cfg.total_kv_tiles % cfg.num_insts_kv) == 0 + ): + # Dense full-tile FP8 can produce packed P and sums in one + # straight-line helper because no S entry is masked. + packed_p[0], packed_p[1], local_sum[0], local_sum[1] = ( + _compute_fp8_p_regs_and_local_sums_dense( + self.scale_softmax_log2, + new_max_arr[0], + new_max_arr[1], + s_arr[0], + s_arr[1], + s_arr[2], + s_arr[3], + s_arr[4], + s_arr[5], + s_arr[6], + s_arr[7], + ) + ) + else: + # General FP8 path preserves masks from softmax by letting the + # helper suppress entries whose S value is -inf. + packed_p[0], packed_p[1], local_sum[0], local_sum[1] = ( + _compute_fp8_p_regs_and_local_sums( + self.scale_softmax_log2, + new_max_arr[0], + new_max_arr[1], + s_arr[0], + s_arr[1], + s_arr[2], + s_arr[3], + s_arr[4], + s_arr[5], + s_arr[6], + s_arr[7], + ) + ) + # Publish sums through TmemS; packed E4M3 P bytes use the + # transposed SMEM tile consumed by BMM2. + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + self.tmem_s_ref.store_p_local_sum(scale_idx, local_sum[scale_idx]) + _store_transposed_smem8b_x2( + self._smem_base_p_i32, + packed_p[0], + packed_p[1], + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.tile_size_kv, + ) + else: + # Tile-Q=8 16-bit path: compute P scalars, accumulate local + # sums, pack to 16-bit, and store a matrix tile into SMEM. + local_sum = cutlass.Array( + Float32, + cfg.num_softmax_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + p_vals = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + local_sum[scale_idx] = Float32(0.0) + for p_idx in cutlass.range_constexpr(8): + p_vals[p_idx] = Float32(0.0) + if cutlass.const_expr( + not self.use_variable_seqlens_kv + and not cfg.use_split_kv + and cfg.total_kv_tiles > 0 + and (cfg.total_kv_tiles % cfg.num_insts_kv) == 0 + ): + # The straight-line helper may synthesize P for an entirely + # masked sparse instance whose maximum stayed at -inf. This is + # intentionally safe: reduce_sums' guarded rescale and the + # correction path's uses_instN gate both key on that sentinel + # and discard the instance before its P/O contribution is visible. + p_result = _compute_p_values_and_local_sums_dense( + self.scale_softmax_log2, + new_max_arr[0], + new_max_arr[1], + s_arr[0], + s_arr[1], + s_arr[2], + s_arr[3], + s_arr[4], + s_arr[5], + s_arr[6], + s_arr[7], + ) + for p_idx in cutlass.range_constexpr(8): + p_vals[p_idx] = p_result[p_idx] + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + local_sum[scale_idx] = p_result[8 + scale_idx] + else: + # General tile-Q=8 path preserves masked S entries as zero P + # contribution by skipping groups whose new max stayed -inf. + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + new_max = new_max_arr[scale_idx] + if new_max != _neg_max_f32(): + for pair_idx in cutlass.range_constexpr(2): + p_base = scale_idx + pair_idx * 2 + scaled_pair = fmul2( + ( + self.scale_softmax_log2, + self.scale_softmax_log2, + ), + fadd2( + (s_arr[p_base], s_arr[p_base + 4]), + (-new_max, -new_max), + ), + ) + p_pair = ( + cute.math.exp2(scaled_pair[0], fastmath=True), + cute.math.exp2(scaled_pair[1], fastmath=True), + ) + p_vals[p_base] = p_pair[0] + p_vals[p_base + 4] = p_pair[1] + local_sum[scale_idx] += p_pair[0] + local_sum[scale_idx] += p_pair[1] + # Pack the P scalars to match the dtype consumed by BMM2. + regs_p = cutlass.Array( + Int32, cfg.num_packed_p_regs, space=cutlass.AddressSpace.rmem + ) + if cutlass.const_expr(cfg.use_bf16_qkv): + for reg_idx in cutlass.range_constexpr(cfg.num_packed_p_regs): + val_base = reg_idx * 2 + regs_p[reg_idx] = _pack_float2_to_bf16( + p_vals[val_base], p_vals[val_base + 1] + ) + else: + for reg_idx in cutlass.range_constexpr(cfg.num_packed_p_regs): + val_base = reg_idx * 2 + regs_p[reg_idx] = _pack_float2_to_fp16( + p_vals[val_base], p_vals[val_base + 1] + ) + # Publish the denominator contribution after P has been computed, + # before the SMEM fence exposes the P tile to the downstream MMA. + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + self.tmem_s_ref.store_p_local_sum(scale_idx, local_sum[scale_idx]) + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + # Compute the stmatrix destination matching the P descriptor + # swizzle and store the register fragment. + slice_idx = warp_idx // Int32(2) + warp_idx_in_slice = warp_idx % Int32(2) + mtx_idx = lane_idx // Int32(8) + thr_row_idx = lane_idx % Int32(8) + mtx_col_idx = warp_idx_in_slice * Int32(4) + (mtx_idx % Int32(4)) + smem_offset_bytes = ( + slice_idx * Int32(8 * 128) + + thr_row_idx * Int32(128) + + ((mtx_col_idx ^ thr_row_idx) * Int32(16)) + ) + smem_dst = ( + self._smem_base_p_i32.subview((smem_offset_bytes >> 2)) + ).data_ptr() + prims.stmatrix( + smem_dst, + regs_p.data_ptr().load(count=4, alignment=4), + prims.MMALayout.COL, + shape=prims.StoreShape.M8N8, + ) + cute.arch.fence_view_async_shared() + if cutlass.const_expr(cfg.use_fp8_qkv): + # FP8 P uses inline STSM stores. Synchronize the producer + # warpgroup before the UMMA-consumer pipeline is committed so + # BMM2 cannot observe a partially written P tile. + prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) + + @consumer_work( + returns=( + p_desc_0_slot, + p_desc_1_slot, + p_tmem_addr_0_slot, + p_tmem_addr_1_slot, + ) + ) + @cute.jit + def p_operands( + self, stage_info: StageInfo + ) -> tuple[ + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc, + Int32, + Int32, + ]: + """Publish the stage-specific P operand consumed by BMM2.""" + cfg = self.cfg + p_desc_0 = prims.Tcgen05SmemDesc(0) + p_desc_1 = prims.Tcgen05SmemDesc(0) + p_tmem_addr_0 = Int32(0) + p_tmem_addr_1 = Int32(0) + if cutlass.const_expr(cfg.uses_tmem_p): + # ConsWork: select the physical TMEM stage paired with the P + # pipeline token that was just waited. The allocation aliases the + # stats-free columns of the corresponding S stage. + p_stage_cols = cfg.tmem_s_cols + if cutlass.const_expr(cfg.streams_tmem_p_fragments): + # KV256's four pipeline stages are K32 fragments of one P + # operand, not four independent full S/P stages. + p_stage_cols = cfg.softmax_score_fragment_regs // 2 + p_tmem_addr = self._tmem_base_addr + Int32( + self._tmem_alloc.offset + stage_info.stage_idx * p_stage_cols + ) + if cutlass.const_expr(self.inst_id == 0): + p_tmem_addr_0 = p_tmem_addr + else: + p_tmem_addr_1 = p_tmem_addr + else: + # ConsWork: build the SMEM descriptor for P. Only the descriptor + # slot corresponding to this resource instance is populated; the + # MmaTask receives both slots and selects the active one for BMM2. + p_desc = prims.Tcgen05SmemDesc.build( + self._smem_base_p, + leading_byte_offset=Int32(cfg.tile_size_q * 128), + stride_byte_offset=1024, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) + if cutlass.const_expr(self.inst_id == 0): + p_desc_0 = p_desc + else: + p_desc_1 = p_desc + return p_desc_0, p_desc_1, p_tmem_addr_0, p_tmem_addr_1 + + @consumer_work(returns=p_tmem_addr_0_slot) + @cute.jit + def wait_p_fragment( + self, + stage_info: StageInfo, + *, + fragment_idx: Constexpr[int], + ) -> Int32: + """Wait for and return the next KV256 P-fragment TMEM address.""" + cfg = self.cfg + _ = stage_info + assert cfg.streams_tmem_p_fragments + _wait_for_mbarrier_phase( + self._fragment_ready.data_ptr() + Int32(fragment_idx), + self.tmem_s_ref.producer_state.phase, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + + fragment_cols = cfg.softmax_score_fragment_regs // 2 + p_tmem_addr = self._tmem_base_addr + Int32( + self._tmem_alloc.offset + fragment_idx * fragment_cols + ) + return p_tmem_addr + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def wait_until_reusable_before_qk(self, stage_info: StageInfo) -> None: + """Wait until the previous same-instance PV has stopped reading P. + + KV256 aliases each streamed P instance with its next S accumulator. + The existing two-stage O pipeline commits stage ``inst_id`` only when + the matching PV completes, so its full barrier is also the P-reuse + credit. The S producer phase supplies the generation: the first QK + waits on the initially complete opposite parity, and every later QK + waits for the preceding PV without another commit or barrier. + """ + _ = stage_info + cfg = self.cfg + assert cfg.streams_tmem_p_fragments + assert cfg.o_stages == cfg.num_insts_kv == 2 + barrier = self.tmem_o_ref.pipeline.sync_object_full.get_barrier( + Int32(self.inst_id) + ) + _wait_for_mbarrier_phase(barrier, self.tmem_s_ref.producer_state.phase) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py new file mode 100644 index 000000000000..20bb78fb810c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py @@ -0,0 +1,2186 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SMEM-side resources for FMHA decode TS kernel. + +Holds ``SmemQResource``, ``SmemPageOffsetsKvResource``, and ``SmemKvResource`` +— the Q tile, paged-KV page-table cache, and shared K/V SMEM ring, +respectively. All three are TMA producers and tcgen05-descriptor consumers. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.experimental import primitives as prims + +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...._block_sparse.common import ( + _block_sparse_kv_atom_size, + _prepared_kv_routes_are_block_aligned, +) +from ..fmha_decode_config import FmhaDecodeConfig +from ..fmha_decode_constants import ( + KV_INST0, + KV_INST1, + KV_KIND_K, + KV_KIND_V, + KV_TILE_256_K_SLOT_FOR_SEMANTIC_ATOM, +) +from ...stage import FmhaStage +from ...tensor_map import transform_ragged_coords +from ...placeholder_helpers import ( + _placeholder_local_array, + _placeholder_smem_array, +) +from .helpers_common import ( + _TASK_CACHE_KV_PAGE_IDX_UB, + _TASK_CACHE_KV_RAW_TILE_BASE, + Constexpr, + DecodeGenResourceBase, + ResourceVars, + _clamp_valid_tile_idx, + _decode_gen_task_cache, + _logical_head_batch, + _logical_q_group_idx, + _major_k_stride_bytes, + _q_group_token_base, + _qkv_smem_swizzle, +) +from .helpers_kv_tile_idx import ( + _load_runtime_seq_len_kv, + _num_skipped_kv_tiles, + _runtime_clamp_valid_tile_idx, + _runtime_last_valid_page_idx, + _runtime_split_kv_global_tile_idx, + _static_split_kv_global_tile_idx, +) + +if TYPE_CHECKING: + from .smem_block_sparse_metadata import SmemBlockSparseKvMetadataResource + + +def _paged_sparse_kv_tma_transaction_geometry( + *, + tile_size_kv: int, + kv_atom_size: int, + head_dim_stage: int, + kv_dtype_bytes: int, +) -> tuple[int, int, int]: + """Return fixed copy count, bytes per copy, and bytes per paged K/V load. + + Every logical atom participates, including atoms mapped to an OOB + coordinate. This keeps the TMA pipeline's expected transaction bytes + independent of route validity. + """ + + chunk_head_dim = min(head_dim_stage, 64) + assert tile_size_kv % kv_atom_size == 0 + assert head_dim_stage % chunk_head_dim == 0 + transactions_per_load = (tile_size_kv // kv_atom_size) * ( + head_dim_stage // chunk_head_dim + ) + transaction_bytes = kv_atom_size * chunk_head_dim * kv_dtype_bytes + return ( + transactions_per_load, + transaction_bytes, + transactions_per_load * transaction_bytes, + ) + + +@cute.jit +def _cp_async_bulk_tensor_4d_shared_cta_global_predicated( + dst_mem: cutlass.Array, + tma_desc: cutlass.Pointer, + coordinates: tuple[Int32, Int32, Int32, Int32], + mbar: cutlass.Array, +) -> None: + """Issue one 4-D TMA load without branching the producer warp.""" + c0, c1, c2, c3 = coordinates + cute.arch.inline_ptx( + "cp.async.bulk.tensor.4d.shared::cta.global.tile.mbarrier::complete_tx::bytes " + "[{$r0}], [{$r1}, {{$r3}, {$r4}, {$r5}, {$r6}}], [{$r2}];", + read_only_args=[ + dst_mem.data_ptr().toint(cutlass.Int32), + tma_desc.toint(cutlass.Int64), + mbar.data_ptr().toint(cutlass.Int32), + c0, + c1, + c2, + c3, + ], + predicate=prims.elect_sync(), + ) + + +@cute.jit +def _local_kv_tile_idx_for_section( + cfg: Constexpr[FmhaDecodeConfig], + stage_info: StageInfo, + inst_id: Constexpr[int], + kv_kind: Constexpr[int], + section: Constexpr[FmhaStage], +) -> Int32: + """Return the local K/V tile index implied by the TS HEAD/LOOP/TAIL cadence. + + The schedule names whether a producer is K0/K1/V0/V1, while the phase + defines where that logical tile sits in the staggered decode pipeline: + + - HEAD produces only initial K tiles. + - LOOP produces V for the current MMA iteration and K for the next one. + - TAIL drains the final V tiles after the last loop iteration. + """ + num_insts_kv = Int32(cfg.num_insts_kv) + if cutlass.const_expr(section == FmhaStage.Head): + return Int32(inst_id) + if cutlass.const_expr(section == FmhaStage.Loop): + base = stage_info.loop_offset * num_insts_kv + Int32(inst_id) + if cutlass.const_expr(kv_kind == KV_KIND_V): + return base + return base + num_insts_kv + return stage_info.loop_end * num_insts_kv + Int32(inst_id) + + +@dataclass(kw_only=True) +class SmemQResource(DecodeGenResourceBase): + """Q tile in SMEM. Loaded once in HEAD.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "q_desc_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "SMEM descriptor for the staged Q tile.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + tma_desc_q: cutlass.Pointer | None = None + h_k_idx: Int32 = None + b_idx: Int32 = None + q_group_idx: Int32 = None + q_token_offset: Int32 = None + seq_len_q: Int32 = None + _alloc: Constexpr[SmemAllocation | None] = None + _smem_base_q: cutlass.Array = None + _q_descs: cutlass.Array = None + q_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def _init_placeholder_state(self) -> None: + """Create placeholder Q SMEM and descriptor state for static analysis.""" + self._smem_base_q = _placeholder_smem_array( + self.cfg.q_dtype, + self.cfg.smem_q_tile_elements * self.cfg.q_stages, + ) + self._q_descs = _placeholder_local_array( + cutlass.Int64, self.cfg.q_stages, alignment=8 + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate one staged Q tile buffer per Q pipeline stage.""" + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.smem_q_tile_bytes * self.cfg.q_stages, + alignment=self.cfg.stensor_align, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Q staging uses SMEM only.""" + return [] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind Q SMEM and build per-stage descriptors for producer/consumer use.""" + if cutlass.const_expr(context is not None and context.smem_base is not None): + # Bind the Q SMEM allocation and create one tcgen05 descriptor per + # pipeline stage. Q is loaded in HEAD and reused by all BMM1 calls. + self._smem_base_q = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.q_dtype, + shape=(self.cfg.smem_q_tile_elements * self.cfg.q_stages,), + addrspace=3, + ) + self._q_descs = cutlass.Array( + cutlass.Int64, + self.cfg.q_stages, + space=cutlass.AddressSpace.rmem, + alignment=8, + ) + stage_elems = self.cfg.smem_q_tile_elements + q_head_dim_partition = min(self.cfg.headdim, 64) + # 16-bit Q uses two 64-column TMA loads for headDim=128. FP8 Q + # uses a descriptor stride based on the actual MMA K-major swizzle. + q_leading_bytes = Int32( + self.cfg.tile_size_q * q_head_dim_partition * self.cfg.q_dtype_bytes + ) + leading_byte_offset = q_leading_bytes + stride_byte_offset = Int32(1024) + if cutlass.const_expr(self.cfg.use_fp8_qkv): + q_head_dim_stage = self.cfg.head_dim_kv_stage + q_tile_bytes = Int32( + self.cfg.tile_size_q * q_head_dim_stage * self.cfg.q_dtype_bytes + ) + leading_byte_offset = q_tile_bytes + stride_byte_offset = Int32( + _major_k_stride_bytes(self.cfg.q_dtype_bytes, self.cfg.headdim) + ) + for stage_idx in cutlass.range_constexpr(self.cfg.q_stages): + # Advance the base address for each stage while preserving the + # same swizzled layout parameters. + smem_ptr = self._smem_base_q.subview(stage_idx * stage_elems) + self._q_descs[stage_idx] = prims.Tcgen05SmemDesc.build( + smem_ptr, + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_swizzle(self.cfg), + ) + return {"q_desc": cutlass.Int64(0)} + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Provide the per-work-tile Q descriptor slot.""" + _ = context + return {"q_desc": cutlass.Int64(0)} + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize Q producer-side SMEM pointers and descriptors.""" + # ProdAuxWork: bind the Q SMEM allocation and precompute descriptor + # bases before the load task starts issuing TMA copies. + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize Q consumer-side descriptor state.""" + # ConsAuxWork: create the same descriptor slots on the MMA side so + # q_desc() can return the stage committed by LoadTask. + self._create_initial_task_locals(stage_info.context) + + @cute.jit + def _tma_coords( + self, + logical_q_group_idx: Int32, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + head_dim_offset: Int32, + ) -> tuple[Int32, Int32, Int32, Int32, Int32]: + """Map a logical Q CTA to token-major tensor-map coordinates.""" + cfg = self.cfg + if cutlass.const_expr(cfg.use_variable_seqlens_q): + if cutlass.const_expr(cfg.groups_tokens_heads_q): + token_idx = logical_q_group_idx * Int32(cfg.q_tokens_per_cta) + global_head_idx = logical_h_k_idx * Int32(cfg.heads_q_per_kv) + tokens_per_box = cfg.q_tokens_per_cta + else: + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + token_idx = logical_q_group_idx // head_ctas_per_token + head_cta_idx = logical_q_group_idx - token_idx * head_ctas_per_token + global_head_idx = logical_h_k_idx * Int32( + cfg.heads_q_per_kv + ) + head_cta_idx * Int32(cfg.tile_size_q) + tokens_per_box = 1 + packed_coords = transform_ragged_coords( + ( + head_dim_offset, + global_head_idx, + self.q_token_offset + token_idx, + ), + ragged_dim_idx=2, + ragged_box_size=tokens_per_box, + ragged_extent=self.seq_len_q - token_idx, + ) + return ( + packed_coords[0], + packed_coords[1], + packed_coords[2], + packed_coords[3], + packed_coords[4], + ) + + if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): + token_idx = logical_q_group_idx * Int32(cfg.q_tokens_per_cta) + return ( + head_dim_offset, + Int32(0), + logical_h_k_idx, + token_idx, + logical_b_idx, + ) + + if cutlass.const_expr(cfg.max_seq_len_q == 1): + return ( + head_dim_offset, + logical_q_group_idx * Int32(cfg.tile_size_q), + logical_h_k_idx, + Int32(0), + logical_b_idx, + ) + + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + token_idx = logical_q_group_idx // head_ctas_per_token + head_cta_idx = logical_q_group_idx - token_idx * head_ctas_per_token + return ( + head_dim_offset, + head_cta_idx * Int32(cfg.tile_size_q), + logical_h_k_idx, + token_idx, + logical_b_idx, + ) + + @cute.jit + def _complete_grouped_q_padding(self, stage_info: StageInfo) -> None: + """Account structural grouped rows that have no corresponding TMA.""" + cfg = self.cfg + if cutlass.const_expr( + cfg.groups_tokens_heads_q and cfg.q_manual_padding_rows > 0 + ): + padding_bytes = cfg.q_manual_padding_rows * cfg.headdim * cfg.q_dtype_bytes + prims.mbarrier_complete_tx(stage_info.barrier, Int32(padding_bytes)) + + @producer_work + @cute.jit + def tma_load(self, stage_info: StageInfo) -> None: + """TMA load the current staged Q tile from GMEM to SMEM.""" + cfg = self.cfg + # ProdWork: issue the Q TMA copies for the current pipeline stage. The + # pipeline barrier in stage_info is the handoff to the MMA consumer. + # Resolve logical coordinates from the work tile when persistent + # scheduling is active; otherwise use the static launch values. + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + # Select the SMEM stage owned by this producer fire. The TS pipeline + # barrier attached to stage_info protects this stage until QK consumes it. + stage_elems = cfg.smem_q_tile_elements + stage_base = self._smem_base_q.subview(stage_info.stage_idx * stage_elems) + if cutlass.const_expr(cfg.use_fp8_qkv): + if prims.elect_sync(): + if cutlass.const_expr(cfg.num_head_dim_stages_kv == 1): + # FP8 with one head-dim stage is one tensor copy into the + # complete staged Q tile. + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base, + self.tma_desc_q, + self._tma_coords( + logical_q_group_idx, + logical_h_k_idx, + logical_b_idx, + Int32(0), + ), + stage_info.barrier, + ) + else: + # H256 FP8 stages Q by head-dim slices so each QK head-dim + # stage sees a contiguous SMEM tile. + q_chunk_dim = cfg.head_dim_kv_stage + for q_chunk_idx in cutlass.range_constexpr( + cfg.num_head_dim_stages_kv + ): + head_dim_offset = q_chunk_idx * q_chunk_dim + smem_offset = head_dim_offset * cfg.tile_size_q + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + self.tma_desc_q, + self._tma_coords( + logical_q_group_idx, + logical_h_k_idx, + logical_b_idx, + Int32(head_dim_offset), + ), + stage_info.barrier, + ) + self._complete_grouped_q_padding(stage_info) + else: + # 16-bit Q is staged in 64-column chunks so the SMEM layout + # matches the tcgen05 descriptor swizzle for H64/H128/H256. + chunk_hd = min(cfg.headdim, 64) + num_chunks = cfg.headdim // chunk_hd + chunk_elems = chunk_hd * cfg.tile_size_q + if prims.elect_sync(): + # One elected lane issues all TMA chunks for this Q tile; the + # async pipeline tracks completion through stage_info.barrier. + for chunk_idx in cutlass.range_constexpr(num_chunks): + head_dim_offset = chunk_idx * chunk_hd + smem_offset = chunk_idx * chunk_elems + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + self.tma_desc_q, + self._tma_coords( + logical_q_group_idx, + logical_h_k_idx, + logical_b_idx, + Int32(head_dim_offset), + ), + stage_info.barrier, + ) + self._complete_grouped_q_padding(stage_info) + + @consumer_work(returns=q_desc_slot) + @cute.jit + def q_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Build Q SMEM descriptor for tcgen05 MMA B operand.""" + # ConsWork: return the descriptor for the stage LoadTask committed. + q_desc = prims.Tcgen05SmemDesc(self._q_descs[Int32(stage_info.stage_idx)]) + return q_desc + + @cute.jit + def current_consumer_q_desc(self) -> prims.Tcgen05SmemDesc: + """Return Q's just-waited stage without creating a routed task local. + + Packed persistent skipping guards the complete data path. Routing a Q + descriptor from guarded HEAD into LOOP would make that descriptor a + cross-region task local. Q's advance-on-wait pipeline already records + the selected stage in ``consumer_work_stage``, so QK work can derive + the same descriptor directly from the shared resource state. + """ + stage_idx = Int32(self.state_src.consumer_work_stage) + return prims.Tcgen05SmemDesc(self._q_descs[stage_idx]) + + +@dataclass(kw_only=True) +class SmemKvTileResource(DecodeGenResourceBase): + """Single K or V tile in SMEM for the split K/V decode schedule.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "kv_desc_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "SMEM descriptor for K loads consumed by QK MMA.", + ), + ( + "v_desc_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "SMEM descriptor for V loads consumed by VP MMA.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + tma_desc_k: cutlass.Pointer | None = None + tma_desc_v: cutlass.Pointer | None = None + tma_desc_k_atom: cutlass.Pointer | None = None + tma_desc_v_atom: cutlass.Pointer | None = None + sparse_kv_metadata: "SmemBlockSparseKvMetadataResource | None" = None + page_offsets_kv: "SmemPageOffsetsKvResource | None" = None + seqlens_kv: cute.Pointer | None = None + max_seq_len_kv: Int32 = None + h_k_idx: Int32 = None + b_idx: Int32 = None + q_group_idx: Int32 = None + seq_len_q: Int32 = None + inst_id: Constexpr[int] = 0 + kv_kind: Constexpr[int] = KV_KIND_K + _alloc: Constexpr[SmemAllocation | None] = None + _smem_base_kv: cutlass.Array = None + _desc_base: prims.Tcgen05SmemDesc = None + kv_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + v_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def _init_placeholder_state(self) -> None: + """Create placeholder split K/V SMEM and descriptor state.""" + num_stages = ( + self.pipeline_config.num_stages if self.pipeline_config is not None else 1 + ) + self._smem_base_kv = _placeholder_smem_array( + self.cfg.kv_dtype, + self.cfg.smem_kv_tile_elements * num_stages, + ) + self._desc_base = prims.Tcgen05SmemDesc(0) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate this split K or V resource's staged tile ring.""" + num_stages = ( + self.pipeline_config.num_stages if self.pipeline_config is not None else 1 + ) + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.smem_kv_tile_bytes * num_stages, + alignment=self.cfg.stensor_align, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Split K/V staging uses SMEM only.""" + return [] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind the split K/V SMEM ring and build its base descriptor.""" + if cutlass.const_expr(context is not None and context.smem_base is not None): + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else 1 + ) + self._smem_base_kv = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.kv_dtype, + shape=(self.cfg.smem_kv_tile_elements * num_stages,), + addrspace=3, + ) + kv_tile_bytes = Int32( + self.cfg.tile_size_kv + * self.cfg.head_dim_kv_stage + * self.cfg.kv_dtype_bytes + ) + leading_byte_offset = Int32( + self.cfg.tile_size_kv + * min(self.cfg.head_dim_kv_stage, 64) + * self.cfg.kv_dtype_bytes + ) + stride_byte_offset = Int32(1024) + if cutlass.const_expr(self.cfg.use_fp8_qkv): + leading_byte_offset = kv_tile_bytes + stride_byte_offset = Int32( + _major_k_stride_bytes( + self.cfg.kv_dtype_bytes, self.cfg.head_dim_kv_stage + ) + ) + if cutlass.const_expr( + self.kv_kind == KV_KIND_V + and (self.cfg.use_fp8_qkv or self.cfg.headdim == 64) + ): + leading_byte_offset = Int32(0) + self._desc_base = prims.Tcgen05SmemDesc.build( + self._smem_base_kv, + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_swizzle(self.cfg), + ) + return {"kv_desc": cutlass.Int64(0), "v_desc": cutlass.Int64(0)} + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Provide descriptor slots for this split K/V work tile.""" + _ = context + return {"kv_desc": cutlass.Int64(0), "v_desc": cutlass.Int64(0)} + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize split K/V producer-side SMEM state.""" + # ProdAuxWork: bind this split K or V SMEM ring and descriptor base + # before the load task issues any staged K/V TMA copies. + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize split K/V consumer-side descriptor state.""" + # ConsAuxWork: initialize descriptor slots for the downstream QK/PV MMA + # task that consumes this split K or V resource. + self._create_initial_task_locals(stage_info.context) + + @cute.jit + def _stage_base(self, stage_info: StageInfo) -> cutlass.Array: + """Return the SMEM base for the current split K/V pipeline stage.""" + stage_elems = self.cfg.smem_kv_tile_bytes // self.cfg.kv_dtype_bytes + return self._smem_base_kv.subview(stage_info.stage_idx * stage_elems) + + @cute.jit + def _local_tile_idx( + self, stage_info: StageInfo, section: Constexpr[FmhaStage] + ) -> Int32: + """Map the schedule phase to the local K or V tile index.""" + return _local_kv_tile_idx_for_section( + self.cfg, stage_info, self.inst_id, self.kv_kind, section + ) + + @cute.jit + def _maybe_runtime_tile_idx(self, stage_info: StageInfo, tile_idx: Int32) -> Int32: + """Apply runtime sequence and split-KV transforms to a local tile index.""" + if cutlass.const_expr(self.cfg.use_paged_kv): + return ( + Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE]) + + tile_idx + ) + if cutlass.const_expr( + self.seqlens_kv is None and not self.cfg.uses_runtime_q_kv_union + ): + if cutlass.const_expr(self.cfg.use_split_kv): + tile_idx = _static_split_kv_global_tile_idx( + self.cfg, stage_info, tile_idx + ) + tile_idx = _clamp_valid_tile_idx(self.cfg, tile_idx) + return tile_idx + Int32(self.cfg.static_num_skipped_kv_tiles) + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + self.h_k_idx, + self.b_idx, + ) + logical_q_group_idx = _logical_q_group_idx( + self.cfg, stage_info, self.q_group_idx + ) + q_token_base = _q_group_token_base(self.cfg, logical_q_group_idx) + tile_idx = _runtime_split_kv_global_tile_idx( + self.cfg, + stage_info, + tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + tile_idx = _runtime_clamp_valid_tile_idx( + self.cfg, + tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + return tile_idx + _num_skipped_kv_tiles( + self.cfg, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + + @cute.jit + def _producer_load( + self, + stage_info: StageInfo, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue the staged TMA load for this split K or V tile.""" + cfg = self.cfg + # Resolve the schedule-local K/V tile, logical head/batch coordinates, + # and descriptor kind before selecting paged or dense addressing. + local_tile_idx = self._local_tile_idx(stage_info, section) + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + tma_desc = ( + self.tma_desc_v + if cutlass.const_expr(self.kv_kind == KV_KIND_V) + else self.tma_desc_k + ) + + if cutlass.const_expr(cfg.use_block_sparse): + assert self.sparse_kv_metadata is not None + assert self.tma_desc_k_atom is not None + assert self.tma_desc_v_atom is not None + # The positional TensorMaps keep the decode ABI stable. The + # primary K/V descriptors are KV128 for coarse routes and one atom + # for fine routes. The auxiliary slots always expose the atom + # descriptor and alias the primary descriptor for fine routes. + tma_desc_atom = ( + self.tma_desc_v_atom + if cutlass.const_expr(self.kv_kind == KV_KIND_V) + else self.tma_desc_k_atom + ) + kv_atom_size = _block_sparse_kv_atom_size(cfg.kv_block_size) + head_dim_stage = cfg.head_dim_kv_stage + head_dim_stage_offset = head_dim_stage_idx * head_dim_stage + chunk_hd = min(head_dim_stage, 64) + num_chunks = head_dim_stage // chunk_hd + tile_chunk_elems = chunk_hd * cfg.tile_size_kv + if cutlass.const_expr(cfg.use_paged_kv): + # Paged sparse routes retain an independent physical page ID + # for every logical atom. Never infer adjacency: issue the + # same atom-sized TMA sequence for valid and invalid atoms, + # with invalid coordinates mapped just beyond page zero so + # the fixed mbarrier transaction count is preserved. + atoms_per_route = cfg.tile_size_kv // kv_atom_size + ( + transactions_per_load, + transaction_bytes, + total_transaction_bytes, + ) = _paged_sparse_kv_tma_transaction_geometry( + tile_size_kv=cfg.tile_size_kv, + kv_atom_size=kv_atom_size, + head_dim_stage=head_dim_stage, + kv_dtype_bytes=cfg.kv_dtype_bytes, + ) + assert total_transaction_bytes == cfg.smem_kv_tile_bytes + num_chunks = transactions_per_load // atoms_per_route + chunk_hd = transaction_bytes // (kv_atom_size * cfg.kv_dtype_bytes) + atom_chunk_elems = transaction_bytes // cfg.kv_dtype_bytes + tile_chunk_elems = chunk_hd * cfg.tile_size_kv + if prims.elect_sync(): + stage_base = self._stage_base(stage_info) + for atom_idx in cutlass.range_constexpr(atoms_per_route): + token_coord, storage_coord = ( + self.sparse_kv_metadata.route_tma_coordinate( + Int32(atom_idx), + logical_b_idx, + ) + ) + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + local_tile_offset = chunk_idx * tile_chunk_elems + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview( + local_tile_offset + atom_idx * atom_chunk_elems + ), + tma_desc_atom, + ( + Int32(global_head_dim_offset), + token_coord, + logical_h_k_idx, + storage_coord, + ), + stage_info.barrier, + ) + elif cutlass.const_expr(kv_atom_size == 64): + # A B128-aligned semantic block keeps every route inside one + # BSR entry, so one KV128 TMA is always legal; TMA OOB fill + # handles a partial physical tail. Other coarse blocks may + # join unrelated entries and must prove physical adjacency. + fragment_chunk_elems = chunk_hd * 64 + if prims.elect_sync(): + origin0, _ = self.sparse_kv_metadata.route_tma_coordinate( + Int32(0), + logical_b_idx, + ) + origin0 = Int32(origin0) + atom_valid_mask = Int32( + self.sparse_kv_metadata.route_atom_valid_mask() + ) + valid0 = cutlass.Boolean((atom_valid_mask & Int32(1)) != Int32(0)) + if not valid0: + origin0 = Int32(self.max_seq_len_kv) + stage_base = self._stage_base(stage_info) + if cutlass.const_expr( + _prepared_kv_routes_are_block_aligned( + cfg.kv_block_size, cfg.tile_size_kv + ) + ): + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + local_tile_offset = chunk_idx * tile_chunk_elems + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(local_tile_offset), + tma_desc, + ( + Int32(global_head_dim_offset), + origin0, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + else: + origin1, _ = self.sparse_kv_metadata.route_tma_coordinate( + Int32(1), + logical_b_idx, + ) + origin1 = Int32(origin1) + valid1 = cutlass.Boolean( + (atom_valid_mask & Int32(2)) != Int32(0) + ) + adjacent = (valid0 & valid1) & cutlass.Boolean( + origin1 == origin0 + Int32(64) + ) + if not valid1: + origin1 = Int32(self.max_seq_len_kv) + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + local_tile_offset = chunk_idx * tile_chunk_elems + if adjacent: + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(local_tile_offset), + tma_desc, + ( + Int32(global_head_dim_offset), + origin0, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + else: + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(local_tile_offset), + tma_desc_atom, + ( + Int32(global_head_dim_offset), + origin0, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview( + local_tile_offset + fragment_chunk_elems + ), + tma_desc_atom, + ( + Int32(global_head_dim_offset), + origin1, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + else: + # Fine routes stay fully general: issue one TMA per route + # atom. Retained metadata has already mapped empty slots to + # an OOB origin, avoiding a repeated predicate here. Keeping + # the load policy independent of adjacency is faster for + # irregular top-k rows. + atom_chunk_elems = chunk_hd * kv_atom_size + atoms_per_route = cfg.tile_size_kv // kv_atom_size + if prims.elect_sync(): + stage_base = self._stage_base(stage_info) + # Reuse each retained origin across all head-dimension + # chunks. The copies still target disjoint SMEM regions + # and share one completion barrier, so only issue order + # changes. + origin, _ = self.sparse_kv_metadata.route_tma_coordinate( + Int32(0), + logical_b_idx, + ) + next_origin, _ = self.sparse_kv_metadata.route_tma_coordinate( + Int32(1), + logical_b_idx, + ) + origin = Int32(origin) + next_origin = Int32(next_origin) + for atom_idx in cutlass.range_constexpr(atoms_per_route): + # Keep two origins ahead of TMA issue so the scalar + # LDS -> uniform-register handoff can overlap a full + # atom's asynchronous copies. + future_origin = next_origin + if cutlass.const_expr(atom_idx + 2 < atoms_per_route): + future_origin, _ = ( + self.sparse_kv_metadata.route_tma_coordinate( + Int32(atom_idx + 2), + logical_b_idx, + ) + ) + future_origin = Int32(future_origin) + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + local_tile_offset = chunk_idx * tile_chunk_elems + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview( + local_tile_offset + atom_idx * atom_chunk_elems + ), + tma_desc_atom, + ( + Int32(global_head_dim_offset), + origin, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + origin = next_origin + next_origin = future_origin + + elif cutlass.const_expr(cfg.use_paged_kv): + # Paged-KV path: the page-offset resource has already staged the + # page IDs for this tile window into SMEM. This producer slices the + # page IDs for one tile and emits one TMA per page fragment. + head_dim_stage = cfg.head_dim_kv_stage + head_dim_stage_offset = head_dim_stage_idx * head_dim_stage + page_fragments = cfg.tile_size_kv // cfg.num_tokens_per_page + tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) + if cutlass.const_expr(cfg.use_fp8_qkv): + if prims.elect_sync(): + # FP8 pages are copied as one contiguous head-dim stage per + # page fragment. + stage_base = self._stage_base(stage_info) + page_ids = self.page_offsets_kv.page_ids(tile_idx) + for page_frag in cutlass.range_constexpr(page_fragments): + page_id = Int32(page_ids[page_frag]) + smem_page_offset = Int32( + page_frag * cfg.num_tokens_per_page * head_dim_stage + ) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_page_offset), + tma_desc, + ( + Int32(head_dim_stage_offset), + Int32(0), + logical_h_k_idx, + page_id, + ), + stage_info.barrier, + ) + else: + # 16-bit pages are copied in 64-column chunks so the SMEM + # layout matches the K/V tcgen05 descriptor swizzle. + chunk_hd = min(head_dim_stage, 64) + num_chunks = head_dim_stage // chunk_hd + tile_chunk_elems = chunk_hd * cfg.tile_size_kv + page_chunk_elems = chunk_hd * cfg.num_tokens_per_page + if prims.elect_sync(): + stage_base = self._stage_base(stage_info) + page_ids = self.page_offsets_kv.page_ids(tile_idx) + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + local_tile_offset = chunk_idx * tile_chunk_elems + for page_frag in cutlass.range_constexpr(page_fragments): + page_id = Int32(page_ids[page_frag]) + smem_page_offset = Int32( + local_tile_offset + page_frag * page_chunk_elems + ) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_page_offset), + tma_desc, + ( + Int32(global_head_dim_offset), + Int32(0), + logical_h_k_idx, + page_id, + ), + stage_info.barrier, + ) + else: + # Dense-KV path: map the runtime-resolved tile directly to the + # tensor's sequence dimension and copy one contiguous tile. + tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) + tile_offset = tile_idx * Int32(cfg.tile_size_kv) + head_dim_stage = cfg.head_dim_kv_stage + head_dim_stage_offset = head_dim_stage_idx * head_dim_stage + if cutlass.const_expr(cfg.use_fp8_qkv): + if prims.elect_sync(): + # FP8 dense K/V needs one tensor copy for the active + # head-dim stage. + prims.cp_async_bulk_tensor_shared_cta_global( + self._stage_base(stage_info), + tma_desc, + ( + Int32(head_dim_stage_offset), + tile_offset, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + else: + # 16-bit dense K/V uses 64-column chunks for the staged + # head-dim slice. + chunk_hd = min(head_dim_stage, 64) + num_chunks = head_dim_stage // chunk_hd + tile_chunk_elems = chunk_hd * cfg.tile_size_kv + if prims.elect_sync(): + stage_base = self._stage_base(stage_info) + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + smem_offset = chunk_idx * tile_chunk_elems + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + tma_desc, + ( + Int32(global_head_dim_offset), + tile_offset, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + + @producer_work + @cute.jit + def load_k0( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the first split K tile for this schedule phase.""" + # ProdWork: K0 uses inst slot 0; the section selects HEAD/LOOP/TAIL + # tile numbering and head_dim_stage_idx selects the H256 slice. + self._producer_load(stage_info, section, head_dim_stage_idx) + + @producer_work + @cute.jit + def load_k1( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the second split K tile for this schedule phase.""" + # ProdWork: K1 uses inst slot 1 but otherwise shares the same staged + # K/V TMA path as K0. + self._producer_load(stage_info, section, head_dim_stage_idx) + + @producer_work + @cute.jit + def load_v0( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the first split V tile for this schedule phase.""" + # ProdWork: V0 publishes the first V descriptor stream consumed by the + # corresponding PV MMA call. + self._producer_load(stage_info, section, head_dim_stage_idx) + + @producer_work + @cute.jit + def load_v1( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the second split V tile for this schedule phase.""" + # ProdWork: V1 publishes the second V descriptor stream consumed by the + # corresponding PV MMA call. + self._producer_load(stage_info, section, head_dim_stage_idx) + + @cute.jit + def _build_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Advance the split K/V base descriptor to the committed stage.""" + stage_offset_bytes = stage_info.stage_idx * Int32(self.cfg.smem_kv_tile_bytes) + return self._desc_base.advance_start_address(stage_offset_bytes) + + @consumer_work(returns=kv_desc_slot) + @cute.jit + def kv_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the K descriptor consumed by QK MMA.""" + # ConsWork: advance the K descriptor to the stage committed by the + # producer and route it through the kv_desc task-local slot. + return self._build_desc(stage_info) + + @consumer_work(returns=v_desc_slot) + @cute.jit + def v_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the V descriptor consumed by PV MMA.""" + # ConsWork: advance the V descriptor to the stage committed by the + # producer and route it through the v_desc task-local slot. + return self._build_desc(stage_info) + + +@dataclass(kw_only=True) +class SmemPageOffsetsKvResource(DecodeGenResourceBase): + """Paged-KV logical-to-physical page IDs staged in SMEM. + + A dedicated warp prefetches the page table entries for the next K/V tile. + The TMA load warp then reads these SMEM-cached offsets when issuing the + page-sized TMA copies, matching the split producer layout. + + Paired K0/K1/V0/V1 schedules publish one stage per logical tile and store + exactly that tile's page IDs. Shared-offset schedules retain a warp-aligned + 32-ID window so one coalesced load can serve adjacent logical tiles. + """ + + cfg: Constexpr[FmhaDecodeConfig] = None + stage_page_ids_per_tile: Constexpr[bool] = False + page_idx_kv: cute.Pointer | None = None + seqlens_kv: cute.Pointer | None = None + use_native_paged_kv: Constexpr[bool] = False + block_tables: cute.Pointer | None = None + block_table_row_stride: cutlass.Int64 = None + max_seq_len_kv: Int32 = None + h_k_idx: Int32 = None + b_idx: Int32 = None + q_group_idx: Int32 = None + seq_len_q: Int32 = None + _alloc: Constexpr[SmemAllocation | None] = None + _smem_page_offsets: cutlass.Array = None + cached_page_ids: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __post_init__(self) -> None: + """Create a shape-stable register slot before task dispatch branches.""" + object.__setattr__( + self, + "cached_page_ids", + TaskLocalVariable( + dtype=cutlass.Array, + default_factory=lambda: cutlass.Array( + Int32, + self.cfg.tile_size_kv // self.cfg.num_tokens_per_page, + space=cutlass.AddressSpace.rmem, + ), + docs="Page IDs reused by every head-dimension stage of one K/V tile.", + ), + ) + self._init_placeholder_state() + + def _init_placeholder_state(self) -> None: + """Create placeholder storage for per-stage page-offset windows.""" + num_stages = ( + self.pipeline_config.num_stages if self.pipeline_config is not None else 1 + ) + pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page + page_ids_per_stage = pages_per_tile if self.stage_page_ids_per_tile else 32 + self._smem_page_offsets = _placeholder_smem_array( + Int32, num_stages * page_ids_per_stage + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate one tile or one held window per page-offset stage.""" + num_stages = ( + self.pipeline_config.num_stages if self.pipeline_config is not None else 1 + ) + pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page + page_ids_per_stage = pages_per_tile if self.stage_page_ids_per_tile else 32 + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=num_stages * page_ids_per_stage * 4, + alignment=16, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Page-offset staging uses SMEM only.""" + return [] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind the page-offset SMEM cache for producer and consumer tasks.""" + if cutlass.const_expr(context is not None and context.smem_base is not None): + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else 1 + ) + pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page + page_ids_per_stage = pages_per_tile if self.stage_page_ids_per_tile else 32 + self._smem_page_offsets = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=cutlass.Int32, + shape=(num_stages * page_ids_per_stage,), + addrspace=3, + ) + return {} + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize producer-side page-offset cache state.""" + # ProdAuxWork: bind the page-offset SMEM cache before the prefetch warp + # starts publishing page-table windows. + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_read_state(self, stage_info: StageInfo) -> None: + """Initialize consumer-side page-offset cache state.""" + # ConsAuxWork: bind the same cache on the load-warp side so K/V TMA + # work can slice out the page IDs for each tile. + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=cached_page_ids) + @cute.jit + def init_cached_read_state(self, stage_info: StageInfo) -> cutlass.Array: + """Initialize the per-tile page-ID register cache.""" + self._create_initial_task_locals(stage_info.context) + return cutlass.Array( + Int32, + self.cfg.tile_size_kv // self.cfg.num_tokens_per_page, + space=cutlass.AddressSpace.rmem, + ) + + @cute.jit + def page_ids(self, tile_idx: Int32) -> cutlass.Array: + """Load the tile's page IDs from its staged cache entry. + + Single-tile stages begin at offset zero. Multi-tile stages use the + runtime-resolved tile index to select from their aligned 32-ID window. + """ + cfg = self.cfg + pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page + if cutlass.const_expr(self.stage_page_ids_per_tile): + offset = self.consumer_work_stage * Int32(pages_per_tile) + else: + group_page_idx = (tile_idx * Int32(pages_per_tile)) & Int32(31) + offset = self.consumer_work_stage * Int32(32) + group_page_idx + if cutlass.const_expr(pages_per_tile in (8, 16)): + # Native shared-memory vector loads top out at four Int32 values. + # Wide page-16 tiles therefore consume their IDs as independently + # aligned 16-byte loads from the same 32-ID cache window. + page_ids = cutlass.Array( + Int32, pages_per_tile, space=cutlass.AddressSpace.rmem + ) + for vector_idx in cutlass.range_constexpr(pages_per_tile // 4): + vector = self._smem_page_offsets.load( + offset + Int32(vector_idx * 4), + vector_size=4, + alignment=16, + ) + for elem_idx in cutlass.range_constexpr(4): + page_ids[vector_idx * 4 + elem_idx] = vector[elem_idx] + return page_ids + if cutlass.const_expr(pages_per_tile == 4): + return self._smem_page_offsets.load(offset, vector_size=4, alignment=16) + if cutlass.const_expr(pages_per_tile == 2): + return self._smem_page_offsets.load(offset, vector_size=2, alignment=8) + return self._smem_page_offsets.load(offset, vector_size=1, alignment=4) + + @consumer_work(returns=cached_page_ids) + @cute.jit + def cache_page_ids( + self, + stage_info: StageInfo, + *, + cached_page_ids: cutlass.Array, + inst_id: Constexpr[int], + kv_kind: Constexpr[int], + section: Constexpr[FmhaStage], + ) -> cutlass.Array: + """Load one tile's page IDs once for all of its head-dim stages.""" + cfg = self.cfg + local_tile_idx = _local_kv_tile_idx_for_section( + cfg, stage_info, inst_id, kv_kind, section + ) + tile_idx = ( + Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE]) + + local_tile_idx + ) + pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page + + # BF16 TMA is issued only by the elected lane, so only that lane needs + # the register cache. FP8's predicated helper builds coordinates in + # every lane and therefore keeps the existing all-lane semantics. + if cutlass.const_expr(cfg.use_fp8_qkv): + fp8_page_ids = self.page_ids(tile_idx) + for page_frag in cutlass.range_constexpr(pages_per_tile): + cached_page_ids[page_frag] = Int32(fp8_page_ids[page_frag]) + elif prims.elect_sync(): + bf16_page_ids = self.page_ids(tile_idx) + for page_frag in cutlass.range_constexpr(pages_per_tile): + cached_page_ids[page_frag] = Int32(bf16_page_ids[page_frag]) + return cached_page_ids + + @cute.jit + def _producer_load_page_offsets( + self, + stage_info: StageInfo, + inst_id: int, + kv_kind: int, + section: Constexpr[FmhaStage], + ) -> None: + """Prefetch one K/V tile or aligned multi-tile window.""" + cfg = self.cfg + local_tile_idx = _local_kv_tile_idx_for_section( + cfg, stage_info, inst_id, kv_kind, section + ) + + # Resolve the logical tile after split-KV and sliding-window + # transforms so the staged page IDs match the K/V TMA descriptor. + tile_idx = ( + Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE]) + + local_tile_idx + ) + _, logical_b_idx = _logical_head_batch(stage_info, self.h_k_idx, self.b_idx) + pages_per_tile = Int32(cfg.tile_size_kv // cfg.num_tokens_per_page) + if cutlass.const_expr(self.use_native_paged_kv): + task_cache = _decode_gen_task_cache(stage_info) + page_idx_ub = Int32(task_cache[_TASK_CACHE_KV_PAGE_IDX_UB]) + page_table_offset = ( + cutlass.Int64(logical_b_idx) * self.block_table_row_stride + ) + page_idx_kv = self.block_tables + else: + if cutlass.const_expr(self.seqlens_kv is None): + page_idx_ub = Int32(cfg.max_num_pages_per_seq_kv - 1) + else: + # Clamp page prefetches to the last valid page; softmax masking + # removes invalid tokens in the final partial tile. + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + self.h_k_idx, + self.b_idx, + ) + page_idx_ub = _runtime_last_valid_page_idx(cfg, seq_len_kv) + + page_table_offset = logical_b_idx * Int32(2 * cfg.max_num_pages_per_seq_kv) + if cutlass.const_expr(kv_kind == KV_KIND_V): + # K and V page tables are stored as two consecutive per-batch ranges. + page_table_offset += Int32(cfg.max_num_pages_per_seq_kv) + page_idx_kv = self.page_idx_kv + smem_page_offsets = self._smem_page_offsets + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + if cutlass.const_expr(self.stage_page_ids_per_tile): + if lane_idx < pages_per_tile: + logical_page_idx = cute.math.min( + tile_idx * pages_per_tile + lane_idx, page_idx_ub + ) + smem_offset = stage_info.stage_idx * pages_per_tile + lane_idx + smem_page_offsets[smem_offset] = Int32( + page_idx_kv[page_table_offset + logical_page_idx] + ) + else: + # Shared-offset schedules use one coalesced warp load for an + # aligned 32-ID window; consumers select their tile within it. + grouped_base_page_idx = ((tile_idx * pages_per_tile) >> Int32(5)) << Int32( + 5 + ) + grouped_logical_page_idx = cute.math.min( + grouped_base_page_idx + lane_idx, page_idx_ub + ) + grouped_smem_offset = stage_info.stage_idx * Int32(32) + lane_idx + smem_page_offsets[grouped_smem_offset] = Int32( + page_idx_kv[page_table_offset + grouped_logical_page_idx] + ) + + @producer_work + @cute.jit + def load_k0(self, stage_info: StageInfo, *, section: Constexpr[FmhaStage]) -> None: + """Produce the first K page-offset stage.""" + # ProdWork: prefetch the page IDs that cover K0's tile. + self._producer_load_page_offsets(stage_info, KV_INST0, KV_KIND_K, section) + + @producer_work + @cute.jit + def load_k1(self, stage_info: StageInfo, *, section: Constexpr[FmhaStage]) -> None: + """Produce the second K page-offset stage.""" + # ProdWork: prefetch the page IDs that cover K1's tile. + self._producer_load_page_offsets(stage_info, KV_INST1, KV_KIND_K, section) + + @producer_work + @cute.jit + def load_v0(self, stage_info: StageInfo, *, section: Constexpr[FmhaStage]) -> None: + """Produce the first V page-offset stage.""" + # ProdWork: prefetch the page IDs that cover V0's tile. + self._producer_load_page_offsets(stage_info, KV_INST0, KV_KIND_V, section) + + @producer_work + @cute.jit + def load_v1(self, stage_info: StageInfo, *, section: Constexpr[FmhaStage]) -> None: + """Produce the second V page-offset stage.""" + # ProdWork: prefetch the page IDs that cover V1's tile. + self._producer_load_page_offsets(stage_info, KV_INST1, KV_KIND_V, section) + + @consumer_work + @cute.jit + def read_offsets(self, stage_info: StageInfo) -> None: + """Consume a generic page-offset window for shared-offset paths.""" + # ConsWork: the page IDs are read directly from SMEM by the K/V load + # resource; this method records the schedule edge only. + _ = stage_info + return + + @consumer_work + @cute.jit + def read_offsets_k0(self, stage_info: StageInfo) -> None: + """Consume the first K page-offset window.""" + # ConsWork: route the K0 offset token to the matching K0 TMA load. + _ = stage_info + return + + @consumer_work + @cute.jit + def read_offsets_k1(self, stage_info: StageInfo) -> None: + """Consume the second K page-offset window.""" + # ConsWork: route the K1 offset token to the matching K1 TMA load. + _ = stage_info + return + + @consumer_work + @cute.jit + def read_offsets_v0(self, stage_info: StageInfo) -> None: + """Consume the first V page-offset window.""" + # ConsWork: route the V0 offset token to the matching V0 TMA load. + _ = stage_info + return + + @consumer_work + @cute.jit + def read_offsets_v1(self, stage_info: StageInfo) -> None: + """Consume the second V page-offset window.""" + # ConsWork: route the V1 offset token to the matching V1 TMA load. + _ = stage_info + return + + +@dataclass(kw_only=True) +class SmemKvResource(DecodeGenResourceBase): + """Shared KV staging resource for the decode kernel. + + K and V loads share one SMEM allocation and one async pipeline/state. + Loads alternate K and V into one ring buffer; the consumer descriptors + target the same allocation. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "kv_desc_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "SMEM descriptor for K loads consumed by QK MMA.", + ), + ( + "v_desc_0_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "First V descriptor consumed by VP MMA.", + ), + ( + "v_desc_1_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "Second V descriptor consumed by VP MMA.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + tma_desc_k: cutlass.Pointer | None = None + tma_desc_v: cutlass.Pointer | None = None + tma_desc_k_atom: cutlass.Pointer | None = None + tma_desc_v_atom: cutlass.Pointer | None = None + sparse_kv_metadata0: "SmemBlockSparseKvMetadataResource | None" = None + sparse_kv_metadata1: "SmemBlockSparseKvMetadataResource | None" = None + page_offsets_kv: SmemPageOffsetsKvResource | None = None + seqlens_kv: cute.Pointer | None = None + max_seq_len_kv: Int32 = None + h_k_idx: Int32 = None + b_idx: Int32 = None + q_group_idx: Int32 = None + seq_len_q: Int32 = None + _alloc: Constexpr[SmemAllocation | None] = None + _smem_base_kv: cutlass.Array = None + _k_desc_base: prims.Tcgen05SmemDesc = None + _v_desc_base: prims.Tcgen05SmemDesc = None + kv_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + v_desc_0_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + v_desc_1_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def _init_placeholder_state(self) -> None: + """Create placeholder state for the shared K/V SMEM ring.""" + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else self.cfg.kv_stages + ) + self._smem_base_kv = _placeholder_smem_array( + self.cfg.kv_dtype, + self.cfg.smem_kv_tile_elements * num_stages, + ) + self._k_desc_base = prims.Tcgen05SmemDesc(0) + self._v_desc_base = prims.Tcgen05SmemDesc(0) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate the shared K/V staged SMEM ring.""" + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else self.cfg.kv_stages + ) + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.smem_kv_tile_bytes * num_stages, + alignment=self.cfg.stensor_align, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Shared K/V staging uses SMEM only.""" + return [] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind the shared K/V ring and build K/V base descriptors.""" + if cutlass.const_expr(context is not None and context.smem_base is not None): + # Bind the shared K/V SMEM ring. K and V descriptors use the same + # allocation but may differ in leading-byte offset. + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else self.cfg.kv_stages + ) + self._smem_base_kv = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.kv_dtype, + shape=(self.cfg.smem_kv_tile_elements * num_stages,), + addrspace=3, + ) + kv_tile_bytes = Int32( + self.cfg.tile_size_kv + * self.cfg.head_dim_kv_stage + * self.cfg.kv_dtype_bytes + ) + # 16-bit K/V stages are made of 64-column chunks. FP8 uses a + # stride derived from the staged head dimension. + k_leading_byte_offset = Int32( + self.cfg.tile_size_kv + * min(self.cfg.head_dim_kv_stage, 64) + * self.cfg.kv_dtype_bytes + ) + stride_byte_offset = Int32(1024) + if cutlass.const_expr(self.cfg.use_fp8_qkv): + k_leading_byte_offset = kv_tile_bytes + stride_byte_offset = Int32( + _major_k_stride_bytes( + self.cfg.kv_dtype_bytes, self.cfg.head_dim_kv_stage + ) + ) + v_leading_byte_offset = k_leading_byte_offset + if cutlass.const_expr(self.cfg.tile_size_kv == 256): + # K spans the complete KV256 row between D64 halves. V is + # staged as four semantic KV64 blocks, each with D/64 adjacent + # D64 halves, so its MMA-K leading step is one KV64 block. + v_leading_byte_offset = Int32(64 * 64 * self.cfg.kv_dtype_bytes) + if cutlass.const_expr(self.cfg.use_fp8_qkv or self.cfg.headdim == 64): + v_leading_byte_offset = Int32(0) + # Descriptor bases are advanced per stage at consumption time; the + # swizzle parameters are invariant for the resource. + self._k_desc_base = prims.Tcgen05SmemDesc.build( + self._smem_base_kv, + leading_byte_offset=k_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_swizzle(self.cfg), + ) + self._v_desc_base = prims.Tcgen05SmemDesc.build( + self._smem_base_kv, + leading_byte_offset=v_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_swizzle(self.cfg), + ) + return { + "kv_desc": cutlass.Int64(0), + "v_desc_0": cutlass.Int64(0), + "v_desc_1": cutlass.Int64(0), + } + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Provide shared K/V descriptor slots for one work tile.""" + _ = context + return { + "kv_desc": cutlass.Int64(0), + "v_desc_0": cutlass.Int64(0), + "v_desc_1": cutlass.Int64(0), + } + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize producer-side shared K/V SMEM state.""" + # ProdAuxWork: bind the shared K/V ring and descriptor bases before + # the load task alternates K and V stages through it. + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize consumer-side shared K/V descriptor state.""" + # ConsAuxWork: initialize descriptor slots for both K and V consumers + # of the shared K/V ring. + self._create_initial_task_locals(stage_info.context) + + @cute.jit + def _stage_base(self, stage_info: StageInfo) -> cutlass.Array: + """Return the SMEM base for the current shared K/V pipeline stage.""" + # Return the base pointer for the producer stage selected by TS. + stage_elems = self.cfg.smem_kv_tile_bytes // self.cfg.kv_dtype_bytes + return self._smem_base_kv.subview(stage_info.stage_idx * stage_elems) + + @cute.jit + def _logical_coords( + self, stage_info: StageInfo, tile_idx: Int32 + ) -> tuple[Int32, Int32, Int32]: + """Return logical head, batch, and token offset for a KV tile.""" + # Translate logical head/batch and tile index into tensor coordinates. + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + tile_offset = tile_idx * Int32(self.cfg.tile_size_kv) + return logical_h_k_idx, logical_b_idx, tile_offset + + @cute.jit + def _maybe_runtime_tile_idx(self, stage_info: StageInfo, tile_idx: Int32) -> Int32: + """Apply runtime sequence and split-KV transforms to a shared K/V tile.""" + if cutlass.const_expr(self.cfg.use_paged_kv): + return ( + Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE]) + + tile_idx + ) + # Apply runtime sequence-length, split-KV, and sliding-window + # transforms to a local prefetch tile. + if cutlass.const_expr( + self.seqlens_kv is None and not self.cfg.uses_runtime_q_kv_union + ): + if cutlass.const_expr(self.cfg.use_split_kv): + tile_idx = _static_split_kv_global_tile_idx( + self.cfg, stage_info, tile_idx + ) + # Clamp out-of-range tile indices so the head/tail cadence stays + # safe for short KV sequences. The softmax mask suppresses the + # duplicated rows downstream. + tile_idx = _clamp_valid_tile_idx(self.cfg, tile_idx) + return tile_idx + Int32(self.cfg.static_num_skipped_kv_tiles) + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + self.h_k_idx, + self.b_idx, + ) + logical_q_group_idx = _logical_q_group_idx( + self.cfg, stage_info, self.q_group_idx + ) + q_token_base = _q_group_token_base(self.cfg, logical_q_group_idx) + tile_idx = _runtime_split_kv_global_tile_idx( + self.cfg, + stage_info, + tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + tile_idx = _runtime_clamp_valid_tile_idx( + self.cfg, + tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + return tile_idx + _num_skipped_kv_tiles( + self.cfg, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + + @cute.jit + def _local_tile_idx( + self, + stage_info: StageInfo, + inst_id: int, + kv_kind: int, + section: Constexpr[FmhaStage], + ) -> Int32: + """Return the local K/V tile index for one shared-ring producer call.""" + return _local_kv_tile_idx_for_section( + self.cfg, stage_info, inst_id, kv_kind, section + ) + + @cute.jit + def _producer_load_kv( + self, + stage_info: StageInfo, + inst_id: int, + kv_kind: int, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + cached_page_ids: cutlass.Array | None = None, + ) -> None: + """Issue one shared-ring K or V TMA load for the schedule phase.""" + cfg = self.cfg + tma_desc = ( + self.tma_desc_v + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self.tma_desc_k + ) + local_tile_idx = self._local_tile_idx(stage_info, inst_id, kv_kind, section) + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + head_dim_stage = cfg.head_dim_kv_stage + head_dim_stage_offset = head_dim_stage_idx * head_dim_stage + + if cutlass.const_expr(cfg.tile_size_kv == 256): + sparse_kv_metadata = ( + self.sparse_kv_metadata0 + if cutlass.const_expr(inst_id == KV_INST0) + else self.sparse_kv_metadata1 + ) + self._producer_load_kv_tile_256( + stage_info, + tma_desc, + local_tile_idx, + logical_h_k_idx, + logical_b_idx, + kv_kind, + cached_page_ids, + sparse_kv_metadata, + ) + elif cutlass.const_expr(cfg.use_paged_kv): + # Paged-KV path: LoadTask consumes pre-staged page IDs and issues + # one TMA per page fragment into the current SMEM stage. Grouped + # cache stages hold 32 page IDs per side; recover the + # runtime-resolved tile_idx so the right per-tile slice is + # selected from the shared window. + page_fragments = cfg.tile_size_kv // cfg.num_tokens_per_page + if cutlass.const_expr(cfg.use_fp8_qkv): + # FP8 pages are contiguous across the staged head dimension. + fp8_stage_base = self._stage_base(stage_info) + if cutlass.const_expr(cached_page_ids is None): + grouped_tile_idx = self._maybe_runtime_tile_idx( + stage_info, local_tile_idx + ) + fp8_page_ids = self.page_offsets_kv.page_ids(grouped_tile_idx) + else: + fp8_page_ids = cached_page_ids + for fp8_page_frag in cutlass.range_constexpr(page_fragments): + fp8_page_id = Int32(fp8_page_ids[fp8_page_frag]) + fp8_smem_page_offset = Int32( + fp8_page_frag * cfg.num_tokens_per_page * head_dim_stage + ) + _cp_async_bulk_tensor_4d_shared_cta_global_predicated( + fp8_stage_base.subview(fp8_smem_page_offset), + tma_desc, + ( + Int32(head_dim_stage_offset), + Int32(0), + logical_h_k_idx, + fp8_page_id, + ), + stage_info.barrier, + ) + else: + # 16-bit pages are split into 64-column chunks inside the + # staged head-dim slice. + chunk_hd = min(head_dim_stage, 64) + num_chunks = head_dim_stage // chunk_hd + tile_chunk_elems = chunk_hd * cfg.tile_size_kv + page_chunk_elems = chunk_hd * cfg.num_tokens_per_page + if cutlass.const_expr(cached_page_ids is None): + # Resolve the tile on every lane before the elected-lane + # branch. The release compiler rejects a local that is + # materialized only on the elected dynamic path. + grouped_tile_idx = self._maybe_runtime_tile_idx( + stage_info, local_tile_idx + ) + if prims.elect_sync(): + stage_base = self._stage_base(stage_info) + if cutlass.const_expr(cached_page_ids is None): + page_ids = self.page_offsets_kv.page_ids(grouped_tile_idx) + else: + page_ids = cached_page_ids + # Consume each cached page ID across every head-dimension + # chunk before advancing. The copies are independent, and + # this order bounds coordinate live ranges in the unrolled + # TMA sequence for every supported page size. + for page_frag in cutlass.range_constexpr(page_fragments): + page_id = Int32(page_ids[page_frag]) + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + local_tile_offset = chunk_idx * tile_chunk_elems + smem_page_offset = Int32( + local_tile_offset + page_frag * page_chunk_elems + ) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_page_offset), + tma_desc, + ( + Int32(global_head_dim_offset), + Int32(0), + logical_h_k_idx, + page_id, + ), + stage_info.barrier, + ) + elif cutlass.const_expr(cfg.use_fp8_qkv): + # Dense FP8 path: one tensor TMA loads the staged K or V tile. + tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) + tile_offset = tile_idx * Int32(cfg.tile_size_kv) + if prims.elect_sync(): + stage_base = self._stage_base(stage_info) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base, + tma_desc, + ( + Int32(head_dim_stage_offset), + tile_offset, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + else: + # Dense 16-bit path: issue 64-column TMA chunks for this staged + # head-dim slice. + chunk_hd = min(head_dim_stage, 64) + num_chunks = head_dim_stage // chunk_hd + tile_chunk_elems = chunk_hd * cfg.tile_size_kv + tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) + tile_offset = tile_idx * Int32(cfg.tile_size_kv) + if prims.elect_sync(): + stage_base = self._stage_base(stage_info) + for chunk_idx in cutlass.range_constexpr(num_chunks): + local_head_dim_offset = chunk_idx * chunk_hd + global_head_dim_offset = ( + head_dim_stage_offset + local_head_dim_offset + ) + smem_offset = chunk_idx * tile_chunk_elems + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + tma_desc, + ( + Int32(global_head_dim_offset), + tile_offset, + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + + @cute.jit + def _producer_load_kv_tile_256( + self, + stage_info: StageInfo, + tma_desc: cutlass.Pointer, + local_tile_idx: Int32, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + kv_kind: Constexpr[int], + cached_page_ids: cutlass.Array | None, + sparse_kv_metadata: "SmemBlockSparseKvMetadataResource | None", + ) -> None: + """Stage one KV256 tile in the physical 2x2-datapath layout. + + The public decode TensorMaps expose KV64 (or one smaller page) + fragments. K places semantic KV64 blocks in physical order + ``(0, 2, 1, 3)`` while V keeps semantic block order with adjacent D64 + halves. Dense and paged profiles derive those fragments from one + contiguous tile; block-sparse profiles consume four prepared KV64 + origins retained by the instruction-local metadata resource. + """ + cfg = self.cfg + grouped_tile_idx = Int32(0) + if cutlass.const_expr(not cfg.use_block_sparse): + grouped_tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) + stage_base = self._stage_base(stage_info) + + if prims.elect_sync(): + # Only the elected TMA issuer needs page IDs. In particular, + # page16/KV256 otherwise makes all 32 load-warp lanes repeat four + # vector loads for the same 16-entry page fragment. + dense_page_ids = cached_page_ids + if cutlass.const_expr( + cfg.use_paged_kv + and not cfg.use_block_sparse + and cached_page_ids is None + ): + assert self.page_offsets_kv is not None + dense_page_ids = self.page_offsets_kv.page_ids(grouped_tile_idx) + for semantic_block in cutlass.range_constexpr(4): + token_coord = Int32(0) + storage_coord = logical_b_idx + if cutlass.const_expr(cfg.use_block_sparse): + assert sparse_kv_metadata is not None + ( + token_coord, + storage_coord, + ) = sparse_kv_metadata.route_tma_coordinate( + Int32(semantic_block), + logical_b_idx, + ) + physical_block = semantic_block + if cutlass.const_expr(kv_kind == KV_KIND_K): + physical_block = KV_TILE_256_K_SLOT_FOR_SEMANTIC_ATOM[ + semantic_block + ] + for dim_half in cutlass.range_constexpr(2): + if cutlass.const_expr(kv_kind == KV_KIND_K): + block_base = ( + dim_half * cfg.tile_size_kv * 64 + physical_block * 64 * 64 + ) + else: + block_base = ( + semantic_block * cfg.headdim * 64 + dim_half * 64 * 64 + ) + + if cutlass.const_expr(cfg.use_block_sparse): + sparse_tma_desc = ( + self.tma_desc_v_atom + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self.tma_desc_k_atom + ) + assert sparse_tma_desc is not None + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(block_base), + sparse_tma_desc, + ( + Int32(dim_half * 64), + token_coord, + logical_h_k_idx, + storage_coord, + ), + stage_info.barrier, + ) + elif cutlass.const_expr(cfg.use_paged_kv): + fragment_tokens = min(cfg.num_tokens_per_page, 64) + fragments_per_block = 64 // fragment_tokens + for fragment in cutlass.range_constexpr(fragments_per_block): + token_in_tile = ( + semantic_block * 64 + fragment * fragment_tokens + ) + logical_page = token_in_tile // cfg.num_tokens_per_page + token_in_page = token_in_tile % cfg.num_tokens_per_page + page_id = Int32(dense_page_ids[logical_page]) + smem_offset = block_base + fragment * fragment_tokens * 64 + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + tma_desc, + ( + Int32(dim_half * 64), + Int32(token_in_page), + logical_h_k_idx, + page_id, + ), + stage_info.barrier, + ) + else: + tile_offset = grouped_tile_idx * Int32(cfg.tile_size_kv) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(block_base), + tma_desc, + ( + Int32(dim_half * 64), + tile_offset + Int32(semantic_block * 64), + logical_h_k_idx, + logical_b_idx, + ), + stage_info.barrier, + ) + + @producer_work + @cute.jit + def load_k0( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the first shared-ring K tile for this schedule phase.""" + # ProdWork: K0 occupies the first K instruction slot in the shared ring. + self._producer_load_kv( + stage_info, KV_INST0, KV_KIND_K, section, head_dim_stage_idx + ) + + @producer_work + @cute.jit + def load_k1( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the second shared-ring K tile for this schedule phase.""" + # ProdWork: K1 occupies the second K instruction slot in the shared ring. + self._producer_load_kv( + stage_info, KV_INST1, KV_KIND_K, section, head_dim_stage_idx + ) + + @producer_work + @cute.jit + def load_v0( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the first shared-ring V tile for this schedule phase.""" + # ProdWork: V0 occupies the first V instruction slot in the shared ring. + self._producer_load_kv( + stage_info, KV_INST0, KV_KIND_V, section, head_dim_stage_idx + ) + + @producer_work + @cute.jit + def load_v1( + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce the second shared-ring V tile for this schedule phase.""" + # ProdWork: V1 occupies the second V instruction slot in the shared ring. + self._producer_load_kv( + stage_info, KV_INST1, KV_KIND_V, section, head_dim_stage_idx + ) + + @producer_work + @cute.jit + def load_k0_cached( + self, + stage_info: StageInfo, + *, + cached_page_ids: cutlass.Array, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce K0 while reusing page IDs across head-dim stages.""" + self._producer_load_kv( + stage_info, + KV_INST0, + KV_KIND_K, + section, + head_dim_stage_idx, + cached_page_ids, + ) + + @producer_work + @cute.jit + def load_k1_cached( + self, + stage_info: StageInfo, + *, + cached_page_ids: cutlass.Array, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce K1 while reusing page IDs across head-dim stages.""" + self._producer_load_kv( + stage_info, + KV_INST1, + KV_KIND_K, + section, + head_dim_stage_idx, + cached_page_ids, + ) + + @producer_work + @cute.jit + def load_v0_cached( + self, + stage_info: StageInfo, + *, + cached_page_ids: cutlass.Array, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce V0 while reusing page IDs across head-dim stages.""" + self._producer_load_kv( + stage_info, + KV_INST0, + KV_KIND_V, + section, + head_dim_stage_idx, + cached_page_ids, + ) + + @producer_work + @cute.jit + def load_v1_cached( + self, + stage_info: StageInfo, + *, + cached_page_ids: cutlass.Array, + section: Constexpr[FmhaStage], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Produce V1 while reusing page IDs across head-dim stages.""" + self._producer_load_kv( + stage_info, + KV_INST1, + KV_KIND_V, + section, + head_dim_stage_idx, + cached_page_ids, + ) + + @cute.jit + def _build_kv_desc( + self, stage_info: StageInfo, kv_kind: int + ) -> prims.Tcgen05SmemDesc: + """Advance the shared K or V descriptor to the committed stage.""" + # Consumers see the same descriptor layout for every stage; only the + # base address advances by the committed SMEM stage index. + stage_offset_bytes = stage_info.stage_idx * Int32(self.cfg.smem_kv_tile_bytes) + return ( + self._v_desc_base.advance_start_address(stage_offset_bytes) + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self._k_desc_base.advance_start_address(stage_offset_bytes) + ) + + @consumer_work(returns=kv_desc_slot) + @cute.jit + def k_desc_0(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the first K descriptor consumed by QK MMA.""" + # ConsWork: expose the committed shared-ring K stage for QK instance 0. + return self._build_kv_desc(stage_info, KV_KIND_K) + + @consumer_work(returns=kv_desc_slot) + @cute.jit + def k_desc_1(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the second K descriptor consumed by QK MMA.""" + # ConsWork: expose the committed shared-ring K stage for QK instance 1. + return self._build_kv_desc(stage_info, KV_KIND_K) + + @consumer_work(returns=v_desc_0_slot) + @cute.jit + def v_desc_0(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the first V descriptor consumed by PV MMA.""" + # ConsWork: expose the committed shared-ring V stage for PV instance 0. + return self._build_kv_desc(stage_info, KV_KIND_V) + + @consumer_work(returns=v_desc_1_slot) + @cute.jit + def v_desc_1(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the second V descriptor consumed by PV MMA.""" + # ConsWork: expose the committed shared-ring V stage for PV instance 1. + return self._build_kv_desc(stage_info, KV_KIND_V) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py new file mode 100644 index 000000000000..c59712aa8f3c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py @@ -0,0 +1,4560 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``TmemCorrResource`` — correction and output resource. + +LOOP stages rescale an in-flight O tile when the running max changes; +TAIL stages combine the two BMM2 instances, normalize by the final +denominator, and either store the final O tile or write partial O for a +split-KV reduction. +""" + +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 + +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + producer_work, +) + +from ..fmha_decode_config import FmhaDecodeConfig +from ...placeholder_helpers import _placeholder_smem_array +from .helpers_common import ( + Constexpr, + DecodeGenResourceBase, + ResourceVars, + fadd2, + ffma2, + fmul2, + _TASK_CACHE_LANE_IDX, + _TASK_CACHE_TMEM_BASE_OFFSET, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _TASK_CACHE_WARP_IDX, + _decode_gen_task_cache, + _keeps_col_base, + _keeps_row_idx, + _keeps_tcgen05_ld, + _keeps_tcgen05_st, + _attention_sink_head_stride, + _local_head_from_q_output_row, + _logical_head_batch, + _logical_q_group_idx, + _q_group_token_base, + _q_physical_output_row, + _q_physical_output_row_from_logical, + _q_row_is_valid_for_seq, + _q_tile_output_row_base, + _q_tile_valid_rows_for_seq, + _neg_max_f32, + _pack_float2_to_bf16, + _pack_float2_to_fp16, +) +from .helpers_kv_tile_idx import ( + _load_runtime_seq_len_kv, + _logical_cta_kv_idx, + _runtime_execution_splits_kv, +) +from .helpers_output import ( + _copy_transposed_smem8b_to_gmem, + _fp16_o_reorg_offsets, + _load_partial_o_vec8_as_f32, + _store_transposed_smem8b, +) +from .helpers_softmax import ( + _attention_sink_for_local_head, + _attention_sink_for_scale_idx, + _pack_float4_to_fp8_e4m3, +) +from .smem_p import SmemPResource +from .tmem_softmax_stats import TmemSoftmaxLocalResource + +_KV_TILE_256_CORRECTION_THREADS = 128 +_KV_TILE_256_LOGICAL_OUTPUT_ROWS = 64 +_KV_TILE_256_EXCHANGE_ROW_STRIDE = 132 +_KV_TILE_256_STATS_PER_THREAD = 4 + + +@dataclass(kw_only=True) +class TmemCorrResource(DecodeGenResourceBase): + """Correction and output resource. + + Loop stages rescale an in-flight O tile when the running softmax max + changes. Tail stages combine the two BMM2 instances, normalize by the final + denominator, and either store the final O tile or write partial results for + a split-KV reduction. + """ + + inst_id: Constexpr[int] = 0 + cfg: Constexpr[FmhaDecodeConfig] = None + softmax_local0_ref: Constexpr[TmemSoftmaxLocalResource] = None + softmax_local1_ref: Constexpr[TmemSoftmaxLocalResource] = None + scale_softmax_log2: Float32 = None + output_scale: Float32 = None + o_ptr: cute.Pointer = None + partial_o_ptr: cute.Pointer = None + partial_stats_ptr: cute.Pointer = None + split_kv_counter_ptr: cute.Pointer = None + attention_sinks_ptr: cute.Pointer = None + seqlens_kv: cute.Pointer = None + max_seq_len_kv: Constexpr[int] = 0 + seq_len_q: Int32 = None + q_token_offset: Int32 = None + num_heads_kv: Int32 = None + h_r: Int32 = None + h_k_idx: Int32 = None + b_idx: Int32 = None + q_group_idx: Int32 = None + active_splits_kv: Int32 = None + static_full_split_prefix: Constexpr[bool] = False + smem_p0_ref: Constexpr[SmemPResource] = None + smem_p1_ref: Constexpr[SmemPResource] = None + tmem_o_ref: object = None + store_barrier_id: Constexpr[int] = 6 + sum_barrier_id: Constexpr[int] = 7 + _alloc: Constexpr[SmemAllocation | None] = None + _sum_alloc: Constexpr[SmemAllocation | None] = None + _gmem_reducer_rank_alloc: Constexpr[SmemAllocation | None] = None + _smem_base_o_i32: cutlass.Array = None + _sum_scratch: cutlass.Array = None + _gmem_reducer_rank: cutlass.Array = None + _cluster_partial_o_alloc: Constexpr[SmemAllocation | None] = None + _cluster_partial_stats_alloc: Constexpr[SmemAllocation | None] = None + _cluster_mbarrier_alloc: Constexpr[SmemAllocation | None] = None + _cluster_partial_o_i32: cutlass.Array = None + _cluster_partial_stats: cutlass.Array = None + _cluster_mbarrier: cutlass.Array = None + _kv_tile_256_exchange_alloc: Constexpr[SmemAllocation | None] = None + _kv_tile_256_exchange: cutlass.Array = None + + def get_o_stage_dtype_bytes(self) -> int: + """Return the element width used by the final O staging buffer.""" + + return ( + 2 + if self.cfg.use_split_kv and self.cfg.use_fp8_output + else self.cfg.o_dtype_bytes + ) + + def _kv_tile_256_exchange_entries(self) -> int: + """Return 128 lane-local stats plus 64 logical output rows.""" + return ( + _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD + + _KV_TILE_256_LOGICAL_OUTPUT_ROWS * _KV_TILE_256_EXCHANGE_ROW_STRIDE + ) + + def _init_placeholder_state(self) -> None: + """Create placeholder SMEM views for correction and split reduction.""" + o_stage_dtype_bytes = self.get_o_stage_dtype_bytes() + o_entries = max( + self.cfg.tile_size_q * self.cfg.headdim * o_stage_dtype_bytes // 4, + 1, + ) + # Keeps reduces each row's denominator in registers (plus the q64 + # xor-16 lane pair) and never enters either Swaps scratch reducer. + # Retain a one-element placeholder for construction-time tracing, but + # do not reserve physical SMEM for it. + sum_entries = max(self.cfg.correction_sum_scratch_entries, 1) + cluster_o_entries = max( + self.cfg.cluster_max_runtime_partial_rows * self.cfg.headdim * 2 // 4, + 1, + ) + cluster_stats_entries = max( + self.cfg.cluster_max_runtime_partial_rows * 2, + 1, + ) + self._smem_base_o_i32 = _placeholder_smem_array(Int32, o_entries) + self._sum_scratch = _placeholder_smem_array(Float32, sum_entries) + self._gmem_reducer_rank = _placeholder_smem_array(Int32, 1) + self._cluster_partial_o_i32 = _placeholder_smem_array(Int32, cluster_o_entries) + self._cluster_partial_stats = _placeholder_smem_array( + Float32, cluster_stats_entries + ) + self._cluster_mbarrier = _placeholder_smem_array(cutlass.Int64, 1) + self._kv_tile_256_exchange = _placeholder_smem_array( + Float32, + self._kv_tile_256_exchange_entries(), + ) + + def _owns_final_epilogue(self) -> bool: + """Whether this correction instance owns final normalization/output.""" + return (self.cfg.num_insts_kv == 1 and self.inst_id == 0) or ( + self.cfg.num_insts_kv != 1 and self.inst_id == 1 + ) + + @cute.jit + def _runtime_splits_kv(self, stage_info: StageInfo) -> Int32: + """Return the useful producer/reducer fanout for this logical tile.""" + cfg = self.cfg + if cutlass.const_expr(self.static_full_split_prefix): + return Int32(cfg.splits_kv) + # Nonpersistent split-KV launches already derived this cluster-uniform + # prefix at kernel entry to prune the physical CTA suffix. Reuse that + # value in correction instead of loading seq_lens and repeating the + # integer partition in every correction lane. + if cutlass.const_expr(self.active_splits_kv is not None): + return self.active_splits_kv + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + self.h_k_idx, + self.b_idx, + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + return _runtime_execution_splits_kv( + cfg, + seq_len_kv, + self.seq_len_q, + _q_group_token_base(cfg, logical_q_group_idx), + ) + + def _can_alias_o_smem(self) -> bool: + """Whether final O staging can reuse the two contiguous P buffers.""" + # Persistent CTAs can start the next work tile while correction is still + # draining the final O staging path. Keep O staging separate from SmemP + # there so the next tile's P producer cannot overwrite the copy source. + o_stage_dtype_bytes = self.get_o_stage_dtype_bytes() + required_o_bytes = self.cfg.tile_size_q * self.cfg.headdim * o_stage_dtype_bytes + available_p_bytes = 2 * self.cfg.smem_p_tile_bytes + return ( + self.smem_p0_ref is not None + and self.smem_p1_ref is not None + and not self.cfg.use_persistent_scheduler + and available_p_bytes >= required_o_bytes + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate correction SMEM scratch, staging, and cluster partial buffers.""" + if not self._owns_final_epilogue(): + return [] + o_stage_dtype_bytes = self.get_o_stage_dtype_bytes() + needs_o_staging = not self.cfg.use_keeps_mma_ab and not self._can_alias_o_smem() + if needs_o_staging and self._alloc is None: + # O staging is only needed for the final correction/output path and + # lives on inst1, which owns the final two-instance reduction. + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.tile_size_q + * self.cfg.headdim + * o_stage_dtype_bytes, + alignment=self.cfg.stensor_align, + ) + sum_entries = self.cfg.correction_sum_scratch_entries + if self._sum_alloc is None and sum_entries > 0: + self._sum_alloc = SmemAllocation( + name=f"{self.name}_sumScratch", + size_bytes=sum_entries * 4, + alignment=16, + ) + if ( + self.cfg.use_split_kv + and not self.cfg.use_separate_reduction_kernel + and not self.cfg.supports_cluster_smem_reduction + and self._gmem_reducer_rank_alloc is None + ): + self._gmem_reducer_rank_alloc = SmemAllocation( + name=f"{self.name}_gmemReducerRank", + size_bytes=4, + alignment=4, + ) + if self.cfg.supports_cluster_smem_reduction: + # Owner-CTA distributed-SMEM staging for the cluster reducer: every + # split's partial O (fp16) and stats (float2) for the rows this CTA + # owns. Peers write into these via prims.mapa; the owner reads them + # locally instead of round-tripping the partials through GMEM. + max_partial_rows = self.cfg.cluster_max_runtime_partial_rows + if self._cluster_partial_o_alloc is None: + self._cluster_partial_o_alloc = SmemAllocation( + name=f"{self.name}_clusterPartialO", + size_bytes=max_partial_rows * self.cfg.headdim * 2, + alignment=16, + ) + if self._cluster_partial_stats_alloc is None: + self._cluster_partial_stats_alloc = SmemAllocation( + name=f"{self.name}_clusterPartialStats", + size_bytes=max_partial_rows * 2 * 4, + alignment=16, + ) + if self._cluster_mbarrier_alloc is None: + # One transaction mbarrier per owner CTA; peers async-store into + # this CTA's partial buffers and signal it. + self._cluster_mbarrier_alloc = SmemAllocation( + name=f"{self.name}_clusterTransactionBarrier", + size_bytes=8, + alignment=8, + ) + if self.cfg.tile_size_kv == 256 and self._kv_tile_256_exchange_alloc is None: + # Tail correction exchanges all lane-local stats, then pipelines + # D32 fragments through 64 logical output rows. Upper lanes publish + # one spatial half while lower lanes retain the matching fragment + # in registers. The dependency graph places this scratch after the + # shared KV ring so it can reuse the dead storage. + payload_bytes = self._kv_tile_256_exchange_entries() * 4 + exchange_bytes = payload_bytes + if self.cfg.uses_rotating_kv256_exchange: + assert payload_bytes <= self.cfg.smem_kv_tile_bytes + # Runtime selects one compact payload inside this explicit + # full-ring alias envelope. The envelope keeps every dynamic + # pointer within a declared allocation while the actual live + # exchange remains only 35,840 B in one 64-KiB stage. + exchange_bytes = self.cfg.smem_kv_tile_bytes * self.cfg.kv_stages + self._kv_tile_256_exchange_alloc = SmemAllocation( + name=f"{self.name}_kvTile256Exchange", + size_bytes=exchange_bytes, + alignment=16, + ) + allocs = [] + if self._alloc is not None: + allocs.append(self._alloc) + if self._sum_alloc is not None: + allocs.append(self._sum_alloc) + if self._gmem_reducer_rank_alloc is not None: + allocs.append(self._gmem_reducer_rank_alloc) + if self._cluster_partial_o_alloc is not None: + allocs.append(self._cluster_partial_o_alloc) + if self._cluster_partial_stats_alloc is not None: + allocs.append(self._cluster_partial_stats_alloc) + if self._cluster_mbarrier_alloc is not None: + allocs.append(self._cluster_mbarrier_alloc) + if self._kv_tile_256_exchange_alloc is not None: + allocs.append(self._kv_tile_256_exchange_alloc) + return allocs + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Correction consumes TMEM O but does not allocate new TMEM.""" + return [] + + @cute.jit + def _split_reduction_rows_per_cta( + self, total_rows: Int32, splits_kv: Int32 + ) -> Int32: + """Return the slice-aligned row capacity owned by each split reducer.""" + rows_per_slice = Int32(self.cfg.split_reduction_rows_per_slice) + num_slices = (total_rows + rows_per_slice - Int32(1)) // rows_per_slice + slices_per_cta = (num_slices + splits_kv - Int32(1)) // splits_kv + return slices_per_cta * rows_per_slice + + @cute.jit + def _split_reduction_row_is_owned( + self, cta_idx_kv: Int32, row_idx: Int32, total_rows: Int32, splits_kv: Int32 + ) -> bool: + """Test whether this split CTA owns one reducer row.""" + rows_per_cta = self._split_reduction_rows_per_cta(total_rows, splits_kv) + row_start = cta_idx_kv * rows_per_cta + row_end = cute.math.min(row_start + rows_per_cta, total_rows) + return row_idx >= row_start and row_idx < row_end + + @cute.jit + def _multi_cta_counter_q_groups(self, h_r: Int32) -> Int32: + """Return the number of split-KV counter groups in the output tile.""" + cfg = self.cfg + if cutlass.const_expr(cfg.groups_tokens_heads_q): + return Int32( + (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta + ) + if cutlass.const_expr( + cfg.max_seq_len_q > 1 + and not cfg.groups_tokens_heads_q + and cfg.heads_q_per_kv != 0 + ): + return Int32( + ((cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q) + * cfg.max_seq_len_q + ) + return (h_r + Int32(cfg.tile_size_q - 1)) // Int32(cfg.tile_size_q) + + @cute.jit + def _gmem_partial_row_offset( + self, + logical_kv_idx: Int32, + cta_idx_kv: Int32, + row_idx: Int32, + ) -> Int64: + """Return a split/row workspace offset using 64-bit arithmetic.""" + + return ( + Int64(logical_kv_idx) * Int64(self.cfg.max_splits_kv) + Int64(cta_idx_kv) + ) * Int64(self.h_r) + Int64(row_idx) + + @cute.jit + def _cluster_reduction_rows_per_cta(self, splits_kv: Int32) -> Int32: + """Return the cluster owner row-band height for active splits.""" + return self._split_reduction_rows_per_cta( + Int32(self.cfg.tile_size_q), splits_kv + ) + + @cute.jit + def _cluster_reduction_cta_for_row(self, row_idx: Int32, splits_kv: Int32) -> Int32: + """Map an output row to its owning cluster split CTA.""" + return row_idx // self._cluster_reduction_rows_per_cta(splits_kv) + + @cute.jit + def _cluster_reduction_local_row(self, row_idx: Int32, splits_kv: Int32) -> Int32: + """Return the row index inside the owning cluster reducer band.""" + rows_per_cta = self._cluster_reduction_rows_per_cta(splits_kv) + return row_idx - (row_idx // rows_per_cta) * rows_per_cta + + @cute.jit + def _cluster_reduction_row_is_owned( + self, cta_idx_kv: Int32, row_idx: Int32, splits_kv: Int32 + ) -> bool: + """Test whether this split CTA owns one cluster reducer row.""" + return self._split_reduction_row_is_owned( + cta_idx_kv, row_idx, Int32(self.cfg.tile_size_q), splits_kv + ) + + @cute.jit + def _cluster_partial_row_idx( + self, split_idx: Int32, local_row_idx: Int32, splits_kv: Int32 + ) -> Int32: + """Linearize a split/local-row pair in cluster partial buffers.""" + return ( + split_idx * self._cluster_reduction_rows_per_cta(splits_kv) + local_row_idx + ) + + @cute.jit + def _cluster_partial_o_i32_offset( + self, row_idx: Int32, col_offset_bytes: Int32 + ) -> Int32: + """Return the int32 offset for a cluster partial-O fragment.""" + return (row_idx * Int32(self.cfg.headdim * 2) + col_offset_bytes) >> Int32(2) + + @cute.jit + def _cluster_partial_stats_offset(self, row_idx: Int32) -> Int32: + """Return the float offset for a cluster partial max/sum pair.""" + _ = self + return row_idx * Int32(2) + + @cute.jit + def _cluster_init_transaction_barrier( + self, mbarrier, transaction_bytes: Constexpr[int] + ) -> None: + """Initialize the owner CTA mbarrier for peer async partial stores.""" + _ = self + tidx, _, _ = cute.arch.thread_idx() + if tidx == Int32(0): + # One lane initializes the owner CTA's transaction mbarrier. The + # expected byte count is the total peer partial-O + stats traffic + # this owner receives before it can reduce local distributed SMEM. + prims.mbarrier_init(mbarrier, 1) + prims.mbarrier_arrive_expect_tx(mbarrier, transaction_bytes) + + @cute.jit + def _cluster_wait_transaction_barrier( + self, + mbarrier, + warp_grp_thread_idx: Int32, + barrier_id: Constexpr[int], + barrier_threads: Constexpr[int], + ) -> None: + """Wait for all peer async cluster partial stores before local reduction.""" + _ = self + # Give every correction lane one nonblocking acquire attempt. When the + # short transaction is already complete this preserves the all-lane + # fast path; otherwise only lane 0 continues polling. The existing + # correction-group barrier publishes lane 0's ready point before any + # reducer lane reads the distributed-SMEM partials. + cluster_transaction_ready = prims.mbarrier_try_wait_parity( + mbarrier, 0, time_limit=0 + ) + if warp_grp_thread_idx == Int32(0): + while not cluster_transaction_ready: + cluster_transaction_ready = prims.mbarrier_try_wait_parity( + mbarrier, 0, time_limit=10_000_000 + ) + prims.barrier_cta_sync(barrier_id, thread_count=barrier_threads) + + @cute.jit + def _cluster_complete_inactive_row_bytes( + self, + mbarrier, + cta_idx_kv: Int32, + splits_kv: Int32, + logical_q_group_idx: Int32, + warp_grp_thread_idx: Int32, + ) -> None: + """Complete transaction bytes suppressed for inactive owner rows. + + The barrier is initialized with the physical owner-band upper bound. + Partial O and stats stores are issued only for valid token/head rows, so + one correction lane reports the remaining bytes before the owner waits. + """ + cfg = self.cfg + rows_per_cta = self._cluster_reduction_rows_per_cta(splits_kv) + valid_tile_rows = _q_tile_valid_rows_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + self.seq_len_q, + ) + owner_row_start = cta_idx_kv * rows_per_cta + valid_owner_rows = cute.math.min( + cute.math.max(valid_tile_rows - owner_row_start, Int32(0)), + rows_per_cta, + ) + bytes_per_row = Int32(cfg.headdim * 2 + 8) + actual_transaction_bytes = splits_kv * valid_owner_rows * bytes_per_row + completion_bytes = ( + Int32(cfg.cluster_transaction_bytes) - actual_transaction_bytes + ) + if warp_grp_thread_idx == Int32(0) and completion_bytes > Int32(0): + prims.mbarrier_complete_tx(mbarrier, completion_bytes) + + @cute.jit + def _cluster_store_async_vec4_i32(self, dst_ptr, vals, mbarrier) -> None: + """Async-store one packed partial-O vector into owner distributed SMEM.""" + _ = self + # Peer CTA writes one 16-byte partial-O vector into owner distributed + # SMEM and charges the bytes to the owner's transaction mbarrier. + # Keep the vectorized publication on the public inline-PTX API so the + # operation remains one 16-byte store instead of four scalar stores. + cute.arch.inline_ptx( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[{$r0}], {{$r1}, {$r2}, {$r3}, {$r4}}, [{$r5}];", + read_only_args=[ + dst_ptr.ir_value(), + vals[0], + vals[1], + vals[2], + vals[3], + mbarrier.ir_value(), + ], + ) + + @cute.jit + def _cluster_store_async_vec2_f32( + self, dst_ptr, val0: Float32, val1: Float32, mbarrier + ) -> None: + """Async-store one partial max/sum pair into owner distributed SMEM.""" + _ = self + # Peer CTA writes one float2 (max, sum) stats record and signals the + # same owner mbarrier used by the matching partial-O vectors. + cute.arch.inline_ptx( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 " + "[{$r0}], {{$r1}, {$r2}}, [{$r3}];", + read_only_args=[ + dst_ptr.ir_value(), + val0, + val1, + mbarrier.ir_value(), + ], + ) + + @cute.jit + def _safe_norm_rcp(self, sum_val: Float32) -> Float32: + """Clamp the softmax denominator before approximate reciprocal.""" + _ = self + return cute.math.rcp( + cute.math.max(sum_val, Float32(1.0e-12), ftz=True), approx=True + ) + + @cute.jit + def _separate_partial_norm_scale(self, sum_val: Float32) -> Float32: + """Return the public output-domain scale for normalized partial O.""" + + # FlashInfer exposes bmm2_scale for every output dtype. Fold it into + # every normalized partial so the shared reducer only merges states. + return self.output_scale * self._safe_norm_rcp(sum_val) + + @cute.jit + def _separate_partial_lse(self, max_val: Float32, sum_val: Float32) -> Float32: + """Convert one local max/sum pair to the shared log2-LSE contract.""" + + lse_val = Float32(-Float32.inf) + if sum_val > Float32(0.0): + lse_val = self.scale_softmax_log2 * max_val + cute.math.log2( + sum_val, fastmath=True + ) + return lse_val + + @cute.jit + def _pack_separate_partial_o_pair(self, val0: Float32, val1: Float32) -> Int32: + """Pack normalized separate partial O in the selected 16-bit type.""" + + if cutlass.const_expr(self.cfg.use_bf16_separate_partial_o): + return _pack_float2_to_bf16(val0, val1) + return _pack_float2_to_fp16(val0, val1) + + @cute.jit + def _swaps_o_stage_base_addr( + self, + tmem_row_base: Int32, + o_base_col: Constexpr[int], + o_stage_idx: Int32, + ) -> Int32: + """Return the TMEM base for one logical Swaps O stage.""" + return ( + tmem_row_base + + Int32(o_base_col) + + o_stage_idx * Int32(self.cfg.tmem_o_stage_cols) + ) + + @cute.jit + def _online_softmax_correction_scale( + self, + old_max: Float32, + new_max: Float32, + ) -> tuple[Float32, cutlass.Boolean]: + """Return the exact max-change scale and whether it is identity.""" + scale_is_identity = old_max == new_max + scale = Float32(1.0) + if not scale_is_identity: + scale = cute.math.exp2( + self.scale_softmax_log2 * (old_max - new_max), + fastmath=True, + ) + return scale, scale_is_identity + + @cute.jit + def _warp_can_skip_o_correction( + self, + lane_scales_are_identity: cutlass.Boolean, + ) -> cutlass.Boolean: + """Return true only when every lane can leave its O fragment unchanged.""" + _ = self + return prims.vote_sync( + cute.arch.FULL_MASK, + lane_scales_are_identity, + prims.VoteSync.ALL, + ) + + @cute.jit + def _swaps_load_o_stage_chunks( + self, + base_addr: Int32, + *, + q_repeats: Constexpr[int], + num_o_chunks: Constexpr[int], + output_f32_regs: Constexpr[int], + ) -> cutlass.Array: + """Load all 64-column chunks from one Swaps O stage into registers.""" + cfg = self.cfg + o_vals = cutlass.Array( + Float32, output_f32_regs, space=cutlass.AddressSpace.rmem + ) + for chunk_idx in cutlass.range_constexpr(num_o_chunks): + loaded = prims.tcgen05_ld( + "16x256b", + prims.make_tmem_ptr( + base_addr + Int32(cfg.swaps_o_chunk_tmem_offset(chunk_idx)), + Float32, + ), + num=q_repeats, + ) + for reg_idx in cutlass.range_constexpr(4 * q_repeats): + o_vals[chunk_idx * 4 * q_repeats + reg_idx] = loaded[reg_idx] + cute.arch.fence_view_async_tmem_load() + return o_vals + + @cute.jit + def _swaps_load_two_o_stage_chunks( + self, + base_addr0: Int32, + base_addr1: Int32, + *, + q_repeats: Constexpr[int], + num_o_chunks: Constexpr[int], + output_f32_regs: Constexpr[int], + ) -> tuple[cutlass.Array, cutlass.Array]: + """Load matching chunks from the two final Swaps O stages.""" + cfg = self.cfg + o0_vals = cutlass.Array( + Float32, output_f32_regs, space=cutlass.AddressSpace.rmem + ) + o1_vals = cutlass.Array( + Float32, output_f32_regs, space=cutlass.AddressSpace.rmem + ) + for chunk_idx in cutlass.range_constexpr(num_o_chunks): + o0_loaded = prims.tcgen05_ld( + "16x256b", + prims.make_tmem_ptr( + base_addr0 + Int32(cfg.swaps_o_chunk_tmem_offset(chunk_idx)), + Float32, + ), + num=q_repeats, + ) + o1_loaded = prims.tcgen05_ld( + "16x256b", + prims.make_tmem_ptr( + base_addr1 + Int32(cfg.swaps_o_chunk_tmem_offset(chunk_idx)), + Float32, + ), + num=q_repeats, + ) + for reg_idx in cutlass.range_constexpr(4 * q_repeats): + o0_vals[chunk_idx * 4 * q_repeats + reg_idx] = o0_loaded[reg_idx] + o1_vals[chunk_idx * 4 * q_repeats + reg_idx] = o1_loaded[reg_idx] + cute.arch.fence_view_async_tmem_load() + return o0_vals, o1_vals + + @cute.jit + def _swaps_store_scaled_o_chunk( + self, + base_addr: Int32, + chunk_idx: Constexpr[int], + chunk_scaled: cutlass.Array, + *, + q_repeats: Constexpr[int], + ) -> None: + """Store one scaled Swaps O chunk back to TMEM.""" + cfg = self.cfg + if cutlass.const_expr(q_repeats == 4): + scaled_vec = cutlass.Vector.from_elements( + ( + chunk_scaled[0], + chunk_scaled[1], + chunk_scaled[2], + chunk_scaled[3], + chunk_scaled[4], + chunk_scaled[5], + chunk_scaled[6], + chunk_scaled[7], + chunk_scaled[8], + chunk_scaled[9], + chunk_scaled[10], + chunk_scaled[11], + chunk_scaled[12], + chunk_scaled[13], + chunk_scaled[14], + chunk_scaled[15], + ), + Float32, + ) + elif cutlass.const_expr(q_repeats == 2): + scaled_vec = cutlass.Vector.from_elements( + ( + chunk_scaled[0], + chunk_scaled[1], + chunk_scaled[2], + chunk_scaled[3], + chunk_scaled[4], + chunk_scaled[5], + chunk_scaled[6], + chunk_scaled[7], + ), + Float32, + ) + else: + scaled_vec = cutlass.Vector.from_elements( + ( + chunk_scaled[0], + chunk_scaled[1], + chunk_scaled[2], + chunk_scaled[3], + ), + Float32, + ) + prims.tcgen05_st( + "16x256b", + prims.make_tmem_ptr( + base_addr + Int32(cfg.swaps_o_chunk_tmem_offset(chunk_idx)), + Float32, + ), + scaled_vec, + ) + + @cute.jit + def _swaps_rescale_o_stage_in_tmem( + self, + base_addr: Int32, + scale_vals: cutlass.Array, + skip_correction: cutlass.Boolean, + *, + q_repeats: Constexpr[int], + num_o_chunks: Constexpr[int], + output_f32_regs: Constexpr[int], + ) -> None: + """Apply online-softmax correction scales to one Swaps O stage.""" + if not skip_correction: + o_vals = self._swaps_load_o_stage_chunks( + base_addr, + q_repeats=q_repeats, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + for chunk_idx in cutlass.range_constexpr(num_o_chunks): + chunk_scaled = cutlass.Array( + Float32, 4 * q_repeats, space=cutlass.AddressSpace.rmem + ) + for reg_pair_idx in cutlass.range_constexpr(2 * q_repeats): + global_pair_idx = chunk_idx * (2 * q_repeats) + reg_pair_idx + scale_base = (reg_pair_idx // 2) * 2 + reg_base = global_pair_idx * 2 + scaled_pair = fmul2( + ( + scale_vals[scale_base], + scale_vals[scale_base + 1], + ), + (o_vals[reg_base], o_vals[reg_base + 1]), + ) + chunk_scaled[reg_pair_idx * 2] = scaled_pair[0] + chunk_scaled[reg_pair_idx * 2 + 1] = scaled_pair[1] + self._swaps_store_scaled_o_chunk( + base_addr, + chunk_idx, + chunk_scaled, + q_repeats=q_repeats, + ) + # Keep the correction task's TMEM ordering point on the no-op path. + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + if not skip_correction: + # FlashInfer's established TMEM path keeps this view fence after + # real stores. Avoid emitting its duplicate wait on a no-op path. + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def _zero_o_vec8(self) -> cutlass.Array: + """Return a zero-initialized 8-element O accumulator fragment.""" + _ = self + output_vals = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) + for elem_idx in cutlass.range_constexpr(8): + output_vals[elem_idx] = Float32(0.0) + return output_vals + + @cute.jit + def _fold_split_o_vec8( + self, + output_vals: cutlass.Array, + sum_val: Float32, + old_max_val: Float32, + max_val: Float32, + local_max: Float32, + local_sum: Float32, + loaded_partial_regs: cutlass.Array, + ) -> tuple[cutlass.Array, Float32, Float32, Float32]: + """Fold one split-KV partial into the online-softmax reducer state.""" + cfg = self.cfg + new_max = cute.math.max(max_val, local_max, ftz=True) + corr_scale0 = cute.math.exp2( + self.scale_softmax_log2 * (old_max_val - new_max), + fastmath=True, + ) + corr_scale1 = cute.math.exp2( + self.scale_softmax_log2 * (local_max - new_max), + fastmath=True, + ) + partial_vals = _load_partial_o_vec8_as_f32( + loaded_partial_regs, + cfg.use_bf16_output and not cfg.use_fp8_output, + ) + sum_val = sum_val * corr_scale0 + local_sum * corr_scale1 + for pair_idx in cutlass.range_constexpr(4): + val_base = pair_idx * 2 + folded_pair = ffma2( + (corr_scale1, corr_scale1), + (partial_vals[val_base], partial_vals[val_base + 1]), + fmul2( + (corr_scale0, corr_scale0), + (output_vals[val_base], output_vals[val_base + 1]), + ), + ) + output_vals[val_base] = folded_pair[0] + output_vals[val_base + 1] = folded_pair[1] + return output_vals, sum_val, new_max, new_max + + @cute.jit + def _store_final_o_vec8( + self, + final_o_dst, + output_vals: cutlass.Array, + norm_scale: Float32, + ) -> None: + """Pack one contiguous 8-element output fragment to the final O dtype.""" + cfg = self.cfg + if cutlass.const_expr(cfg.use_fp8_output): + final_pairs = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) + for pair_idx in cutlass.range_constexpr(4): + val_base = pair_idx * 2 + pair = fmul2( + (norm_scale, norm_scale), + (output_vals[val_base], output_vals[val_base + 1]), + ) + final_pairs[val_base] = pair[0] + final_pairs[val_base + 1] = pair[1] + final_fp8_regs = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + final_fp8_regs[0] = _pack_float4_to_fp8_e4m3( + final_pairs[0], + final_pairs[1], + final_pairs[2], + final_pairs[3], + ) + final_fp8_regs[1] = _pack_float4_to_fp8_e4m3( + final_pairs[4], + final_pairs[5], + final_pairs[6], + final_pairs[7], + ) + final_o_dst.store( + final_fp8_regs.data_ptr().load(count=2, alignment=4), + alignment=8, + ) + else: + final_regs = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(4): + pair = fmul2( + (norm_scale, norm_scale), + ( + output_vals[reg_idx * 2], + output_vals[reg_idx * 2 + 1], + ), + ) + if cutlass.const_expr(cfg.use_bf16_output): + final_regs[reg_idx] = _pack_float2_to_bf16(pair[0], pair[1]) + else: + final_regs[reg_idx] = _pack_float2_to_fp16(pair[0], pair[1]) + final_o_dst.store( + final_regs.data_ptr().load(count=4, alignment=4), + alignment=16, + ) + + @cute.jit + def _store_softmax_normalized_o_vec8( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + reduce_row_idx: Int32, + reduce_col_idx: Int32, + output_vals: cutlass.Array, + sum_val: Float32, + max_val: Float32, + ) -> None: + """Normalize one O fragment, apply attention sink, and store to GMEM.""" + dst_row_base, norm_scale = self._softmax_output_row_state( + logical_h_k_idx, + logical_b_idx, + reduce_row_idx, + sum_val, + max_val, + ) + dst_offset = dst_row_base + reduce_col_idx * Int32(self.cfg.o_dtype_bytes) + final_o_dst = cutlass.inttoptr( + self.o_ptr.toint() + cutlass.Int64(dst_offset), + mem_space=1, + dtype=Int32, + ) + self._store_final_o_vec8(final_o_dst, output_vals, norm_scale) + + @cute.jit + def _softmax_output_row_state( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + reduce_row_idx: Int32, + sum_val: Float32, + max_val: Float32, + ) -> tuple[Int64, Float32]: + """Resolve one logical row's output address and normalization once.""" + cfg = self.cfg + attention_sink_h_r = _attention_sink_head_stride(cfg, self.h_r) + attention_sink_head_idx = _local_head_from_q_output_row( + cfg, self.h_r, reduce_row_idx + ) + sum_val += _attention_sink_for_local_head( + cfg, + self.attention_sinks_ptr, + self.scale_softmax_log2, + max_val, + logical_h_k_idx, + attention_sink_h_r, + self.num_heads_kv, + attention_sink_head_idx, + ) + # ``output_scale`` is the public bmm2_scale. Fold it into the final + # normalization for every output dtype; split partials reach this + # helper only after the cross-CTA reduction has completed. + norm_scale = self.output_scale * self._safe_norm_rcp(sum_val) + if cutlass.const_expr(cfg.use_fp8_qkv): + # Since P is scaled to [0, 448] for Fused GMEM/cluster FP8-Q, + # divide the partial O by 448 before narrowing it to 16 bits, + # and restore after the partials have been reduced in FP32. + norm_scale *= Float32(448.0) + physical_dst_row_idx = _q_physical_output_row_from_logical( + cfg, + self.h_r, + self.num_heads_kv, + logical_b_idx, + logical_h_k_idx, + reduce_row_idx, + self.q_token_offset, + ) + dst_row_base = Int64(physical_dst_row_idx) * Int64( + cfg.headdim * cfg.o_dtype_bytes + ) + return dst_row_base, norm_scale + + @cute.jit + def _warp_reduce_col_group_sum_pair( + self, sum_pair: tuple[Float32, Float32] + ) -> tuple[Float32, Float32]: + """Reduce two scale sums across lanes in one correction column group.""" + _ = self + for shfl_idx in cutlass.range_constexpr(3): + shfl_mask = 16 >> shfl_idx + other_pair = ( + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=sum_pair[0], + offset=shfl_mask, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=sum_pair[1], + offset=shfl_mask, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ) + sum_pair = fadd2(sum_pair, other_pair) + return sum_pair + + def _resolve_o_smem_offset(self) -> int | None: + """Return the aliased SmemP offset used for final O staging, if valid.""" + # Non-persistent kernels can reuse the two contiguous SmemP allocations + # as final O staging. Persistent kernels allocate separate O staging to + # avoid overlap with the next work tile's P producer. + if not self._can_alias_o_smem(): + return None + if self.smem_p0_ref is None or self.smem_p1_ref is None: + return None + p0_alloc = getattr(self.smem_p0_ref, "_alloc", None) + p1_alloc = getattr(self.smem_p1_ref, "_alloc", None) + if p0_alloc is None or p1_alloc is None: + return None + lo_alloc, hi_alloc = (p0_alloc, p1_alloc) + if hi_alloc.offset < lo_alloc.offset: + lo_alloc, hi_alloc = hi_alloc, lo_alloc + if hi_alloc.offset != lo_alloc.offset + lo_alloc.size_bytes: + raise ValueError( + "TmemCorrResource expected contiguous SmemP allocations for O aliasing" + ) + return lo_alloc.offset + + @cute.jit + def create_function_variables( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind epilogue SMEM buffers and initialize cluster reduction barriers.""" + alias_offset = self._resolve_o_smem_offset() + o_stage_dtype_bytes = self.get_o_stage_dtype_bytes() + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and (alias_offset is not None or self._alloc is not None) + ): + # SMEM O staging is addressed as int32 because both stmatrix output + # and vectorized GMEM stores move packed 16-byte fragments. + self._smem_base_o_i32 = cutlass.Array( + context.smem_base.data_ptr() + + (alias_offset if alias_offset is not None else self._alloc.offset), + dtype=Int32, + shape=( + self.cfg.tile_size_q * self.cfg.headdim * o_stage_dtype_bytes // 4, + ), + addrspace=3, + ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self._sum_alloc is not None + ): + # Sum scratch combines the four correction warps' denominator + # partials before normalization. + self._sum_scratch = cutlass.Array( + context.smem_base.data_ptr() + self._sum_alloc.offset, + dtype=Float32, + shape=(self._sum_alloc.size_bytes // 4,), + addrspace=3, + ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self._gmem_reducer_rank_alloc is not None + ): + self._gmem_reducer_rank = cutlass.Array( + context.smem_base.data_ptr() + self._gmem_reducer_rank_alloc.offset, + dtype=Int32, + shape=(1,), + addrspace=3, + ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self._cluster_partial_o_alloc is not None + ): + # cluster distributed-SMEM partial-O staging, addressed as int32 because + # peers deliver packed 16-byte fp16 fragments via prims.mapa. + self._cluster_partial_o_i32 = cutlass.Array( + context.smem_base.data_ptr() + self._cluster_partial_o_alloc.offset, + dtype=Int32, + shape=(self._cluster_partial_o_alloc.size_bytes // 4,), + addrspace=3, + ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self._cluster_partial_stats_alloc is not None + ): + # cluster distributed-SMEM stats staging (float2 max/sum per owned row + # and split). + self._cluster_partial_stats = cutlass.Array( + context.smem_base.data_ptr() + self._cluster_partial_stats_alloc.offset, + dtype=Float32, + shape=(self._cluster_partial_stats_alloc.size_bytes // 4,), + addrspace=3, + ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self._cluster_mbarrier_alloc is not None + ): + # Per-owner transaction mbarrier, initialised once to expect the full + # cross-split partial byte count before any peer async-stores into it + # (the cluster_arrive/cluster_wait at kernel start makes it visible). + self._cluster_mbarrier = cutlass.Array( + context.smem_base.data_ptr() + self._cluster_mbarrier_alloc.offset, + dtype=cutlass.Int64, + shape=(1,), + addrspace=3, + ) + self._cluster_init_transaction_barrier( + self._cluster_mbarrier.data_ptr(), + self.cfg.cluster_transaction_bytes, + ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self._kv_tile_256_exchange_alloc is not None + ): + self._kv_tile_256_exchange = cutlass.Array( + context.smem_base.data_ptr() + self._kv_tile_256_exchange_alloc.offset, + dtype=Float32, + shape=(self._kv_tile_256_exchange_entries(),), + addrspace=3, + ) + return {} + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_epilogue_state(self, stage_info: StageInfo) -> None: + """Preserve the correction init schedule slot after eager SMEM binding.""" + # ProdAuxWork: function variables are materialized before TaskManager.run() + # so cluster mbarriers are visible before peer async stores. Keep this as + # a captured-schedule placeholder for existing task structure. + return + + @cute.jit + def _sync_gmem_split_reducers( + self, + counter_group_idx: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + ) -> Int32: + """Return the arrival-ranked owner index for GMEM split-KV reduction.""" + cfg = self.cfg + # All correction lanes must finish publishing this CTA's partial O/stats + # before lane 0 participates in the global completion counter. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + counter_ptr = cutlass.inttoptr( + self.split_kv_counter_ptr.toint() + + cutlass.Int64(counter_group_idx * Int32(4)), + mem_space=1, + dtype=Uint32, + ) + if warp_grp_thread_idx == Int32(0): + inc_limit = prims.mov_b32(splits_kv - Int32(1), target_type=Uint32) + # Global wrapping INC counts split-CTA completion. Map arrivals in + # reverse order so the last publisher becomes owner 0 instead of + # handing reduction back to a fixed logical CTA. Non-owners can + # leave immediately. If multiple row-slice owners are needed, only + # the earlier owner arrivals wait for the last arrival to publish + # every split and wrap the counter to zero. + old_complete_u32 = prims.atomicrmw( + prims.AtomicOp.INC, + counter_ptr, + inc_limit, + mem_order=prims.MemOrder.ACQ_REL, + syncscope=prims.MemScope.GPU, + ) + old_complete = prims.mov_b32(old_complete_u32, target_type=Int32) + reducer_cta_idx = splits_kv - Int32(1) - old_complete + if reducer_cta_idx < Int32( + cfg.cluster_reduction_num_owner_ctas + ) and old_complete < splits_kv - Int32(1): + while prims.load_ext( + counter_ptr, + order=prims.MemOrder.ACQUIRE, + scope=prims.MemScope.GPU, + ) != Uint32(0): + pass + self._gmem_reducer_rank.store(reducer_cta_idx, 0, alignment=4) + # Broadcast the arrival-ranked owner and, for owner CTAs, make the + # all-arrived state visible before reading split partials. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + return self._gmem_reducer_rank.load(0, alignment=4) + + @cute.jit + def _reduce_and_store_gmem_split_o_vec8( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + splits_kv: Int32, + reduce_row_idx: Int32, + reduce_col_idx: Int32, + valid_reduce_row, + *, + full_prefix: Constexpr[bool], + ) -> None: + """Reduce one GMEM split-KV O fragment and write the final output.""" + cfg = self.cfg + reduction_splits_kv = splits_kv + if cutlass.const_expr(full_prefix): + reduction_splits_kv = Int32(cfg.splits_kv) + output_vals = self._zero_o_vec8() + sum_val = Float32(0.0) + old_max_val = _neg_max_f32() + max_val = _neg_max_f32() + partial_max = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + partial_sum = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + partial_regs = cutlass.Array(Int32, 16, space=cutlass.AddressSpace.rmem) + for split_base_i in cutlass.range_constexpr(0, cfg.max_splits_kv, 4): + split_base = Int32(split_base_i) + # Load up to four split partials first. Keeping the stats and O + # loads grouped lets the fold loop below operate on registers only. + for jj in cutlass.range_constexpr(4): + if cutlass.const_expr(full_prefix): + split_idx = Int32(split_base_i + jj) + valid_split_idx = cutlass.const_expr( + split_base_i + jj < cfg.splits_kv + ) + else: + split_idx = split_base + Int32(jj) + valid_split_idx = split_idx < reduction_splits_kv + if cutlass.const_expr(cfg.max_splits_kv % 4 != 0): + split_idx = cute.math.min( + split_idx, reduction_splits_kv - Int32(1) + ) + if valid_split_idx and valid_reduce_row: + workspace_row = self._gmem_partial_row_offset( + logical_kv_idx, + split_idx, + reduce_row_idx, + ) + stats_offset = workspace_row * Int64(2 * 4) + stats_src = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_offset, + mem_space=1, + dtype=Float32, + ) + stats_pair = stats_src.load(count=2, alignment=8) + partial_max[jj] = stats_pair[0] + partial_sum[jj] = stats_pair[1] + partial_o_offset = workspace_row * Int64(cfg.headdim * 2) + Int64( + reduce_col_idx + ) * Int64(2) + partial_o_src = cutlass.inttoptr( + self.partial_o_ptr.toint() + partial_o_offset, + mem_space=1, + dtype=Int32, + ) + loaded_partial_regs = partial_o_src.load(count=4, alignment=16) + for reg_idx in cutlass.range_constexpr(4): + partial_regs[jj * 4 + reg_idx] = loaded_partial_regs[reg_idx] + + # Fold the loaded partials with the online log-sum-exp recurrence: + # update max/sum and rescale the accumulated O vector each time a + # split contributes a larger max. + for jj in cutlass.range_constexpr(4): + valid_split_for_apply = split_base + Int32(jj) < reduction_splits_kv + if cutlass.const_expr(full_prefix): + valid_split_for_apply = cutlass.const_expr( + split_base_i + jj < cfg.splits_kv + ) + if valid_split_for_apply and valid_reduce_row: + regs_base = jj * 4 + loaded_partial_regs = cutlass.Vector.from_elements( + ( + partial_regs[regs_base], + partial_regs[regs_base + 1], + partial_regs[regs_base + 2], + partial_regs[regs_base + 3], + ), + Int32, + ) + output_vals, sum_val, old_max_val, max_val = ( + self._fold_split_o_vec8( + output_vals, + sum_val, + old_max_val, + max_val, + partial_max[jj], + partial_sum[jj], + loaded_partial_regs, + ) + ) + + if valid_reduce_row: + self._store_softmax_normalized_o_vec8( + logical_h_k_idx, + logical_b_idx, + reduce_row_idx, + reduce_col_idx, + output_vals, + sum_val, + max_val, + ) + + @cute.jit + def _reduce_fused_gmem_partial_segment( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + q_row_offset: Int32, + logical_q_group_idx: Int32, + reducer_cta_idx: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + rows_per_reducer: Int32, + segment_idx: Int32, + *, + full_prefix: Constexpr[bool], + ) -> None: + """Reduce one correction-warpgroup slice from fused GMEM partials.""" + cfg = self.cfg + partial_row_bytes = Int32(cfg.headdim * 2) + reduce_base_offset = warp_grp_thread_idx * Int32(16) + Int32( + segment_idx * cfg.split_reduction_slice_bytes + ) + reduce_tile_row_idx = reduce_base_offset // partial_row_bytes + reduce_local_row_idx = reducer_cta_idx * rows_per_reducer + reduce_tile_row_idx + reduce_output_row_idx = q_row_offset + reduce_local_row_idx + reduce_col_idx = (reduce_base_offset % partial_row_bytes) >> Int32(1) + valid_reduce_row = ( + reduce_tile_row_idx < rows_per_reducer + and _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + reduce_local_row_idx, + self.seq_len_q, + ) + ) + self._reduce_and_store_gmem_split_o_vec8( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + splits_kv, + reduce_output_row_idx, + reduce_col_idx, + valid_reduce_row, + full_prefix=full_prefix, + ) + + @cute.jit + def _reduce_fused_gmem_partials_impl( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + q_row_offset: Int32, + logical_q_group_idx: Int32, + counter_group_idx: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + *, + full_prefix: Constexpr[bool], + ) -> None: + """Elect fused-GMEM owners and run static or contracted geometry.""" + cfg = self.cfg + reduction_splits_kv = splits_kv + if cutlass.const_expr(full_prefix): + reduction_splits_kv = Int32(cfg.splits_kv) + reducer_cta_idx = self._sync_gmem_split_reducers( + counter_group_idx, + reduction_splits_kv, + warp_grp_thread_idx, + ) + + if reducer_cta_idx < Int32(cfg.cluster_reduction_num_owner_ctas): + if cutlass.const_expr(full_prefix): + rows_per_reducer = Int32(cfg.cluster_reduction_rows_per_cta) + for segment_idx in cutlass.range_constexpr( + cfg.split_reduction_slices_per_cta + ): + self._reduce_fused_gmem_partial_segment( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + reducer_cta_idx, + reduction_splits_kv, + warp_grp_thread_idx, + rows_per_reducer, + Int32(segment_idx), + full_prefix=True, + ) + else: + rows_per_reducer = self._split_reduction_rows_per_cta( + Int32(cfg.tile_size_q), reduction_splits_kv + ) + runtime_segments = rows_per_reducer // Int32( + cfg.split_reduction_rows_per_slice + ) + for segment_idx in cutlass.range(runtime_segments, unroll=1): + self._reduce_fused_gmem_partial_segment( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + reducer_cta_idx, + reduction_splits_kv, + warp_grp_thread_idx, + rows_per_reducer, + segment_idx, + full_prefix=False, + ) + + @cute.jit + def _reduce_fused_gmem_partials( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + q_row_offset: Int32, + logical_q_group_idx: Int32, + counter_group_idx: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + ) -> None: + """Select the full-prefix or contracted fused-GMEM reducer.""" + cfg = self.cfg + if cutlass.const_expr(self.static_full_split_prefix): + self._reduce_fused_gmem_partials_impl( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + counter_group_idx, + Int32(cfg.splits_kv), + warp_grp_thread_idx, + full_prefix=True, + ) + return + # The producer prefix is CTA-uniform. The full path therefore keeps + # every correction-barrier participant on the same static schedule. + if splits_kv == Int32(cfg.splits_kv): + self._reduce_fused_gmem_partials_impl( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + counter_group_idx, + Int32(cfg.splits_kv), + warp_grp_thread_idx, + full_prefix=True, + ) + else: + self._reduce_fused_gmem_partials_impl( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + counter_group_idx, + splits_kv, + warp_grp_thread_idx, + full_prefix=False, + ) + + @cute.jit + def _reduce_cluster_partial_segment( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + q_row_offset: Int32, + logical_q_group_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + rows_per_cta: Int32, + segment_idx: Int32, + *, + full_prefix: Constexpr[bool], + ) -> None: + """Reduce one correction-warpgroup slice from cluster partial SMEM.""" + cfg = self.cfg + reduction_splits_kv = splits_kv + if cutlass.const_expr(full_prefix): + reduction_splits_kv = Int32(cfg.splits_kv) + row_bytes = Int32(cfg.headdim * 2) + reduce_base_offset = ( + warp_grp_thread_idx * Int32(16) + + cta_idx_kv * rows_per_cta * row_bytes + + Int32(segment_idx * cfg.split_reduction_slice_bytes) + ) + reduce_row_idx = reduce_base_offset // row_bytes + reduce_col_idx = (reduce_base_offset % row_bytes) >> Int32(1) + valid_reduce_row = self._cluster_reduction_row_is_owned( + cta_idx_kv, reduce_row_idx, reduction_splits_kv + ) and _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + reduce_row_idx, + self.seq_len_q, + ) + local_row_idx = self._cluster_reduction_local_row( + reduce_row_idx, reduction_splits_kv + ) + output_vals = self._zero_o_vec8() + sum_val = Float32(0.0) + old_max_val = _neg_max_f32() + max_val = _neg_max_f32() + partial_max = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + partial_sum = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + partial_regs = cutlass.Array(Int32, 16, space=cutlass.AddressSpace.rmem) + for split_base_i in cutlass.range_constexpr(0, cfg.max_splits_kv, 4): + split_base = Int32(split_base_i) + # Load up to four split partials first. The full-prefix + # specialization resolves every slot predicate at compile time; + # contracted prefixes retain the runtime bound and safe clamp. + for jj in cutlass.range_constexpr(4): + if cutlass.const_expr(full_prefix): + split_idx = Int32(split_base_i + jj) + valid_split_idx = cutlass.const_expr( + split_base_i + jj < cfg.splits_kv + ) + else: + split_idx = split_base + Int32(jj) + valid_split_idx = split_idx < splits_kv + if cutlass.const_expr(cfg.max_splits_kv % 4 != 0): + split_idx = cute.math.min(split_idx, splits_kv - Int32(1)) + if valid_split_idx and valid_reduce_row: + partial_row_idx = self._cluster_partial_row_idx( + split_idx, local_row_idx, reduction_splits_kv + ) + stats_pair = ( + self._cluster_partial_stats.subview( + self._cluster_partial_stats_offset(partial_row_idx) + ) + .data_ptr() + .load(count=2, alignment=8) + ) + partial_max[jj] = stats_pair[0] + partial_sum[jj] = stats_pair[1] + loaded_partial_regs = ( + self._cluster_partial_o_i32.subview( + self._cluster_partial_o_i32_offset( + partial_row_idx, + reduce_col_idx << Int32(1), + ) + ) + .data_ptr() + .load(count=4, alignment=16) + ) + for reg_idx in cutlass.range_constexpr(4): + partial_regs[jj * 4 + reg_idx] = loaded_partial_regs[reg_idx] + + # Preserve online-softmax arithmetic order after grouped loads. + for jj in cutlass.range_constexpr(4): + valid_split_for_apply = split_base + Int32(jj) < splits_kv + if cutlass.const_expr(full_prefix): + valid_split_for_apply = cutlass.const_expr( + split_base_i + jj < cfg.splits_kv + ) + if valid_split_for_apply and valid_reduce_row: + regs_base = jj * 4 + loaded_partial_regs = cutlass.Vector.from_elements( + ( + partial_regs[regs_base], + partial_regs[regs_base + 1], + partial_regs[regs_base + 2], + partial_regs[regs_base + 3], + ), + Int32, + ) + output_vals, sum_val, old_max_val, max_val = ( + self._fold_split_o_vec8( + output_vals, + sum_val, + old_max_val, + max_val, + partial_max[jj], + partial_sum[jj], + loaded_partial_regs, + ) + ) + if valid_reduce_row: + self._store_softmax_normalized_o_vec8( + logical_h_k_idx, + logical_b_idx, + q_row_offset + reduce_row_idx, + reduce_col_idx, + output_vals, + sum_val, + max_val, + ) + + @cute.jit + def _reduce_cluster_partials_impl( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + q_row_offset: Int32, + logical_q_group_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + *, + full_prefix: Constexpr[bool], + ) -> None: + """Wait for cluster publication and run static or contracted geometry.""" + cfg = self.cfg + reduction_splits_kv = splits_kv + if cutlass.const_expr(full_prefix): + reduction_splits_kv = Int32(cfg.splits_kv) + # Split CTAs always publish their partials, but only CTAs with a + # physical output-row band participate in the owner-side wait and + # reduction. Ownership is allocated in correction-warpgroup slices, + # so structurally empty ranks never wait on an unused local barrier. + if cta_idx_kv < Int32(cfg.cluster_reduction_num_owner_ctas): + self._cluster_complete_inactive_row_bytes( + self._cluster_mbarrier.data_ptr(), + cta_idx_kv, + reduction_splits_kv, + logical_q_group_idx, + warp_grp_thread_idx, + ) + self._cluster_wait_transaction_barrier( + self._cluster_mbarrier.data_ptr(), + warp_grp_thread_idx, + self.store_barrier_id, + cfg.correction_barrier_threads, + ) + if cutlass.const_expr(full_prefix): + rows_per_cta = Int32(cfg.cluster_reduction_rows_per_cta) + for segment_idx in cutlass.range_constexpr( + cfg.split_reduction_slices_per_cta + ): + self._reduce_cluster_partial_segment( + logical_h_k_idx, + logical_b_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + reduction_splits_kv, + warp_grp_thread_idx, + rows_per_cta, + Int32(segment_idx), + full_prefix=True, + ) + else: + rows_per_cta = self._cluster_reduction_rows_per_cta(reduction_splits_kv) + runtime_segments = rows_per_cta // Int32( + cfg.split_reduction_rows_per_slice + ) + for segment_idx in cutlass.range(runtime_segments, unroll=1): + self._reduce_cluster_partial_segment( + logical_h_k_idx, + logical_b_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + reduction_splits_kv, + warp_grp_thread_idx, + rows_per_cta, + segment_idx, + full_prefix=False, + ) + + @cute.jit + def _stage_fp16_o_regs_to_smem( + self, + cfg: Constexpr[FmhaDecodeConfig], + regs_o: cutlass.Array, + warp_grp_thread_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + *, + output_pair_regs: Constexpr[int], + ) -> None: + """Stage packed FP16/BF16 O fragments through the STSM layout. + + Correction lanes hold O as packed register pairs, but final GMEM stores + need contiguous row segments. This first operation writes the register + fragments into the swizzled SMEM layout before a second operation + reloads contiguous vectors. + """ + if cutlass.const_expr(output_pair_regs < 4): + smem_offset_bytes, _, _, _ = _fp16_o_reorg_offsets( + cfg, warp_grp_thread_idx, warp_idx, lane_idx, 0, 0 + ) + smem_dst = self._smem_base_o_i32.subview( + smem_offset_bytes >> Int32(2) + ).data_ptr() + prims.stmatrix( + smem_dst, + regs_o.data_ptr().load(count=output_pair_regs, alignment=4), + prims.MMALayout.COL, + shape=prims.StoreShape.M8N8, + ) + else: + for stsm_group_idx in cutlass.range_constexpr(output_pair_regs // 4): + smem_offset_bytes, _, _, _ = _fp16_o_reorg_offsets( + cfg, + warp_grp_thread_idx, + warp_idx, + lane_idx, + stsm_group_idx, + 0, + ) + smem_dst = self._smem_base_o_i32.subview( + smem_offset_bytes >> Int32(2) + ).data_ptr() + prims.stmatrix( + smem_dst, + (regs_o.data_ptr() + stsm_group_idx * 4).load( + count=4, + alignment=4, + ), + prims.MMALayout.COL, + shape=prims.StoreShape.M8N8, + ) + cute.arch.fence_view_async_shared() + # The copy-out phase below reloads SMEM written by all correction + # warps, so synchronize the correction warpgroup after the STSM stores. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + @cute.jit + def _copy_staged_fp16_o_segments( + self, + cfg: Constexpr[FmhaDecodeConfig], + logical_kv_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + row_offset: Int32, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_q_group_idx: Int32, + warp_grp_thread_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + *, + num_copy_segments: Constexpr[int], + copy_to_partial: Constexpr[bool], + enable_cluster: Constexpr[bool], + full_prefix: Constexpr[bool], + ) -> None: + """Copy STSM-reorganized FP16/BF16 O fragments to partial or final GMEM. + + Each 2 KiB segment is covered by 128 lanes x 16 bytes. For split-KV the + destination is either GMEM partial storage or owner-CTA distributed SMEM + for cluster reduction; otherwise it is the final output tensor. + """ + publication_splits_kv = splits_kv + if cutlass.const_expr(enable_cluster and full_prefix): + publication_splits_kv = Int32(cfg.splits_kv) + for copy_segment_idx in cutlass.range_constexpr(num_copy_segments): + _, load_smem_offset, dst_row_idx, dst_col_offset = _fp16_o_reorg_offsets( + cfg, + warp_grp_thread_idx, + warp_idx, + lane_idx, + 0, + copy_segment_idx, + ) + if cutlass.const_expr(copy_to_partial): + partial_o_row_idx = row_offset + dst_row_idx + valid_partial_row = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + dst_row_idx, + self.seq_len_q, + ) + if valid_partial_row: + smem_src = self._smem_base_o_i32.subview( + load_smem_offset >> Int32(2) + ).data_ptr() + if cutlass.const_expr(enable_cluster): + # Peers async-store partial O into the owner CTA's SMEM + # and signal its transaction mbarrier. + cluster_owner = self._cluster_reduction_cta_for_row( + dst_row_idx, publication_splits_kv + ) + cluster_local = self._cluster_reduction_local_row( + dst_row_idx, publication_splits_kv + ) + partial_o_dst = prims.mapa( + self._cluster_partial_o_i32.subview( + self._cluster_partial_o_i32_offset( + self._cluster_partial_row_idx( + cta_idx_kv, + cluster_local, + publication_splits_kv, + ), + dst_col_offset, + ) + ).data_ptr(), + cluster_owner, + ) + self._cluster_store_async_vec4_i32( + partial_o_dst, + smem_src.load(count=4, alignment=16), + prims.mapa( + self._cluster_mbarrier.data_ptr(), + cluster_owner, + ), + ) + else: + partial_o_row_base = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + partial_o_row_idx, + ) * Int64(cfg.headdim * 2) + partial_o_dst = cutlass.inttoptr( + self.partial_o_ptr.toint() + + partial_o_row_base + + Int64(dst_col_offset), + mem_space=1, + dtype=Int32, + ) + partial_o_dst.store( + smem_src.load(count=4, alignment=16), + alignment=16, + ) + else: + if _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + dst_row_idx, + self.seq_len_q, + ): + smem_src = self._smem_base_o_i32.subview( + load_smem_offset >> Int32(2) + ).data_ptr() + physical_dst_row_idx = _q_physical_output_row( + cfg, + self.h_r, + self.num_heads_kv, + logical_b_idx, + logical_h_k_idx, + logical_q_group_idx, + dst_row_idx, + self.q_token_offset, + ) + dst_row_base = Int64(physical_dst_row_idx) * Int64( + cfg.headdim * cfg.o_dtype_bytes + ) + dst_ptr = cutlass.inttoptr( + self.o_ptr.toint() + dst_row_base + Int64(dst_col_offset), + mem_space=1, + dtype=Int32, + ) + dst_ptr.store(smem_src.load(count=4, alignment=16), alignment=16) + + @cute.jit + def _copy_staged_fp16_o_to_partial( + self, + cfg: Constexpr[FmhaDecodeConfig], + logical_kv_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + row_offset: Int32, + logical_q_group_idx: Int32, + warp_grp_thread_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + *, + num_copy_segments: Constexpr[int], + enable_cluster: Constexpr[bool], + full_prefix: Constexpr[bool], + ) -> None: + """Publish staged partial O to GMEM, or owner CTA SMEM for cluster reduction.""" + self._copy_staged_fp16_o_segments( + cfg, + logical_kv_idx, + cta_idx_kv, + splits_kv, + row_offset, + Int32(0), + Int32(0), + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + num_copy_segments=num_copy_segments, + copy_to_partial=True, + enable_cluster=enable_cluster, + full_prefix=full_prefix, + ) + + @cute.jit + def _copy_staged_fp16_o_to_output( + self, + cfg: Constexpr[FmhaDecodeConfig], + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_q_group_idx: Int32, + warp_grp_thread_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + *, + num_copy_segments: Constexpr[int], + ) -> None: + """Copy staged FP16/BF16 O fragments to the final output tensor.""" + self._copy_staged_fp16_o_segments( + cfg, + Int32(0), + Int32(0), + Int32(1), + Int32(0), + logical_h_k_idx, + logical_b_idx, + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + num_copy_segments=num_copy_segments, + copy_to_partial=False, + enable_cluster=False, + full_prefix=False, + ) + + @cute.jit + def _stage_and_copy_swaps_partial_o( + self, + cfg: Constexpr[FmhaDecodeConfig], + regs_partial_o: cutlass.Array, + logical_kv_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + row_offset: Int32, + logical_q_group_idx: Int32, + warp_grp_thread_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + *, + output_pair_regs: Constexpr[int], + num_copy_segments: Constexpr[int], + enable_cluster: Constexpr[bool], + full_prefix: Constexpr[bool], + ) -> None: + """Stage Swaps partial O through SMEM and publish it for split-KV.""" + # Operation 1: reformat per-lane O registers into contiguous row + # fragments via SMEM. Operation 2: publish those fragments to the + # selected split-KV staging backend. + self._stage_fp16_o_regs_to_smem( + cfg, + regs_partial_o, + warp_grp_thread_idx, + warp_idx, + lane_idx, + output_pair_regs=output_pair_regs, + ) + self._copy_staged_fp16_o_to_partial( + cfg, + logical_kv_idx, + cta_idx_kv, + splits_kv, + row_offset, + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + num_copy_segments=num_copy_segments, + enable_cluster=enable_cluster, + full_prefix=full_prefix, + ) + + @cute.jit + def _store_split_stats_from_scale_arrays( + self, + cfg: Constexpr[FmhaDecodeConfig], + logical_kv_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + row_offset: Int32, + logical_q_group_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + final_max: cutlass.Array, + reduced_sum: cutlass.Array, + *, + num_scale_groups: Constexpr[int], + enable_cluster: Constexpr[bool], + full_prefix: Constexpr[bool], + ) -> None: + """Publish split-KV state for Swaps tail epilogues. + + Separate-GMEM writes one FP32 log2-LSE value per row. Fused GMEM/cluster + writes a float2(max, sum) record. + """ + publication_splits_kv = splits_kv + if cutlass.const_expr(enable_cluster and full_prefix): + publication_splits_kv = Int32(cfg.splits_kv) + # Only one warp publishes per-row state. The lane mapping keeps each + # record next to the matching partial-O row. + if warp_idx == Int32(0) and lane_idx < Int32(4 * num_scale_groups): + stats_idx = lane_idx >> Int32(2) + quad_thread_idx = lane_idx & Int32(3) + stats_row_base = ( + quad_thread_idx * Int32(2) + + ((stats_idx >> Int32(1)) * Int32(8)) + + (stats_idx & Int32(1)) + ) + stats_row_idx = row_offset + stats_row_base + if _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + stats_row_base, + self.seq_len_q, + ): + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + lse_row = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + stats_row_idx, + ) + lse_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + lse_row * Int64(4), + mem_space=1, + dtype=Float32, + ) + lse_ptr.store( + self._separate_partial_lse( + final_max[stats_idx], reduced_sum[stats_idx] + ), + alignment=4, + ) + elif cutlass.const_expr(enable_cluster): + cluster_owner = self._cluster_reduction_cta_for_row( + stats_row_base, publication_splits_kv + ) + cluster_local = self._cluster_reduction_local_row( + stats_row_base, publication_splits_kv + ) + stats_ptr = prims.mapa( + self._cluster_partial_stats.subview( + self._cluster_partial_stats_offset( + self._cluster_partial_row_idx( + cta_idx_kv, + cluster_local, + publication_splits_kv, + ) + ) + ).data_ptr(), + cluster_owner, + ) + self._cluster_store_async_vec2_f32( + stats_ptr, + final_max[stats_idx], + reduced_sum[stats_idx], + prims.mapa( + self._cluster_mbarrier.data_ptr(), + cluster_owner, + ), + ) + else: + stats_row = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + stats_row_idx, + ) + stats_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_row * Int64(2 * 4), + mem_space=1, + dtype=Float32, + ) + stats_ptr.store( + cutlass.Vector.from_elements( + (final_max[stats_idx], reduced_sum[stats_idx]), + Float32, + ), + alignment=8, + ) + + @cute.jit + def _publish_and_reduce_cluster_swaps_partials_impl( + self, + regs_partial_o: cutlass.Array, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + q_row_offset: Int32, + logical_q_group_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + final_max: cutlass.Array, + reduced_sum: cutlass.Array, + *, + output_pair_regs: Constexpr[int], + num_copy_segments: Constexpr[int], + num_scale_groups: Constexpr[int], + full_prefix: Constexpr[bool], + ) -> None: + """Publish and reduce one static or contracted cluster split prefix.""" + cfg = self.cfg + publication_splits_kv = splits_kv + if cutlass.const_expr(full_prefix): + publication_splits_kv = Int32(cfg.splits_kv) + self._stage_and_copy_swaps_partial_o( + cfg, + regs_partial_o, + logical_kv_idx, + cta_idx_kv, + publication_splits_kv, + q_row_offset, + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + output_pair_regs=output_pair_regs, + num_copy_segments=num_copy_segments, + enable_cluster=True, + full_prefix=full_prefix, + ) + self._store_split_stats_from_scale_arrays( + cfg, + logical_kv_idx, + cta_idx_kv, + publication_splits_kv, + q_row_offset, + logical_q_group_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + num_scale_groups=num_scale_groups, + enable_cluster=True, + full_prefix=full_prefix, + ) + self._reduce_cluster_partials_impl( + logical_h_k_idx, + logical_b_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + publication_splits_kv, + warp_grp_thread_idx, + full_prefix=full_prefix, + ) + + @cute.jit + def _publish_and_reduce_cluster_swaps_partials( + self, + regs_partial_o: cutlass.Array, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + q_row_offset: Int32, + logical_q_group_idx: Int32, + cta_idx_kv: Int32, + splits_kv: Int32, + warp_grp_thread_idx: Int32, + warp_idx: Int32, + lane_idx: Int32, + final_max: cutlass.Array, + reduced_sum: cutlass.Array, + *, + output_pair_regs: Constexpr[int], + num_copy_segments: Constexpr[int], + num_scale_groups: Constexpr[int], + ) -> None: + """Select full-prefix or contracted cluster publication and reduction.""" + cfg = self.cfg + if cutlass.const_expr(self.static_full_split_prefix): + self._publish_and_reduce_cluster_swaps_partials_impl( + regs_partial_o, + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + Int32(cfg.splits_kv), + warp_grp_thread_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + output_pair_regs=output_pair_regs, + num_copy_segments=num_copy_segments, + num_scale_groups=num_scale_groups, + full_prefix=True, + ) + return + # Runtime pruning produces one prefix for the complete physical + # cluster. Every correction lane and every active rank takes the same + # branch before any distributed-SMEM address or mbarrier operation. + if splits_kv == Int32(cfg.splits_kv): + self._publish_and_reduce_cluster_swaps_partials_impl( + regs_partial_o, + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + Int32(cfg.splits_kv), + warp_grp_thread_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + output_pair_regs=output_pair_regs, + num_copy_segments=num_copy_segments, + num_scale_groups=num_scale_groups, + full_prefix=True, + ) + else: + self._publish_and_reduce_cluster_swaps_partials_impl( + regs_partial_o, + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + splits_kv, + warp_grp_thread_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + output_pair_regs=output_pair_regs, + num_copy_segments=num_copy_segments, + num_scale_groups=num_scale_groups, + full_prefix=False, + ) + + @cute.jit + def _kv_tile_256_exchange_for_stage( + self, + stage_info: StageInfo, + scratch_stage: Int32 | None, + ) -> cutlass.Array: + """Return the fixed exchange or its dynamically selected KV stage.""" + if cutlass.const_expr(scratch_stage is None): + return self._kv_tile_256_exchange + return cutlass.Array( + stage_info.context.smem_base.data_ptr() + + self._kv_tile_256_exchange_alloc.offset + + scratch_stage * Int32(self.cfg.smem_kv_tile_bytes), + dtype=Float32, + shape=(self._kv_tile_256_exchange_entries(),), + addrspace=3, + ) + + @cute.jit + def _kv_tile_256_temporal_fragment( + self, + *, + base_addr0: Int32, + base_addr1: Int32, + fragment_col: Constexpr[int], + weight00: Float32, + weight10: Float32, + ): + """Load and combine one D32 fragment from the two temporal stages.""" + cfg = self.cfg + o0_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr0 + Int32(fragment_col), Float32), + num=32, + offset=cfg.headdim // 2, + ) + o1_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr1 + Int32(fragment_col), Float32), + num=32, + offset=cfg.headdim // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + combined: tuple = () + for elem in cutlass.range_constexpr(0, 32, 2): + combined += ffma2( + (weight10, weight10), + (o1_vals[elem], o1_vals[elem + 1]), + fmul2( + (weight00, weight00), + (o0_vals[elem], o0_vals[elem + 1]), + ), + ) + return cutlass.Vector.from_elements(combined, Float32) + + @cute.jit + def _kv_tile_256_merge_spatial_output( + self, + *, + exchange: cutlass.Array, + exchange_idx: Int32, + base_addr0: Int32, + base_addr1: Int32, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + cta_idx_kv: Int32, + logical_q_group_idx: Int32, + q_row_offset: Int32, + weight00: Float32, + weight10: Float32, + denominator: Float32, + global_max: Float32, + ) -> None: + """Merge KV256 spatial halves one D32 fragment at a time. + + All four correction warps compute their temporal partial before each + barrier. Upper lanes publish their D32 fragment while lower lanes keep + the matching D32 in registers, then consume the peer and write either + the direct or split-KV ABI. This preserves the 64-row exchange without + serializing the complete D128 upper and lower halves. + """ + cfg = self.cfg + partial_o_uses_bf16 = ( + cfg.use_bf16_separate_partial_o + if cfg.use_separate_reduction_kernel + else cfg.use_bf16_output + ) + output_exchange_base = Int32( + _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD + ) + output_lane = exchange_idx < Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) + exchange_row_idx = exchange_idx & Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS - 1) + output_exchange_row_base = output_exchange_base + exchange_row_idx * Int32( + _KV_TILE_256_EXCHANGE_ROW_STRIDE + ) + logical_output_row_idx = q_row_offset + exchange_row_idx + valid_output_row = cutlass.Boolean(False) + + if cutlass.const_expr(cfg.use_split_kv): + partial_scale = Float32(1.0) + partial_row_base = Int64(0) + if output_lane: + valid_output_row = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + exchange_row_idx, + self.seq_len_q, + ) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + partial_scale = self._separate_partial_norm_scale(denominator) + partial_row_base = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) * Int64(cfg.headdim * 2) + else: + dst_row_base = Int64(0) + norm_scale = Float32(1.0) + if output_lane: + valid_output_row = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + exchange_row_idx, + self.seq_len_q, + ) + dst_row_base, norm_scale = self._softmax_output_row_state( + logical_h_k_idx, + logical_b_idx, + logical_output_row_idx, + denominator, + global_max, + ) + + for fragment in cutlass.range_constexpr(cfg.headdim // 32): + fragment_col = fragment * 32 + own_vals = self._kv_tile_256_temporal_fragment( + base_addr0=base_addr0, + base_addr1=base_addr1, + fragment_col=fragment_col, + weight00=weight00, + weight10=weight10, + ) + if exchange_idx >= Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS): + ( + exchange.data_ptr() + output_exchange_row_base + Int32(fragment_col) + ).store(own_vals, alignment=16) + + # Lower lanes keep ``own_vals`` live across the barrier. Once every + # lane arrives, upper lanes may prepare the next fragment while + # lower lanes consume the current peer fragment. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + if output_lane: + peer_vals = ( + exchange.data_ptr() + output_exchange_row_base + Int32(fragment_col) + ).load(count=32, alignment=16) + for vector_idx in cutlass.range_constexpr(4): + vector_col = vector_idx * 8 + output_vals = cutlass.Array( + Float32, + 8, + space=cutlass.AddressSpace.rmem, + ) + for elem in cutlass.range_constexpr(0, 8, 2): + value_idx = vector_col + elem + merged = fadd2( + (own_vals[value_idx], own_vals[value_idx + 1]), + ( + Float32(peer_vals[value_idx]), + Float32(peer_vals[value_idx + 1]), + ), + ) + output_vals[elem] = merged[0] + output_vals[elem + 1] = merged[1] + if valid_output_row: + output_col = fragment_col + vector_col + if cutlass.const_expr(cfg.use_split_kv): + scaled_values: tuple = () + for elem in cutlass.range_constexpr(0, 8, 2): + scaled_values += fmul2( + (partial_scale, partial_scale), + (output_vals[elem], output_vals[elem + 1]), + ) + scaled_vector = cutlass.Vector.from_elements( + scaled_values, Float32 + ) + if cutlass.const_expr(partial_o_uses_bf16): + packed = scaled_vector.to(cutlass.BFloat16).bitcast( + Int32 + ) + else: + packed = scaled_vector.to(cutlass.Float16).bitcast( + Int32 + ) + partial_o_dst = cutlass.inttoptr( + self.partial_o_ptr.toint() + + partial_row_base + + Int64(output_col * cfg.o_dtype_bytes), + mem_space=1, + dtype=Int32, + ) + partial_o_dst.store(packed, alignment=16) + else: + dst_offset = dst_row_base + Int32( + output_col * cfg.o_dtype_bytes + ) + final_o_dst = cutlass.inttoptr( + self.o_ptr.toint() + cutlass.Int64(dst_offset), + mem_space=1, + dtype=Int32, + ) + self._store_final_o_vec8( + final_o_dst, + output_vals, + norm_scale, + ) + + if cutlass.const_expr(cfg.use_split_kv): + if valid_output_row: + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + stats_row = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) + stats_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_row * Int64(4), + mem_space=1, + dtype=Float32, + ) + stats_ptr.store( + self._separate_partial_lse(global_max, denominator), + alignment=4, + ) + else: + stats_row = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) + stats_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_row * Int64(2 * 4), + mem_space=1, + dtype=Float32, + ) + stats_ptr.store( + cutlass.Vector.from_elements( + (global_max, denominator), + Float32, + ), + alignment=8, + ) + + @cute.jit + def _kv_tile_256_tail_epilogue( + self, + stage_info: StageInfo, + *, + scratch_stage: Int32 | None, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + inst0_new_max_arr: cutlass.Array, + inst0_sum_arr: cutlass.Array, + inst1_new_max_arr: cutlass.Array, + inst1_sum_arr: cutlass.Array, + tmem_row_base: Int32, + o_base_col: Constexpr[int], + warp_grp_thread_idx: Int32, + ) -> None: + """Merge KV256's two temporal and two spatial output partials. + + The standard decode schedule still owns the two temporal instances. + KV256 adds one physical spatial split per instance. Correction exchanges + their stats, stages one spatial half in SMEM after the shared KV ring is + dead, then publishes the ordinary logical Q64xD128 output. + """ + cfg = self.cfg + assert cfg.headdim == 128 + exchange = self._kv_tile_256_exchange_for_stage( + stage_info, + scratch_stage, + ) + + exchange_idx = warp_grp_thread_idx + peer_idx = exchange_idx ^ Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) + stats_base = exchange_idx * Int32(_KV_TILE_256_STATS_PER_THREAD) + exchange[stats_base] = inst0_new_max_arr[0] + exchange[stats_base + Int32(1)] = inst0_sum_arr[0] + exchange[stats_base + Int32(2)] = inst1_new_max_arr[0] + exchange[stats_base + Int32(3)] = inst1_sum_arr[0] + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + peer_stats_base = peer_idx * Int32(_KV_TILE_256_STATS_PER_THREAD) + max00 = inst0_new_max_arr[0] + sum00 = inst0_sum_arr[0] + max10 = inst1_new_max_arr[0] + sum10 = inst1_sum_arr[0] + max01 = Float32(exchange[peer_stats_base]) + sum01 = Float32(exchange[peer_stats_base + Int32(1)]) + max11 = Float32(exchange[peer_stats_base + Int32(2)]) + sum11 = Float32(exchange[peer_stats_base + Int32(3)]) + + uses00 = max00 != _neg_max_f32() + uses10 = max10 != _neg_max_f32() + uses01 = max01 != _neg_max_f32() + uses11 = max11 != _neg_max_f32() + global_max = _neg_max_f32() + if uses00: + global_max = max00 + if uses10: + global_max = cute.math.max(global_max, max10, ftz=True) + if uses01: + global_max = cute.math.max(global_max, max01, ftz=True) + if uses11: + global_max = cute.math.max(global_max, max11, ftz=True) + + weight00 = Float32(0.0) + weight10 = Float32(0.0) + denominator = Float32(0.0) + if uses00: + weight00 = cute.math.exp2( + self.scale_softmax_log2 * (max00 - global_max), + fastmath=True, + ) + denominator += sum00 * weight00 + if uses10: + weight10 = cute.math.exp2( + self.scale_softmax_log2 * (max10 - global_max), + fastmath=True, + ) + denominator += sum10 * weight10 + if uses01: + denominator += sum01 * cute.math.exp2( + self.scale_softmax_log2 * (max01 - global_max), + fastmath=True, + ) + if uses11: + denominator += sum11 * cute.math.exp2( + self.scale_softmax_log2 * (max11 - global_max), + fastmath=True, + ) + + base_addr0 = tmem_row_base + Int32( + o_base_col + tail_o_stage_idx_0 * cfg.tmem_o_stage_cols + ) + base_addr1 = tmem_row_base + Int32( + o_base_col + tail_o_stage_idx_1 * cfg.tmem_o_stage_cols + ) + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + q_row_offset = _q_tile_output_row_base(cfg, logical_q_group_idx) + logical_kv_idx = logical_b_idx * self.num_heads_kv + logical_h_k_idx + splits_kv = Int32(1) + cta_idx_kv = Int32(0) + if cutlass.const_expr(cfg.use_split_kv): + splits_kv = self._runtime_splits_kv(stage_info) + cta_idx_kv = _logical_cta_kv_idx(cfg, stage_info) + + self._kv_tile_256_merge_spatial_output( + exchange=exchange, + exchange_idx=exchange_idx, + base_addr0=base_addr0, + base_addr1=base_addr1, + logical_h_k_idx=logical_h_k_idx, + logical_b_idx=logical_b_idx, + logical_kv_idx=logical_kv_idx, + cta_idx_kv=cta_idx_kv, + logical_q_group_idx=logical_q_group_idx, + q_row_offset=q_row_offset, + weight00=weight00, + weight10=weight10, + denominator=denominator, + global_max=global_max, + ) + + if cutlass.const_expr( + cfg.use_split_kv and not cfg.use_separate_reduction_kernel + ): + # Every correction lane joins the generic fused-GMEM completion + # protocol after the low 64 lanes publish one full D128 row each. + counter_q_groups = self._multi_cta_counter_q_groups(self.h_r) + counter_group_idx = logical_kv_idx * counter_q_groups + logical_q_group_idx + self._reduce_fused_gmem_partials( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + counter_group_idx, + splits_kv, + warp_grp_thread_idx, + ) + + # The next persistent work tile may reuse this exchange allocation. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + @cute.jit + def _keeps_tail_epilogue( + self, + stage_info: StageInfo, + *, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + inst0_new_max_arr: cutlass.Array, + inst0_sum_arr: cutlass.Array, + inst1_new_max_arr: cutlass.Array, + inst1_sum_arr: cutlass.Array, + tmem_row_base: Int32, + o_base_col: Constexpr[int], + output_pair_regs: Constexpr[int], + keeps_o_ldst_offset: Constexpr[int], + row_idx: Int32, + col_base: Int32, + warp_grp_thread_idx: Int32, + ) -> None: + """Finish the KeepsMmaAb tail stage and optional split-KV reduction.""" + cfg = self.cfg + # Tail stats are already final per-instance sums and maxima. Combine + # inst0/inst1 first, then either publish split-KV partials or normalize + # directly to the output tensor. + inst0_sum_0 = inst0_sum_arr[0] + inst0_new_max_0 = inst0_new_max_arr[0] + inst1_sum_0 = inst1_sum_arr[0] + inst1_new_max_0 = inst1_new_max_arr[0] + + if cutlass.const_expr(cfg.tile_size_q == 64): + # The q64 16dp32bitx2 layout keeps the two 64-K half-row sums + # separate throughout the online recurrence. Pair them only at + # finalization, after which both lanes use the same denominator + # while continuing to own disjoint O-column halves. + inst0_sum_0 += Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=inst0_sum_0, + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + if cutlass.const_expr(cfg.num_insts_kv != 1): + inst1_sum_0 += Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=inst1_sum_0, + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + + if cutlass.const_expr( + cfg.use_fp8_qkv + and cfg.use_fp8_output + and cfg.tile_size_q == 128 + and cfg.headdim == 128 + and cfg.num_insts_kv == 2 + and cfg.max_seq_len_q == 1 + and cfg.has_static_dense_full_kv_tiles + ): + final_max_0 = cute.math.max(inst0_new_max_0, inst1_new_max_0, ftz=True) + exp_scale0_0 = cute.math.exp2( + self.scale_softmax_log2 * (inst0_new_max_0 - final_max_0), + fastmath=True, + ) + exp_scale1_0 = cute.math.exp2( + self.scale_softmax_log2 * (inst1_new_max_0 - final_max_0), + fastmath=True, + ) + reduced_sum_0 = inst0_sum_0 * exp_scale0_0 + inst1_sum_0 * exp_scale1_0 + else: + uses_inst0 = inst0_new_max_0 != _neg_max_f32() + uses_inst1 = False + if cutlass.const_expr(cfg.num_insts_kv != 1): + uses_inst1 = inst1_new_max_0 != _neg_max_f32() + final_max_0 = _neg_max_f32() + if uses_inst0: + final_max_0 = inst0_new_max_0 + if uses_inst1: + final_max_0 = cute.math.max(final_max_0, inst1_new_max_0, ftz=True) + + exp_scale0_0 = Float32(0.0) + exp_scale1_0 = Float32(0.0) + reduced_sum_0 = Float32(0.0) + if uses_inst0: + exp_scale0_0 = cute.math.exp2( + self.scale_softmax_log2 * (inst0_new_max_0 - final_max_0), + fastmath=True, + ) + reduced_sum_0 += inst0_sum_0 * exp_scale0_0 + if uses_inst1: + exp_scale1_0 = cute.math.exp2( + self.scale_softmax_log2 * (inst1_new_max_0 - final_max_0), + fastmath=True, + ) + reduced_sum_0 += inst1_sum_0 * exp_scale1_0 + + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + # The physical MMA tile may end with structural padding. Scratch rows + # omit that padding, so advance Q groups by the complete token/head rows + # represented by the Q tensor map rather than by ``tile_size_q``. + q_row_offset = _q_tile_output_row_base(cfg, logical_q_group_idx) + logical_output_row_idx = q_row_offset + row_idx + valid_output_row = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + row_idx, + self.seq_len_q, + ) + if cutlass.const_expr(cfg.use_split_kv): + logical_kv_idx = logical_b_idx * self.num_heads_kv + logical_h_k_idx + splits_kv = self._runtime_splits_kv(stage_info) + cta_idx_kv = _logical_cta_kv_idx(cfg, stage_info) + base_addr0 = tmem_row_base + Int32( + o_base_col + tail_o_stage_idx_0 * cfg.tmem_o_stage_cols + ) + base_addr1 = tmem_row_base + Int32( + o_base_col + tail_o_stage_idx_1 * cfg.tmem_o_stage_cols + ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + base_addr1 = base_addr0 + partial_dst_col_offset = col_base * Int32(2) + partial_norm_scale = Float32(1.0) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + partial_norm_scale = self._separate_partial_norm_scale(reduced_sum_0) + elif cutlass.const_expr(cfg.use_fp8_qkv): + partial_norm_scale = Float32(1.0 / 448.0) + regs_o_chunk = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + partial_o_row_base = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) * Int64(cfg.headdim * 2) + for store_idx in cutlass.range_constexpr(output_pair_regs // 4): + chunk_col = store_idx * 8 + # Load O0/O1 from the tail TMEM slots, combine with + # max-correction scales, then store normalized 16-bit O for + # standalone reduction or unnormalized O for fused reduction. + o0_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr0 + Int32(chunk_col), Float32), + num=8, + offset=keeps_o_ldst_offset, + ) + if cutlass.const_expr(cfg.num_insts_kv != 1): + o1_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr1 + Int32(chunk_col), Float32), + num=8, + offset=keeps_o_ldst_offset, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for chunk_idx in cutlass.range_constexpr(4): + reg_base = chunk_idx * 2 + if cutlass.const_expr(cfg.num_insts_kv == 1): + partial_pair = fmul2( + (exp_scale0_0, exp_scale0_0), + ( + o0_vals[reg_base], + o0_vals[reg_base + 1], + ), + ) + else: + partial_pair = ffma2( + (exp_scale1_0, exp_scale1_0), + (o1_vals[reg_base], o1_vals[reg_base + 1]), + fmul2( + (exp_scale0_0, exp_scale0_0), + ( + o0_vals[reg_base], + o0_vals[reg_base + 1], + ), + ), + ) + if cutlass.const_expr( + cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + ): + partial_pair = fmul2( + (partial_norm_scale, partial_norm_scale), partial_pair + ) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + regs_o_chunk[chunk_idx] = self._pack_separate_partial_o_pair( + partial_pair[0], partial_pair[1] + ) + elif cutlass.const_expr(cfg.use_bf16_output): + regs_o_chunk[chunk_idx] = _pack_float2_to_bf16( + partial_pair[0], partial_pair[1] + ) + else: + regs_o_chunk[chunk_idx] = _pack_float2_to_fp16( + partial_pair[0], partial_pair[1] + ) + partial_o_dst = cutlass.inttoptr( + self.partial_o_ptr.toint() + + partial_o_row_base + + Int64(partial_dst_col_offset) + + Int64(store_idx * 16), + mem_space=1, + dtype=Int32, + ) + if valid_output_row: + partial_o_dst.store( + regs_o_chunk.data_ptr().load(count=4, alignment=4), + alignment=16, + ) + + if col_base == Int32(0) and valid_output_row: + # One state record per output row accompanies all partial-O + # vector chunks for that row. + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + stats_row = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) + stats_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_row * Int64(4), + mem_space=1, + dtype=Float32, + ) + stats_ptr.store( + self._separate_partial_lse(final_max_0, reduced_sum_0), + alignment=4, + ) + else: + stats_row = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) + stats_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_row * Int64(2 * 4), + mem_space=1, + dtype=Float32, + ) + stats_ptr.store( + cutlass.Vector.from_elements( + (final_max_0, reduced_sum_0), Float32 + ), + alignment=8, + ) + + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + # Separate reducer kernel will consume the partial O/stats + # written above; this CTA has no in-kernel reduction work. + return + + # In-kernel GMEM reducer: all split CTAs publish partials, then + # synchronize before each owner CTA reduces its 2-KiB slice band. + counter_q_groups = self._multi_cta_counter_q_groups(self.h_r) + counter_group_idx = logical_kv_idx * counter_q_groups + logical_q_group_idx + self._reduce_fused_gmem_partials( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + counter_group_idx, + splits_kv, + warp_grp_thread_idx, + ) + return + + attention_sink_h_r = _attention_sink_head_stride(cfg, self.h_r) + attention_sink_head_idx = _local_head_from_q_output_row( + cfg, self.h_r, logical_output_row_idx + ) + reduced_sum_0 += _attention_sink_for_local_head( + cfg, + self.attention_sinks_ptr, + self.scale_softmax_log2, + final_max_0, + logical_h_k_idx, + attention_sink_h_r, + self.num_heads_kv, + attention_sink_head_idx, + ) + # Apply public bmm2_scale in the final direct-output normalization. + # The split-KV branch above returns before reaching this point, so its + # partial O remains unscaled for the selected reducer. + norm_scale_0 = self.output_scale * self._safe_norm_rcp(reduced_sum_0) + # Direct-output path: fold attention sinks into the denominator, then + # bake normalization and per-instance max correction into the O scales. + final_scale0_0 = norm_scale_0 * exp_scale0_0 + final_scale1_0 = norm_scale_0 * exp_scale1_0 + + base_addr0 = tmem_row_base + Int32( + o_base_col + tail_o_stage_idx_0 * cfg.tmem_o_stage_cols + ) + base_addr1 = tmem_row_base + Int32( + o_base_col + tail_o_stage_idx_1 * cfg.tmem_o_stage_cols + ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + base_addr1 = base_addr0 + physical_dst_row_idx = _q_physical_output_row( + cfg, + self.h_r, + self.num_heads_kv, + logical_b_idx, + logical_h_k_idx, + logical_q_group_idx, + row_idx, + self.q_token_offset, + ) + dst_row_base = Int64(physical_dst_row_idx) * Int64( + cfg.headdim * cfg.o_dtype_bytes + ) + dst_col_offset = Int64(col_base) * Int64(cfg.o_dtype_bytes) + if cutlass.const_expr( + cfg.use_fp8_output and cfg.tile_size_q == 128 and cfg.headdim == 128 + ): + regs_o_chunk = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for store_idx in cutlass.range_constexpr(output_pair_regs // 8): + chunk_col = store_idx * 16 + # FP8 output path: load a wider O chunk, apply final scales, + # pack four values per register, and store directly to GMEM. + o0_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr0 + Int32(chunk_col), Float32), + num=16, + offset=keeps_o_ldst_offset, + ) + if cutlass.const_expr(cfg.num_insts_kv != 1): + o1_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr1 + Int32(chunk_col), Float32), + num=16, + offset=keeps_o_ldst_offset, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for packed_idx in cutlass.range_constexpr(4): + reg_base = packed_idx * 4 + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair0 = fmul2( + (final_scale0_0, final_scale0_0), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ) + final_pair1 = fmul2( + (final_scale0_0, final_scale0_0), + ( + o0_vals[reg_base + 2], + o0_vals[reg_base + 3], + ), + ) + else: + final_pair0 = ffma2( + (final_scale1_0, final_scale1_0), + (o1_vals[reg_base], o1_vals[reg_base + 1]), + fmul2( + (final_scale0_0, final_scale0_0), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ), + ) + final_pair1 = ffma2( + (final_scale1_0, final_scale1_0), + ( + o1_vals[reg_base + 2], + o1_vals[reg_base + 3], + ), + fmul2( + (final_scale0_0, final_scale0_0), + ( + o0_vals[reg_base + 2], + o0_vals[reg_base + 3], + ), + ), + ) + regs_o_chunk[packed_idx] = _pack_float4_to_fp8_e4m3( + final_pair0[0], + final_pair0[1], + final_pair1[0], + final_pair1[1], + ) + dst_ptr = cutlass.inttoptr( + self.o_ptr.toint() + + dst_row_base + + dst_col_offset + + Int64(store_idx * 16), + mem_space=1, + dtype=Int32, + ) + if valid_output_row: + dst_ptr.store( + regs_o_chunk.data_ptr().load(count=4, alignment=4), + alignment=16, + ) + return + + regs_o_chunk = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for store_idx in cutlass.range_constexpr(output_pair_regs // 4): + chunk_col = store_idx * 8 + # FP16/BF16 or generic FP8 path: load one output chunk from each + # tail O stage, combine the two instances, then pack to output type. + o0_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr0 + Int32(chunk_col), Float32), + num=8, + offset=keeps_o_ldst_offset, + ) + if cutlass.const_expr(cfg.num_insts_kv != 1): + o1_vals = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr1 + Int32(chunk_col), Float32), + num=8, + offset=keeps_o_ldst_offset, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + if cutlass.const_expr(cfg.use_fp8_output): + for packed_idx in cutlass.range_constexpr(2): + reg_base = packed_idx * 4 + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair0 = fmul2( + (final_scale0_0, final_scale0_0), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ) + final_pair1 = fmul2( + (final_scale0_0, final_scale0_0), + ( + o0_vals[reg_base + 2], + o0_vals[reg_base + 3], + ), + ) + else: + final_pair0 = ffma2( + (final_scale1_0, final_scale1_0), + (o1_vals[reg_base], o1_vals[reg_base + 1]), + fmul2( + (final_scale0_0, final_scale0_0), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ), + ) + final_pair1 = ffma2( + (final_scale1_0, final_scale1_0), + ( + o1_vals[reg_base + 2], + o1_vals[reg_base + 3], + ), + fmul2( + (final_scale0_0, final_scale0_0), + ( + o0_vals[reg_base + 2], + o0_vals[reg_base + 3], + ), + ), + ) + regs_o_chunk[packed_idx] = _pack_float4_to_fp8_e4m3( + final_pair0[0], + final_pair0[1], + final_pair1[0], + final_pair1[1], + ) + else: + for chunk_idx in cutlass.range_constexpr(4): + reg_base = chunk_idx * 2 + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair = fmul2( + (final_scale0_0, final_scale0_0), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ) + else: + final_pair = ffma2( + (final_scale1_0, final_scale1_0), + (o1_vals[reg_base], o1_vals[reg_base + 1]), + fmul2( + (final_scale0_0, final_scale0_0), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ), + ) + if cutlass.const_expr(cfg.use_bf16_output): + regs_o_chunk[chunk_idx] = _pack_float2_to_bf16( + final_pair[0], final_pair[1] + ) + else: + regs_o_chunk[chunk_idx] = _pack_float2_to_fp16( + final_pair[0], final_pair[1] + ) + dst_ptr = cutlass.inttoptr( + self.o_ptr.toint() + + dst_row_base + + dst_col_offset + + Int64(store_idx * (8 if cfg.use_fp8_output else 16)), + mem_space=1, + dtype=Int32, + ) + if cutlass.const_expr(cfg.use_fp8_output): + if valid_output_row: + dst_ptr.store( + regs_o_chunk.data_ptr().load(count=2, alignment=4), + alignment=8, + ) + else: + if valid_output_row: + dst_ptr.store( + regs_o_chunk.data_ptr().load(count=4, alignment=4), + alignment=16, + ) + return + + @cute.jit + def _swaps_wide_tail_epilogue( + self, + stage_info: StageInfo, + *, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + inst0_new_max_arr: cutlass.Array, + inst0_sum_arr: cutlass.Array, + inst1_new_max_arr: cutlass.Array, + inst1_sum_arr: cutlass.Array, + tmem_row_base: Int32, + o_base_col: Constexpr[int], + q_repeats: Constexpr[int], + num_scale_groups: Constexpr[int], + output_pair_regs: Constexpr[int], + output_f32_regs: Constexpr[int], + num_o_chunks: Constexpr[int], + ) -> None: + """Finish the SwapsMmaAb TileSizeQ 16/32 tail stage.""" + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + # Tail ProdWork: combine inst0 and inst1 stats, reduce the + # final denominator across correction warps, then normalize + # and output the final O tile. + final_max = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + final_sum = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + exp_scale0 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + exp_scale1 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + reduced_sum = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + final_scale0 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + final_scale1 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + col_group_idx = warp_grp_thread_idx & Int32(0x3) + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 2): + final_sum_pair_arr = cutlass.Array( + Float32, 2, space=cutlass.AddressSpace.rmem + ) + # Merge the two softmax instances with the standard + # log-sum-exp correction. Some tail stages can be inactive, + # so skip entries whose max is still the sentinel value. + for pair_idx in cutlass.range_constexpr(2): + scale_idx = scale_base + pair_idx + inst0_max = inst0_new_max_arr[scale_idx] + inst1_max = inst1_new_max_arr[scale_idx] + uses_inst0 = inst0_max != _neg_max_f32() + uses_inst1 = inst1_max != _neg_max_f32() + + final_max_val = _neg_max_f32() + if uses_inst0: + final_max_val = inst0_max + if uses_inst1: + final_max_val = cute.math.max(final_max_val, inst1_max, ftz=True) + + exp_scale0_val = Float32(0.0) + exp_scale1_val = Float32(0.0) + final_sum_val = Float32(0.0) + if uses_inst0: + exp_scale0_val = cute.math.exp2( + self.scale_softmax_log2 * (inst0_max - final_max_val), + fastmath=True, + ) + final_sum_val += inst0_sum_arr[scale_idx] * exp_scale0_val + if uses_inst1: + exp_scale1_val = cute.math.exp2( + self.scale_softmax_log2 * (inst1_max - final_max_val), + fastmath=True, + ) + final_sum_val += inst1_sum_arr[scale_idx] * exp_scale1_val + + final_max[scale_idx] = final_max_val + exp_scale0[scale_idx] = exp_scale0_val + exp_scale1[scale_idx] = exp_scale1_val + final_sum_pair_arr[pair_idx] = final_sum_val + + final_sum_pair = self._warp_reduce_col_group_sum_pair( + (final_sum_pair_arr[0], final_sum_pair_arr[1]) + ) + final_sum[scale_base] = final_sum_pair[0] + final_sum[scale_base + 1] = final_sum_pair[1] + + warp_store_base = warp_idx * Int32( + 4 * num_scale_groups + ) + col_group_idx * Int32(num_scale_groups) + if lane_idx < Int32(4): + # Store one warp partial per column group. The following CTA + # barrier and reload combine the four correction warps for the + # final denominator. + if cutlass.const_expr(num_scale_groups == 8): + self._sum_scratch.store( + ( + final_sum[0], + final_sum[1], + final_sum[2], + final_sum[3], + final_sum[4], + final_sum[5], + final_sum[6], + final_sum[7], + ), + warp_store_base, + alignment=16, + ) + else: + self._sum_scratch.store( + ( + final_sum[0], + final_sum[1], + final_sum[2], + final_sum[3], + ), + warp_store_base, + alignment=16, + ) + prims.barrier_cta_sync( + self.sum_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + reduce_base = col_group_idx * Int32(num_scale_groups) + reduced_vec = self._sum_scratch.load( + reduce_base, + vector_size=num_scale_groups, + alignment=16, + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + reduced_sum[scale_idx] = reduced_vec[scale_idx] + for warp_offset in cutlass.range_constexpr(1, 4): + # Combine the four warp partials for this column group. + other_vec = self._sum_scratch.load( + reduce_base + warp_offset * Int32(4 * num_scale_groups), + vector_size=num_scale_groups, + alignment=16, + ) + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 2): + reduced_pair = fadd2( + ( + reduced_sum[scale_base], + reduced_sum[scale_base + 1], + ), + (other_vec[scale_base], other_vec[scale_base + 1]), + ) + reduced_sum[scale_base] = reduced_pair[0] + reduced_sum[scale_base + 1] = reduced_pair[1] + if cutlass.const_expr(not cfg.use_split_kv): + logical_h_k_idx, _ = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + # Attention sinks add to the denominator before the + # final normalization scale is computed. + reduced_sum[scale_idx] += _attention_sink_for_scale_idx( + cfg, + self.attention_sinks_ptr, + self.scale_softmax_log2, + final_max[scale_idx], + logical_h_k_idx, + self.h_r, + self.num_heads_kv, + logical_q_group_idx, + col_group_idx, + scale_idx, + ) + + base_addr0 = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, tail_o_stage_idx_0 + ) + base_addr1 = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, tail_o_stage_idx_1 + ) + o0_vals, o1_vals = self._swaps_load_two_o_stage_chunks( + base_addr0, + base_addr1, + q_repeats=q_repeats, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + + if cutlass.const_expr(cfg.use_split_kv): + # Split-KV tail: separate reduction stores normalized 16-bit O + # plus log2-LSE; fused GMEM/cluster retains unnormalized O plus + # max/sum state. + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + q_row_offset = _q_tile_output_row_base(cfg, logical_q_group_idx) + logical_kv_idx = logical_b_idx * self.num_heads_kv + logical_h_k_idx + splits_kv = self._runtime_splits_kv(stage_info) + cta_idx_kv = _logical_cta_kv_idx(cfg, stage_info) + + if cutlass.const_expr(cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv): + for scale_idx in cutlass.range_constexpr(num_scale_groups): + norm_scale = Float32(1.0 / 448.0) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + norm_scale = self._separate_partial_norm_scale( + reduced_sum[scale_idx] + ) + final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] + final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] + + regs_partial_o = cutlass.Array( + Int32, output_pair_regs, space=cutlass.AddressSpace.rmem + ) + for pair_idx in cutlass.range_constexpr(output_pair_regs): + # Separate reduction includes this split's reciprocal sum; + # fused reduction delays normalization until the final merge. + scale_base = ((pair_idx % (2 * q_repeats)) // 2) * 2 + reg_base = pair_idx * 2 + partial_scale0 = ( + (final_scale0[scale_base], final_scale0[scale_base + 1]) + if cutlass.const_expr( + cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + ) + else (exp_scale0[scale_base], exp_scale0[scale_base + 1]) + ) + partial_scale1 = ( + (final_scale1[scale_base], final_scale1[scale_base + 1]) + if cutlass.const_expr( + cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + ) + else (exp_scale1[scale_base], exp_scale1[scale_base + 1]) + ) + partial_pair = ffma2( + partial_scale0, + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + partial_scale1, + (o1_vals[reg_base], o1_vals[reg_base + 1]), + ), + ) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + regs_partial_o[pair_idx] = self._pack_separate_partial_o_pair( + partial_pair[0], partial_pair[1] + ) + elif cutlass.const_expr(cfg.use_bf16_output and not cfg.use_fp8_output): + regs_partial_o[pair_idx] = _pack_float2_to_bf16( + partial_pair[0], partial_pair[1] + ) + else: + regs_partial_o[pair_idx] = _pack_float2_to_fp16( + partial_pair[0], partial_pair[1] + ) + + num_copy_segments = max( + (cfg.tile_size_q * cfg.headdim * 2 + 2047) // 2048, + 1, + ) + if cutlass.const_expr(cfg.supports_cluster_smem_reduction): + self._publish_and_reduce_cluster_swaps_partials( + regs_partial_o, + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + splits_kv, + warp_grp_thread_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + output_pair_regs=output_pair_regs, + num_copy_segments=num_copy_segments, + num_scale_groups=num_scale_groups, + ) + return + + self._stage_and_copy_swaps_partial_o( + cfg, + regs_partial_o, + logical_kv_idx, + cta_idx_kv, + splits_kv, + q_row_offset, + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + output_pair_regs=output_pair_regs, + num_copy_segments=num_copy_segments, + enable_cluster=False, + full_prefix=False, + ) + + self._store_split_stats_from_scale_arrays( + cfg, + logical_kv_idx, + cta_idx_kv, + splits_kv, + q_row_offset, + logical_q_group_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + num_scale_groups=num_scale_groups, + enable_cluster=False, + full_prefix=False, + ) + + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + # The standalone reducer consumes the partial O/stats above; + # do not enter the in-kernel completion/election protocol. + return + + counter_q_groups = self._multi_cta_counter_q_groups(self.h_r) + counter_group_idx = logical_kv_idx * counter_q_groups + logical_q_group_idx + self._reduce_fused_gmem_partials( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + counter_group_idx, + splits_kv, + warp_grp_thread_idx, + ) + return + + for scale_idx in cutlass.range_constexpr(num_scale_groups): + # Direct final-output scales include both the reciprocal softmax + # denominator and public bmm2_scale. Split partials returned above + # use only the unnormalized ``exp_scale*`` arrays. + norm_scale = self.output_scale * self._safe_norm_rcp(reduced_sum[scale_idx]) + final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] + final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] + + regs_o = cutlass.Array( + Int32, + cfg.num_fp8_output_regs if cfg.use_fp8_output else output_pair_regs, + space=cutlass.AddressSpace.rmem, + ) + if cutlass.const_expr(cfg.use_fp8_output): + # Fold and pack two adjacent pairs immediately. Keeping a second + # full FP32 output array live increases peak register pressure and + # spilling before the transposed STSM. + for packed_idx in cutlass.range_constexpr(cfg.num_fp8_output_regs): + pair_idx0 = packed_idx * 2 + pair_idx1 = pair_idx0 + 1 + scale_base0 = ((pair_idx0 % (2 * q_repeats)) // 2) * 2 + scale_base1 = ((pair_idx1 % (2 * q_repeats)) // 2) * 2 + reg_base0 = pair_idx0 * 2 + reg_base1 = pair_idx1 * 2 + final_pair0 = ffma2( + (final_scale0[scale_base0], final_scale0[scale_base0 + 1]), + (o0_vals[reg_base0], o0_vals[reg_base0 + 1]), + fmul2( + ( + final_scale1[scale_base0], + final_scale1[scale_base0 + 1], + ), + (o1_vals[reg_base0], o1_vals[reg_base0 + 1]), + ), + ) + final_pair1 = ffma2( + (final_scale0[scale_base1], final_scale0[scale_base1 + 1]), + (o0_vals[reg_base1], o0_vals[reg_base1 + 1]), + fmul2( + ( + final_scale1[scale_base1], + final_scale1[scale_base1 + 1], + ), + (o1_vals[reg_base1], o1_vals[reg_base1 + 1]), + ), + ) + # bmm2_scale is already folded into ``final_scale*`` above. + regs_o[packed_idx] = _pack_float4_to_fp8_e4m3( + final_pair0[0], + final_pair0[1], + final_pair1[0], + final_pair1[1], + ) + else: + for pair_idx in cutlass.range_constexpr(output_pair_regs): + # non-split-KV path: combine inst0/inst1 O and pack directly + # to the final output dtype. + scale_base = ((pair_idx % (2 * q_repeats)) // 2) * 2 + reg_base = pair_idx * 2 + final_pair = ffma2( + (final_scale0[scale_base], final_scale0[scale_base + 1]), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + ( + final_scale1[scale_base], + final_scale1[scale_base + 1], + ), + (o1_vals[reg_base], o1_vals[reg_base + 1]), + ), + ) + if cutlass.const_expr(cfg.use_bf16_output): + regs_o[pair_idx] = _pack_float2_to_bf16( + final_pair[0], final_pair[1] + ) + else: + regs_o[pair_idx] = _pack_float2_to_fp16( + final_pair[0], final_pair[1] + ) + + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + if cutlass.const_expr(cfg.use_fp8_output): + # FP8 final output is staged with transposed 8-bit + # STSM and then copied to GMEM as contiguous vectors. This two-step + # layout is required because correction registers are in MMA + # fragment order, not output row-major order. + _store_transposed_smem8b( + self._smem_base_o_i32, + regs_o.data_ptr().load( + count=cfg.num_fp8_output_regs, + alignment=4, + ), + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.headdim, + cfg.num_fp8_output_regs, + ) + cute.arch.fence_view_async_shared() + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + _copy_transposed_smem8b_to_gmem( + self._smem_base_o_i32, + self.o_ptr, + cfg, + logical_h_k_idx, + logical_b_idx, + logical_q_group_idx, + self.h_r, + self.num_heads_kv, + self.seq_len_q, + self.q_token_offset, + warp_grp_thread_idx, + cfg.fp8_copy_can_use_full_tile_fast_path, + ) + if cutlass.const_expr(cfg.use_persistent_scheduler): + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + return + + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + self._stage_fp16_o_regs_to_smem( + cfg, + regs_o, + warp_grp_thread_idx, + warp_idx, + lane_idx, + output_pair_regs=output_pair_regs, + ) + num_copy_segments = max( + (cfg.tile_size_q * cfg.headdim * cfg.o_dtype_bytes + 2047) // 2048, + 1, + ) + self._copy_staged_fp16_o_to_output( + cfg, + logical_h_k_idx, + logical_b_idx, + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + num_copy_segments=num_copy_segments, + ) + if cutlass.const_expr(cfg.use_persistent_scheduler): + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + return + + @cute.jit + def _swaps_q8_tail_epilogue( + self, + stage_info: StageInfo, + *, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + inst0_new_max_arr: cutlass.Array, + inst0_sum_arr: cutlass.Array, + inst1_new_max_arr: cutlass.Array, + inst1_sum_arr: cutlass.Array, + tmem_row_base: Int32, + o_base_col: Constexpr[int], + ) -> None: + """Finish the SwapsMmaAb TileSizeQ 8 tail stage.""" + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + # Tile-Q=8 tail ProdWork: combine the two final O stages, + # normalize, and either emit final O or reduce split-KV + # partials. + num_scale_groups = 2 + final_max = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + final_sum = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + exp_scale0 = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + exp_scale1 = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + reduced_sum = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + final_scale0 = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + final_scale1 = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + + # Compute per-instance exp corrections relative to the final + # max for the two scale groups. + for scale_idx in cutlass.range_constexpr(num_scale_groups): + inst0_max = inst0_new_max_arr[scale_idx] + inst1_max = inst1_new_max_arr[scale_idx] + uses_inst0 = inst0_max != _neg_max_f32() + uses_inst1 = inst1_max != _neg_max_f32() + final_max[scale_idx] = _neg_max_f32() + if uses_inst0: + final_max[scale_idx] = inst0_max + if uses_inst1: + final_max[scale_idx] = cute.math.max( + final_max[scale_idx], inst1_max, ftz=True + ) + exp_scale0[scale_idx] = Float32(0.0) + exp_scale1[scale_idx] = Float32(0.0) + if uses_inst0: + exp_scale0[scale_idx] = cute.math.exp2( + self.scale_softmax_log2 * (inst0_max - final_max[scale_idx]), + fastmath=True, + ) + if uses_inst1: + exp_scale1[scale_idx] = cute.math.exp2( + self.scale_softmax_log2 * (inst1_max - final_max[scale_idx]), + fastmath=True, + ) + reduced_sum[scale_idx] = Float32(0.0) + final_scale0[scale_idx] = Float32(0.0) + final_scale1[scale_idx] = Float32(0.0) + + final_sums = ffma2( + (exp_scale0[0], exp_scale0[1]), + (inst0_sum_arr[0], inst0_sum_arr[1]), + fmul2( + (exp_scale1[0], exp_scale1[1]), + (inst1_sum_arr[0], inst1_sum_arr[1]), + ), + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + final_sum[scale_idx] = final_sums[scale_idx] + + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + col_group_idx = warp_grp_thread_idx & Int32(0x3) + + # Reduce within each warp across lanes that share the same + # column group, then combine the 4 warp partials through + # shared memory. + final_sum_pair = self._warp_reduce_col_group_sum_pair( + (final_sum[0], final_sum[1]) + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + final_sum[scale_idx] = final_sum_pair[scale_idx] + + warp_store_base = warp_idx * Int32(8) + col_group_idx * Int32(2) + if lane_idx < Int32(4): + self._sum_scratch.store( + (final_sum[0], final_sum[1]), + warp_store_base, + alignment=8, + ) + prims.barrier_cta_sync( + self.sum_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + reduce_base = col_group_idx * Int32(2) + reduced_pair = self._sum_scratch.load( + reduce_base, + vector_size=2, + alignment=8, + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + reduced_sum[scale_idx] = reduced_pair[scale_idx] + for warp_offset in cutlass.range_constexpr(1, 4): + other_pair = self._sum_scratch.load( + reduce_base + warp_offset * Int32(8), + vector_size=2, + alignment=8, + ) + reduced_pair = fadd2( + (reduced_sum[0], reduced_sum[1]), + (other_pair[0], other_pair[1]), + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + reduced_sum[scale_idx] = reduced_pair[scale_idx] + + if cutlass.const_expr(not cfg.use_split_kv): + logical_h_k_idx, _ = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + reduced_sum[scale_idx] += _attention_sink_for_scale_idx( + cfg, + self.attention_sinks_ptr, + self.scale_softmax_log2, + final_max[scale_idx], + logical_h_k_idx, + self.h_r, + self.num_heads_kv, + logical_q_group_idx, + col_group_idx, + scale_idx, + ) + + base_addr0 = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, tail_o_stage_idx_0 + ) + base_addr1 = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, tail_o_stage_idx_1 + ) + output_pair_regs = cfg.num_fp16_output_regs + output_f32_regs = output_pair_regs * 2 + num_o_chunks = cfg.headdim // 64 + o0_vals, o1_vals = self._swaps_load_two_o_stage_chunks( + base_addr0, + base_addr1, + q_repeats=1, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + + if cutlass.const_expr(cfg.use_split_kv): + # Publish this CTA's partial output and statistics. Standalone + # reduction returns after publication; cluster peers reduce through + # DSMEM, while fused GMEM elects the final producer CTA. + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + q_row_offset = _q_tile_output_row_base(cfg, logical_q_group_idx) + logical_kv_idx = logical_b_idx * self.num_heads_kv + logical_h_k_idx + splits_kv = self._runtime_splits_kv(stage_info) + cta_idx_kv = _logical_cta_kv_idx(cfg, stage_info) + + if cutlass.const_expr(cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv): + for scale_idx in cutlass.range_constexpr(num_scale_groups): + norm_scale = Float32(1.0 / 448.0) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + norm_scale = self._separate_partial_norm_scale( + reduced_sum[scale_idx] + ) + final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] + final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] + + # Store normalized 16-bit O for the standalone reducer, or preserve + # unnormalized 16-bit O for fused GMEM/cluster reduction. + regs_partial_o = cutlass.Array( + Int32, + cfg.num_fp16_output_regs, + space=cutlass.AddressSpace.rmem, + ) + partial_scale0_pair = ( + (final_scale0[0], final_scale0[1]) + if cutlass.const_expr( + cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + ) + else (exp_scale0[0], exp_scale0[1]) + ) + partial_scale1_pair = ( + (final_scale1[0], final_scale1[1]) + if cutlass.const_expr( + cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + ) + else (exp_scale1[0], exp_scale1[1]) + ) + for reg_idx in cutlass.range_constexpr(cfg.num_fp16_output_regs): + reg_base = reg_idx * 2 + partial_pair = ffma2( + partial_scale0_pair, + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + partial_scale1_pair, + (o1_vals[reg_base], o1_vals[reg_base + 1]), + ), + ) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + regs_partial_o[reg_idx] = self._pack_separate_partial_o_pair( + partial_pair[0], partial_pair[1] + ) + elif cutlass.const_expr(cfg.use_bf16_output and not cfg.use_fp8_output): + regs_partial_o[reg_idx] = _pack_float2_to_bf16( + partial_pair[0], partial_pair[1] + ) + else: + regs_partial_o[reg_idx] = _pack_float2_to_fp16( + partial_pair[0], partial_pair[1] + ) + + num_partial_o_segments = max( + (cfg.tile_size_q * cfg.headdim * 2 + 2047) // 2048, + 1, + ) + if cutlass.const_expr(cfg.supports_cluster_smem_reduction): + self._publish_and_reduce_cluster_swaps_partials( + regs_partial_o, + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + cta_idx_kv, + splits_kv, + warp_grp_thread_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + output_pair_regs=output_pair_regs, + num_copy_segments=num_partial_o_segments, + num_scale_groups=2, + ) + return + + self._stage_and_copy_swaps_partial_o( + cfg, + regs_partial_o, + logical_kv_idx, + cta_idx_kv, + splits_kv, + q_row_offset, + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + output_pair_regs=output_pair_regs, + num_copy_segments=num_partial_o_segments, + enable_cluster=False, + full_prefix=False, + ) + + self._store_split_stats_from_scale_arrays( + cfg, + logical_kv_idx, + cta_idx_kv, + splits_kv, + q_row_offset, + logical_q_group_idx, + warp_idx, + lane_idx, + final_max, + reduced_sum, + num_scale_groups=2, + enable_cluster=False, + full_prefix=False, + ) + + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + # The standalone reducer consumes the partial O/stats. Do not + # also execute the fused in-kernel counter/cluster reduction. + return + + counter_q_groups = self._multi_cta_counter_q_groups(self.h_r) + counter_group_idx = logical_kv_idx * counter_q_groups + logical_q_group_idx + self._reduce_fused_gmem_partials( + logical_h_k_idx, + logical_b_idx, + logical_kv_idx, + q_row_offset, + logical_q_group_idx, + counter_group_idx, + splits_kv, + warp_grp_thread_idx, + ) + return + + for scale_idx in cutlass.range_constexpr(num_scale_groups): + # Direct final-output scales include both the reciprocal softmax + # denominator and public bmm2_scale. Split partials returned above + # use only the unnormalized ``exp_scale*`` arrays. + norm_scale = self.output_scale * self._safe_norm_rcp(reduced_sum[scale_idx]) + final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] + final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] + + regs_o = cutlass.Array( + Int32, + cfg.num_fp8_output_regs if cfg.use_fp8_output else cfg.num_fp16_output_regs, + space=cutlass.AddressSpace.rmem, + ) + # Form the final O registers from the two tail stages using the + # normalized instance scales. + if cutlass.const_expr(cfg.use_fp8_output): + for packed_idx in cutlass.range_constexpr(cfg.num_fp8_output_regs): + pair_idx0 = packed_idx * 2 + pair_idx1 = pair_idx0 + 1 + src_idx0 = pair_idx0 * 2 + src_idx1 = pair_idx1 * 2 + final_pair0 = ffma2( + (final_scale0[0], final_scale0[1]), + (o0_vals[src_idx0], o0_vals[src_idx0 + 1]), + fmul2( + (final_scale1[0], final_scale1[1]), + (o1_vals[src_idx0], o1_vals[src_idx0 + 1]), + ), + ) + final_pair1 = ffma2( + (final_scale0[0], final_scale0[1]), + (o0_vals[src_idx1], o0_vals[src_idx1 + 1]), + fmul2( + (final_scale1[0], final_scale1[1]), + (o1_vals[src_idx1], o1_vals[src_idx1 + 1]), + ), + ) + regs_o[packed_idx] = _pack_float4_to_fp8_e4m3( + final_pair0[0], + final_pair0[1], + final_pair1[0], + final_pair1[1], + ) + else: + for reg_idx in cutlass.range_constexpr(cfg.num_fp16_output_regs): + reg_base = reg_idx * 2 + final_pair = ffma2( + (final_scale0[0], final_scale0[1]), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + (final_scale1[0], final_scale1[1]), + (o1_vals[reg_base], o1_vals[reg_base + 1]), + ), + ) + if cutlass.const_expr(cfg.use_bf16_output): + regs_o[reg_idx] = _pack_float2_to_bf16(final_pair[0], final_pair[1]) + else: + regs_o[reg_idx] = _pack_float2_to_fp16(final_pair[0], final_pair[1]) + + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + if cutlass.const_expr(cfg.use_fp8_output): + # FP8 output uses transposed 8-bit staging before the final + # vectorized GMEM copy. + _store_transposed_smem8b( + self._smem_base_o_i32, + regs_o.data_ptr().load( + count=cfg.num_fp8_output_regs, + alignment=4, + ), + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.headdim, + cfg.num_fp8_output_regs, + ) + cute.arch.fence_view_async_shared() + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + _copy_transposed_smem8b_to_gmem( + self._smem_base_o_i32, + self.o_ptr, + cfg, + logical_h_k_idx, + logical_b_idx, + logical_q_group_idx, + self.h_r, + self.num_heads_kv, + self.seq_len_q, + self.q_token_offset, + warp_grp_thread_idx, + cfg.fp8_copy_can_use_full_tile_fast_path, + ) + if cutlass.const_expr(cfg.use_persistent_scheduler): + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + else: + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + self._stage_fp16_o_regs_to_smem( + cfg, + regs_o, + warp_grp_thread_idx, + warp_idx, + lane_idx, + output_pair_regs=cfg.num_fp16_output_regs, + ) + num_copy_segments = max( + (cfg.tile_size_q * cfg.headdim * cfg.o_dtype_bytes + 2047) // 2048, + 1, + ) + self._copy_staged_fp16_o_to_output( + cfg, + logical_h_k_idx, + logical_b_idx, + logical_q_group_idx, + warp_grp_thread_idx, + warp_idx, + lane_idx, + num_copy_segments=num_copy_segments, + ) + if cutlass.const_expr(cfg.use_persistent_scheduler): + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + @producer_work + @cute.jit + def correction_loop_epilogue( + self, + stage_info: StageInfo, + *, + o_stage_idx: Int32, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + ) -> None: + """Rescale the live O stage before later BMM2 waves accumulate into it.""" + # ProdWork: correction owns the live O TMEM stage after loop stats + # arrive. Rescale it in place so later PV waves accumulate in the + # updated online-softmax frame. + cfg = self.cfg + # Resolve the live TMEM O stage and default SwapsMmaAb column base. The + # KeepsMmaAb path overrides the base because O is allocated by TmemO. + task_cache = _decode_gen_task_cache(stage_info) + tmem_row_base = task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + o_base_col = 2 * cfg.tmem_s_cols + 2 * cfg.tmem_stats_cols + + if cutlass.const_expr(cfg.use_keeps_mma_ab): + # KeepsMmaAb has one softmax scale group for this path. Compute the + # rescale once, then apply it independently to every P-by-V + # head-dimension slice in TMEM. + o_base_col = self.tmem_o_ref._alloc.offset + corr_chunk_regs = cfg.keeps_loop_correction_chunk_regs + + old_max_0 = old_max_arr[0] + new_max_0 = new_max_arr[0] + scale_0, scale_is_identity = self._online_softmax_correction_scale( + old_max_0, + new_max_0, + ) + skip_correction = self._warp_can_skip_o_correction(scale_is_identity) + scale_pair = (scale_0, scale_0) + scaled_chunk = cutlass.Array( + Float32, corr_chunk_regs, space=cutlass.AddressSpace.rmem + ) + if not skip_correction: + # Each correction warp owns disjoint Keeps rows. When every + # row retains its running maximum, the scale is exactly one + # and the in-place TMEM rescale is a no-op. + for ( + head_dim_stage_offset, + keeps_o_ldst_offset, + num_corr_chunks, + ) in cfg.keeps_loop_correction_stage_layout: + stage_base_addr = tmem_row_base + Int32( + o_base_col + + o_stage_idx * cfg.tmem_o_stage_cols + + head_dim_stage_offset + ) + for chunk_idx in cutlass.range_constexpr(num_corr_chunks): + # Load one O chunk, multiply by + # exp(old_max-new_max), and store it before the next PV + # wave accumulates into this head-dimension slice. + chunk_col = chunk_idx * corr_chunk_regs + chunk_addr = stage_base_addr + Int32(chunk_col) + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(chunk_addr, Float32), + num=corr_chunk_regs, + offset=keeps_o_ldst_offset, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for reg_pair_idx in cutlass.range_constexpr( + corr_chunk_regs // 2 + ): + reg_base = reg_pair_idx * 2 + scaled_pair = fmul2( + scale_pair, + (loaded[reg_base], loaded[reg_base + 1]), + ) + scaled_chunk[reg_base] = scaled_pair[0] + scaled_chunk[reg_base + 1] = scaled_pair[1] + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(chunk_addr, Float32), + scaled_chunk.data_ptr().load( + count=corr_chunk_regs, alignment=4 + ), + offset=keeps_o_ldst_offset, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + if skip_correction: + # Preserve the correction task's TMEM ordering when this warp + # issues no correction transaction. + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + return + + if cutlass.const_expr(cfg.tile_size_q == 16): + q_repeats = max(cfg.tile_size_q // 8, 1) + num_scale_groups = cfg.num_softmax_scale_groups + output_pair_regs = cfg.num_fp16_output_regs + output_f32_regs = output_pair_regs * 2 + num_o_chunks = cfg.headdim // 64 + scale_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + lane_scales_are_identity = cutlass.Boolean(True) + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 2): + old_max_0 = old_max_arr[scale_base] + old_max_1 = old_max_arr[scale_base + 1] + new_max_0 = new_max_arr[scale_base] + new_max_1 = new_max_arr[scale_base + 1] + scale_0_is_identity = old_max_0 == new_max_0 + scale_1_is_identity = old_max_1 == new_max_1 + scale_0 = Float32(1.0) + scale_1 = Float32(1.0) + # Preserve the packed arithmetic used by the original path; + # only the exponentials and TMEM transaction are conditional. + max_diff_pair = fadd2( + (old_max_0, old_max_1), + (-new_max_0, -new_max_1), + ) + scale_pair = fmul2( + (self.scale_softmax_log2, self.scale_softmax_log2), + max_diff_pair, + ) + if not scale_0_is_identity: + scale_0 = cute.math.exp2(scale_pair[0], fastmath=True) + if not scale_1_is_identity: + scale_1 = cute.math.exp2(scale_pair[1], fastmath=True) + scale_vals[scale_base] = scale_0 + scale_vals[scale_base + 1] = scale_1 + lane_scales_are_identity = ( + lane_scales_are_identity & scale_0_is_identity & scale_1_is_identity + ) + + skip_correction = self._warp_can_skip_o_correction(lane_scales_are_identity) + + base_addr = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, o_stage_idx + ) + self._swaps_rescale_o_stage_in_tmem( + base_addr, + scale_vals, + skip_correction, + q_repeats=q_repeats, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + return + + if cutlass.const_expr(cfg.tile_size_q == 32): + # FlashInfer keeps Q32 straight-line: the warp vote and dynamic + # TMEM gate regress the wider Swaps pipeline on B200. + q_repeats = max(cfg.tile_size_q // 8, 1) + num_scale_groups = cfg.num_softmax_scale_groups + output_pair_regs = cfg.num_fp16_output_regs + output_f32_regs = output_pair_regs * 2 + num_o_chunks = cfg.headdim // 64 + scale_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 2): + old_max_0 = old_max_arr[scale_base] + old_max_1 = old_max_arr[scale_base + 1] + new_max_0 = new_max_arr[scale_base] + new_max_1 = new_max_arr[scale_base + 1] + max_diff_pair = fadd2((old_max_0, old_max_1), (-new_max_0, -new_max_1)) + scale_pair = fmul2( + (self.scale_softmax_log2, self.scale_softmax_log2), + max_diff_pair, + ) + scale_vals[scale_base] = cute.math.exp2(scale_pair[0], fastmath=True) + scale_vals[scale_base + 1] = cute.math.exp2( + scale_pair[1], fastmath=True + ) + + base_addr = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, o_stage_idx + ) + self._swaps_rescale_o_stage_in_tmem( + base_addr, + scale_vals, + cutlass.Boolean(False), + q_repeats=q_repeats, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + return + + # Tile-Q=8 loop ProdWork: rescale the current O stage in place when + # the row max changes. + old_max_0 = old_max_arr[0] + old_max_1 = old_max_arr[1] + new_max_0 = new_max_arr[0] + new_max_1 = new_max_arr[1] + scale_0_is_identity = old_max_0 == new_max_0 + scale_1_is_identity = old_max_1 == new_max_1 + scale_0 = Float32(1.0) + scale_1 = Float32(1.0) + max_diff_pair = fadd2((old_max_0, old_max_1), (-new_max_0, -new_max_1)) + scale_pair = fmul2( + (self.scale_softmax_log2, self.scale_softmax_log2), max_diff_pair + ) + if not scale_0_is_identity: + scale_0 = cute.math.exp2(scale_pair[0], fastmath=True) + if not scale_1_is_identity: + scale_1 = cute.math.exp2(scale_pair[1], fastmath=True) + skip_correction = self._warp_can_skip_o_correction( + scale_0_is_identity & scale_1_is_identity + ) + + scale_vals = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + scale_vals[0] = scale_0 + scale_vals[1] = scale_1 + base_addr = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, o_stage_idx + ) + output_pair_regs = cfg.num_fp16_output_regs + output_f32_regs = output_pair_regs * 2 + num_o_chunks = cfg.headdim // 64 + self._swaps_rescale_o_stage_in_tmem( + base_addr, + scale_vals, + skip_correction, + q_repeats=1, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + + @cute.jit + def _correction_tail_epilogue_impl( + self, + stage_info: StageInfo, + *, + scratch_stage: Int32 | None, + o_stage_idx: Int32, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + inst0_new_max_arr: cutlass.Array, + inst0_sum_arr: cutlass.Array, + inst1_new_max_arr: cutlass.Array, + inst1_sum_arr: cutlass.Array, + ) -> None: + """Normalize final O stages and store or publish the output tile.""" + _ = o_stage_idx + _ = old_max_arr + _ = new_max_arr + cfg = self.cfg + # ProdWork: consume final per-instruction softmax stats, normalize the + # tail O stages, then route the tile to direct output or the active + # split-KV reduction path. + # Final per-instance stats are explicit correction-task loop-carried + # values. Keeping them in the captured schedule avoids relying on + # private resource attributes across a persistent work-tile loop. + task_cache = _decode_gen_task_cache(stage_info) + tmem_row_base = task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + o_base_col = 2 * cfg.tmem_s_cols + 2 * cfg.tmem_stats_cols + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + + if cutlass.const_expr(cfg.use_keeps_mma_ab): + # KeepsMmaAb finalization is owned by the last active K/V instance. + # Earlier instances publish stats but do not write the final output. + if cutlass.const_expr( + (cfg.num_insts_kv == 1 and self.inst_id == 0) + or (cfg.num_insts_kv != 1 and self.inst_id == 1) + ): + # Derive the per-lane output row/column ownership before + # entering the common Keeps tail helper. + o_base_col = self.tmem_o_ref._alloc.offset + output_f32_regs = cfg.keeps_output_f32_regs + output_pair_regs = output_f32_regs // 2 + keeps_o_ldst_offset = cfg.headdim // 2 + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) + col_base = _keeps_col_base(cfg, lane_idx, cfg.headdim // 2) + if cutlass.const_expr(cfg.tile_size_kv == 256): + self._kv_tile_256_tail_epilogue( + stage_info, + scratch_stage=scratch_stage, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + tmem_row_base=tmem_row_base, + o_base_col=o_base_col, + warp_grp_thread_idx=warp_grp_thread_idx, + ) + return + self._keeps_tail_epilogue( + stage_info, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + tmem_row_base=tmem_row_base, + o_base_col=o_base_col, + output_pair_regs=output_pair_regs, + keeps_o_ldst_offset=keeps_o_ldst_offset, + row_idx=row_idx, + col_base=col_base, + warp_grp_thread_idx=warp_grp_thread_idx, + ) + return + + if cutlass.const_expr(self.inst_id != 1): + # SwapsMmaAb tail uses instance 1 to combine inst0/inst1 final O + # stages. Instance 0 exits after publishing its stats. + return + + if cutlass.const_expr(cfg.tile_size_q in (16, 32)): + # Wide tile-Q Swaps path has multiple softmax scale groups and may + # span several 64-column O chunks. + q_repeats = max(cfg.tile_size_q // 8, 1) + num_scale_groups = cfg.num_softmax_scale_groups + output_pair_regs = cfg.num_fp16_output_regs + output_f32_regs = output_pair_regs * 2 + num_o_chunks = cfg.headdim // 64 + self._swaps_wide_tail_epilogue( + stage_info, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + tmem_row_base=tmem_row_base, + o_base_col=o_base_col, + q_repeats=q_repeats, + num_scale_groups=num_scale_groups, + output_pair_regs=output_pair_regs, + output_f32_regs=output_f32_regs, + num_o_chunks=num_o_chunks, + ) + return + + # Tile-Q=8 Swaps path has exactly two scale groups and a compact tail + # helper specialized for that register layout. + self._swaps_q8_tail_epilogue( + stage_info, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + tmem_row_base=tmem_row_base, + o_base_col=o_base_col, + ) + return + + # Task Scheduling routes every non-constexpr work argument as a required + # data-flow token. Keep separate fixed/rotating entry points so only the + # latter consumes ``scratch_stage``; both still share the implementation. + @producer_work + @cute.jit + def correction_tail_epilogue( + self, + stage_info: StageInfo, + *, + o_stage_idx: Int32, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + inst0_new_max_arr: cutlass.Array, + inst0_sum_arr: cutlass.Array, + inst1_new_max_arr: cutlass.Array, + inst1_sum_arr: cutlass.Array, + ) -> None: + """Run the ordinary fixed-exchange tail epilogue.""" + self._correction_tail_epilogue_impl( + stage_info, + scratch_stage=None, + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + + @producer_work + @cute.jit + def correction_tail_epilogue_rotating_exchange( + self, + stage_info: StageInfo, + *, + scratch_stage: Int32, + o_stage_idx: Int32, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + inst0_new_max_arr: cutlass.Array, + inst0_sum_arr: cutlass.Array, + inst1_new_max_arr: cutlass.Array, + inst1_sum_arr: cutlass.Array, + ) -> None: + """Run persistent direct output in the stage named by its credit.""" + self._correction_tail_epilogue_impl( + stage_info, + scratch_stage=scratch_stage, + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py new file mode 100644 index 000000000000..fcf2cbe71730 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py @@ -0,0 +1,518 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``TmemOResource`` — TMEM O accumulator for BMM2. + +Producer (MmaTask): P × V MMA → O. Consumer (Correction): tracks which +O stage is ready (``o_stage_idx`` plus tail stage indices) so the in-place +rescale path can find the correct columns. +""" + +from dataclasses import dataclass +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass.experimental import primitives as prims + +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import SmemAllocation, TmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ..fmha_decode_constants import KV_INST0 +from ..fmha_decode_config import FmhaDecodeConfig +from ...tcgen05_compat import tcgen05_mma_ws +from .helpers_common import ( + Constexpr, + DecodeGenResourceBase, + DescriptorValue, + _TASK_CACHE_TMEM_BASE_OFFSET, + _decode_gen_task_cache, + _freeze_smem_descriptor, + _mma_k_step, + _mma_kind_for_qkv, +) + + +def _pv_mma_operand_contract_for_config( + cfg: FmhaDecodeConfig, +) -> tuple[bool, int, int, int, int]: + """Return ``(P-is-A, M, N, A-major, B-major)`` for BMM2.""" + active_head_dim = ( + cfg.headdim if cfg.head_dim_per_stage_kv == 0 else cfg.head_dim_kv_stage + ) + if cfg.use_keeps_mma_ab: + if cfg.tile_size_kv == 256: + # The WS 2x2 PV instruction exposes two spatial D128 partials as + # one physical KV256 operation. Correction merges those spatial + # halves after the two temporal decode streams are complete. + return True, cfg.tile_size_q, cfg.tile_size_kv, 0, 1 + return True, cfg.tile_size_q, active_head_dim, 0, 1 + return False, active_head_dim, cfg.tile_size_q, 1, 0 + + +@dataclass(kw_only=True) +class TmemOResource(DecodeGenResourceBase): + """TMEM O accumulator for BMM2. + + Producers run logical P x V MMA into staged TMEM O columns. Correction + consumers track which O stage is ready and rescale or output it. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("o_stage_idx", Int32, Int32(0), "Current O TMEM stage index."), + ( + "tail_o_stage_idx_0", + Int32, + Int32(0), + "O TMEM stage index for the first tail VP wave.", + ), + ( + "tail_o_stage_idx_1", + Int32, + Int32(1), + "O TMEM stage index for the second tail VP wave.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + scale_softmax_log2: Float32 = None + _alloc: Constexpr[TmemAllocation | None] = None + o_stage_idx: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + tail_o_stage_idx_0: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + tail_o_stage_idx_1: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def get_smem_requirements(self) -> list[SmemAllocation]: + """O accumulation lives in TMEM and needs no SMEM allocation.""" + return [] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Allocate staged TMEM columns for logical O accumulators.""" + cfg = self.cfg + if self._alloc is None: + self._alloc = TmemAllocation( + name=f"{self.name}", + num_columns=cfg.tmem_o_stage_cols * cfg.o_stages, + ) + return [self._alloc] + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1), + ) + @cute.jit + def init_stage_state(self, stage_info: StageInfo) -> tuple[Int32, Int32, Int32]: + """Initialize live and tail O-stage indices for correction.""" + # ConsAuxWork: seed correction's O-stage tracker before any PV MMA + # stages have been committed. + del stage_info + return Int32(0), Int32(0), Int32(1) + + @producer_work + @cute.jit + def vp_mma_loop( + self, + stage_info: StageInfo, + *, + v_desc_0: DescriptorValue, + v_desc_1: DescriptorValue, + p_desc_0: DescriptorValue, + p_desc_1: DescriptorValue, + p_tmem_addr_0: Int32, + p_tmem_addr_1: Int32, + inst_idx: Constexpr[int], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue loop PV MMA with overwrite on the first K/V tile only.""" + # ProdWork: loop PV accumulates into the live O stage; later correction + # work consumes this stage before another PV wave reuses it. + self._vp_mma( + stage_info, + v_desc_0=v_desc_0, + v_desc_1=v_desc_1, + p_desc_0=p_desc_0, + p_desc_1=p_desc_1, + p_tmem_addr_0=p_tmem_addr_0, + p_tmem_addr_1=p_tmem_addr_1, + initial_scale_d=stage_info.loop_offset != Int32(0), + inst_idx=inst_idx, + head_dim_stage_idx=head_dim_stage_idx, + ) + + @producer_work + @cute.jit + def vp_mma_tail( + self, + stage_info: StageInfo, + *, + v_desc_0: DescriptorValue, + v_desc_1: DescriptorValue, + p_desc_0: DescriptorValue, + p_desc_1: DescriptorValue, + p_tmem_addr_0: Int32, + p_tmem_addr_1: Int32, + inst_idx: Constexpr[int], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue tail PV MMA after all loop K tiles have been launched.""" + # ProdWork: tail PV publishes one of the final O stages consumed by the + # tail correction path. + self._vp_mma( + stage_info, + v_desc_0=v_desc_0, + v_desc_1=v_desc_1, + p_desc_0=p_desc_0, + p_desc_1=p_desc_1, + p_tmem_addr_0=p_tmem_addr_0, + p_tmem_addr_1=p_tmem_addr_1, + initial_scale_d=stage_info.loop_end != Int32(0), + inst_idx=inst_idx, + head_dim_stage_idx=head_dim_stage_idx, + ) + + @producer_work + @cute.jit + def vp_mma_loop_fragment( + self, + stage_info: StageInfo, + *, + v_desc: DescriptorValue, + p_tmem_addr: Int32, + fragment_idx: Constexpr[int], + ) -> None: + """Issue one K32 fragment of a KV256 loop PV tile.""" + self._vp_mma_fragment( + stage_info, + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=fragment_idx, + initial_scale_d=stage_info.loop_offset != Int32(0), + ) + + @producer_work + @cute.jit + def vp_mma_tail_fragment( + self, + stage_info: StageInfo, + *, + v_desc: DescriptorValue, + p_tmem_addr: Int32, + fragment_idx: Constexpr[int], + ) -> None: + """Issue one K32 fragment of the final KV256 PV tile.""" + self._vp_mma_fragment( + stage_info, + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=fragment_idx, + initial_scale_d=stage_info.loop_end != Int32(0), + ) + + @cute.jit + def _vp_mma_fragment( + self, + stage_info: StageInfo, + *, + v_desc: DescriptorValue, + p_tmem_addr: Int32, + fragment_idx: Constexpr[int], + initial_scale_d, + ) -> None: + """Issue the two WS MMA steps covered by one KV256 P fragment. + + ``p_tmem_addr`` is already the base of the fragment selected by + ``wait_p_fragment``. Only the two local K-step offsets are added here; + ``fragment_idx`` must not be applied to the TMEM address again. + """ + cfg = self.cfg + assert cfg.tile_size_kv == 256 and cfg.uses_two_inst_tmem_p + v_desc = _freeze_smem_descriptor(v_desc) + + task_cache = _decode_gen_task_cache(stage_info) + tmem_col = prims.make_tmem_ptr( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + stage_info.stage_idx * cfg.tmem_o_cols, + Float32, + ) + _, mma_m, mma_n, a_major, b_major = _pv_mma_operand_contract_for_config(cfg) + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=cfg.q_dtype, + b_dtype=cfg.q_dtype, + a_major=a_major, + b_major=b_major, + n_dim=mma_n, + m_dim=mma_m, + ) + first_k_step = fragment_idx * 2 + if prims.elect_sync(): + for local_k_step in cutlass.range_constexpr(2): + k_step = first_k_step + local_k_step + p_operand = prims.make_tmem_ptr( + p_tmem_addr + Int32(local_k_step * 8), Int32 + ) + iter_v_desc = v_desc + Int32( + (k_step // 4) * cfg.headdim * 16 + (k_step % 4) * 128 + ) + tcgen05_mma_ws( + _mma_kind_for_qkv(cfg), + tmem_col, + p_operand, + iter_v_desc, + idesc, + initial_scale_d or fragment_idx != 0 or local_k_step != 0, + ) + + @cute.jit + def _vp_mma( + self, + stage_info: StageInfo, + *, + v_desc_0: DescriptorValue, + v_desc_1: DescriptorValue, + p_desc_0: DescriptorValue, + p_desc_1: DescriptorValue, + p_tmem_addr_0: Int32, + p_tmem_addr_1: Int32, + initial_scale_d, + inst_idx: Constexpr[int], + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue one non-fragmented PV MMA wave for loop or tail work.""" + cfg = self.cfg + # KV256 is always streamed through _vp_mma_fragment so no full 128-P + # row is kept live. Keep this routine as the sole generic PV path. + assert cfg.tile_size_kv != 256 + # Select the descriptor pair for this BMM2 call. With staged head + # dimensions, consecutive calls belong to the same KV instance and + # different head-dim slices. + if cutlass.const_expr(inst_idx == KV_INST0): + v_desc = v_desc_0 + p_desc = p_desc_0 + p_tmem_addr = p_tmem_addr_0 + else: + v_desc = v_desc_1 + p_desc = p_desc_1 + p_tmem_addr = p_tmem_addr_1 + v_desc = _freeze_smem_descriptor(v_desc) + if cutlass.const_expr(not cfg.uses_tmem_p): + p_desc = _freeze_smem_descriptor(p_desc) + task_cache = _decode_gen_task_cache(stage_info) + p_is_a, mma_m, mma_n, a_major, b_major = _pv_mma_operand_contract_for_config( + cfg + ) + + if cutlass.const_expr(cfg.head_dim_per_stage_kv == 0): + # Full-headDim path: all V columns for this O stage live in one + # contiguous TMEM stage selected by the TS pipeline stage index. + base_addr = task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32( + self._alloc.offset + stage_info.stage_idx * cfg.tmem_o_cols + ) + tmem_col = cutlass.inttoptr( + base_addr, + 6, + Float32, + ) + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=cfg.q_dtype, + b_dtype=cfg.q_dtype, + a_major=a_major, + b_major=b_major, + n_dim=mma_n, + m_dim=mma_m, + ) + + if prims.elect_sync(): + # Accumulate into O after the first K/V tile for this output + # stage; the first wave overwrites the TMEM O stage. + scale_d = initial_scale_d + pv_k_steps = cfg.tile_size_kv // _mma_k_step(cfg) + for ki in cutlass.range_constexpr(pv_k_steps): + # Keeps computes P x V (A=P, B=V); Swaps computes the + # transposed V^T x P^T tile (A=V, B=P). + if cutlass.const_expr(cfg.uses_tmem_p): + # TMEM P stores two 16-bit values per column (or four + # FP8 values), so each 16-wide MMA-K step advances by + # the corresponding packed-column count. + p_cols_per_k_step = _mma_k_step(cfg) * cfg.q_dtype_bytes // 4 + p_operand = prims.make_tmem_ptr( + p_tmem_addr + Int32(ki * p_cols_per_k_step), + Int32, + ) + else: + p_operand = p_desc + if cutlass.const_expr(p_is_a): + a_desc, b_desc = p_operand, v_desc + else: + a_desc, b_desc = v_desc, p_operand + prims.tcgen05_mma( + _mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + a_desc, + b_desc, + idesc, + scale_d, + ) + scale_d = True + if cutlass.const_expr(ki + 1 < pv_k_steps): + # Advance V and P descriptors to the next MMA-K + # slice, including the 16-bit 128-token jump across + # split SMEM rows. + v_desc = v_desc + Int32( + (cfg.headdim * 2) if cfg.use_fp8_qkv else 128 + ) + if cutlass.const_expr(not cfg.uses_tmem_p): + if cutlass.const_expr( + not cfg.use_fp8_qkv + and cfg.tile_size_kv == 128 + and ki == 3 + ): + if cutlass.const_expr(cfg.tile_size_q >= 16): + p_desc = p_desc + Int32(8 * cfg.tile_size_q - 6) + else: + p_desc = p_desc + Int32(58) + else: + p_desc = p_desc + Int32(2) + else: + # Staged-headDim path: each call owns one head-dim slice of the same + # logical O stage. The TMEM offset selects that slice before MMA. + head_dim_stage_tmem_offset = cfg.pv_head_dim_stage_tmem_offset( + head_dim_stage_idx + ) + base_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32( + self._alloc.offset + stage_info.stage_idx * cfg.tmem_o_stage_cols + ) + + Int32(head_dim_stage_tmem_offset) + ) + tmem_col = cutlass.inttoptr( + base_addr, + 6, + Float32, + ) + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=cfg.q_dtype, + b_dtype=cfg.q_dtype, + a_major=a_major, + b_major=b_major, + n_dim=mma_n, + m_dim=mma_m, + ) + + if prims.elect_sync(): + # Issue one MMA per K step. The first wave may overwrite the + # slice, while later loop/tail waves accumulate into it. + scale_d = initial_scale_d + for ki in cutlass.range_constexpr(cfg.tile_size_kv // _mma_k_step(cfg)): + if cutlass.const_expr(cfg.uses_tmem_p): + p_cols_per_k_step = _mma_k_step(cfg) * cfg.q_dtype_bytes // 4 + p_operand = prims.make_tmem_ptr( + p_tmem_addr + Int32(ki * p_cols_per_k_step), + Int32, + ) + else: + p_operand = p_desc + if cutlass.const_expr(p_is_a): + a_desc, b_desc = p_operand, v_desc + else: + a_desc, b_desc = v_desc, p_operand + prims.tcgen05_mma( + _mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + a_desc, + b_desc, + idesc, + scale_d, + ) + scale_d = True + if cutlass.const_expr( + ki + 1 < cfg.tile_size_kv // _mma_k_step(cfg) + ): + # Advance V and P to the next MMA-K slice inside the + # staged head-dim tile. + v_desc = v_desc + Int32( + (cfg.head_dim_kv_stage * 2) if cfg.use_fp8_qkv else 128 + ) + if cutlass.const_expr(not cfg.uses_tmem_p): + if cutlass.const_expr( + not cfg.use_fp8_qkv + and cfg.tile_size_kv == 128 + and ki == 3 + ): + if cutlass.const_expr(cfg.tile_size_q >= 16): + p_desc = p_desc + Int32(8 * cfg.tile_size_q - 6) + else: + p_desc = p_desc + Int32(58) + else: + p_desc = p_desc + Int32(2) + + @cute.jit + def _return_o_stage_state( + self, + o_stage_idx: Int32, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + ) -> tuple[object, object, object]: + """Return the O-stage task-local tuple in scheduler order.""" + return o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1 + + @consumer_work(returns=(o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1)) + @cute.jit + def update_o_stage_loop( + self, + stage_info: StageInfo, + *, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + ) -> tuple[object, object, object]: + """Record the loop O stage that correction will rescale in place.""" + # ConsWork: loop correction records which O stage became available for + # in-place rescale. + o_stage_idx = stage_info.stage_idx + return self._return_o_stage_state( + o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1 + ) + + @consumer_work(returns=(o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1)) + @cute.jit + def update_o_stage_tail( + self, + stage_info: StageInfo, + *, + tail_o_stage_idx_0: Int32, + tail_o_stage_idx_1: Int32, + inst_idx: Constexpr[int], + ) -> tuple[object, object, object]: + """Record the two final O stages consumed by tail correction.""" + # ConsWork: tail correction consumes one completed O stage per K/V + # instance and records which TMEM columns hold inst0/inst1. + o_stage_idx = stage_info.stage_idx + if cutlass.const_expr(inst_idx == KV_INST0): + tail_o_stage_idx_0 = stage_info.stage_idx + else: + tail_o_stage_idx_1 = stage_info.stage_idx + return self._return_o_stage_state( + o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1 + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py new file mode 100644 index 000000000000..703134a5b61d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py @@ -0,0 +1,2615 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``TmemSResource`` — BMM1 accumulator / softmax input. + +Producer: QK MMA → S in TMEM. Consumer: load S to registers, maintain +running row max/sum, apply optional causal / sliding-window / sink masks, +publish softmax stats for ``SmemPResource``. +""" + +from dataclasses import dataclass +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Boolean, Float32, Int32, Uint32 +from cutlass.experimental import primitives as prims + +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...._block_sparse.common import _MAX_KV_ATOM_SIZE +from ...._block_sparse.prepared import _PREPARED_ROUTE_IS_FULL_FLAG +from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig +from ..fmha_decode_constants import KV_TILE_256_RESCALE_THRESHOLD_LOG2 +from ...tcgen05_compat import tcgen05_mma_ws +from ...placeholder_helpers import ( + _placeholder_local_array, + _placeholder_smem_array, +) +from .helpers_common import ( + Constexpr, + DecodeGenResourceBase, + ResourceVars, + ffma2, + _TASK_CACHE_LANE_IDX, + _TASK_CACHE_KV_RAW_TILE_BASE, + _TASK_CACHE_KV_VALID_TILE_END, + _TASK_CACHE_KV_WINDOW_START, + _TASK_CACHE_TMEM_BASE_OFFSET, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _TASK_CACHE_WARP_IDX, + _clamp_valid_tile_idx, + _decode_gen_task_cache, + _freeze_smem_descriptor, + _is_last_loop_iteration, + _keeps_col_base, + _keeps_row_idx, + _keeps_score_col, + _keeps_tcgen05_ld, + _keeps_tcgen05_st, + _logical_q_group_idx, + _mma_k_step, + _mma_kind_for_qkv, + _neg_max_f32, + _softmax_scale_pair_width, + _q_row_is_valid_for_seq, + _q_row_token_and_local_head, + _q_group_token_base, + _softmax_tile_idx, +) +from .smem_block_sparse_metadata import ( + _SOFTMAX_TOKEN_MASK_IS_FULL_FLAG, + _swaps_forwards_packed_route_full, +) +from .helpers_kv_tile_idx import ( + _kv_tile_is_fully_unmasked_for_q_group, + _load_runtime_seq_len_kv, + _num_skipped_kv_tiles, + _runtime_clamp_valid_tile_idx, + _runtime_split_kv_global_tile_idx, + _runtime_total_kv_tiles, + _sliding_window_start_idx, + _static_split_kv_global_tile_idx, +) +from .helpers_softmax import ( + _float_to_u32_for_atomic_max, + _init_softmax_scratch_u32, + _smem_atomic_max_u32, + _u32_to_float_for_atomic_max, + _wspro_reduce_max4, +) + +# A block-sparse route often changes the exact row maximum without changing it +# enough to justify rescaling the live O tile. Keeping the prior anchor within +# this bound makes the correction scale exactly one and bounds FP16/BF16 P by +# 2**8. As in the FlashInfer/TRT-LLM policy, this assumes normal model logits +# rather than adversarial values outside the qualified probability bound. +_BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2 = 8.0 + + +def _swaps_uses_origin0_k32_full_guard(cfg: FmhaDecodeConfig) -> bool: + """Whether one staged origin can prove this warp's K32 slice valid.""" + + return ( + cfg.kv_block_size >= 32 + and not cfg.use_kv_valid_bits + and not cfg.uses_uniform_causal_mask + and not cfg.uses_per_row_causal_mask + ) + + +def _swaps_token_word_covers_kv_tail(cfg: FmhaDecodeConfig) -> bool: + """Whether SWAP's prepared token word covers the logical KV tail.""" + + return ( + cfg.use_kv_valid_bits + and not cfg.uses_uniform_causal_mask + and not cfg.uses_per_row_causal_mask + ) + + +def _swaps_uses_token_only_score_validity(cfg: FmhaDecodeConfig) -> bool: + """Whether prepared token words replace SWAP's atom-origin guard.""" + + return ( + cfg.use_block_sparse + and _swaps_token_word_covers_kv_tail(cfg) + and cfg.tile_size_q < 64 + and cfg.use_persistent_scheduler + and (cfg.kv_block_size >= 16 or cfg.use_parallel_sparse_kv_loads) + ) + + +@cute.jit +def _can_skip_sparse_keeps_structural_mask( + q_row_is_valid: Boolean, + origin0: Int32, + origin1: Int32, + valid0: Int32, + valid1: Int32, + seq_len_kv: Int32, + causal_end: Int32, + *, + apply_causal_mask: cutlass.Constexpr[bool], +) -> Boolean: + """Return whether one Keeps row needs no Q/tail/causal predicate. + + Token-bit masking is independent. Comparing against the last complete + KV64 origin avoids overflowing an origin near the Int32 upper bound. + """ + + fragment_size = Int32(_MAX_KV_ATOM_SIZE) + last_complete_origin = seq_len_kv - fragment_size + can_skip = Boolean( + q_row_is_valid + and valid0 != Int32(0) + and valid1 != Int32(0) + and origin0 <= last_complete_origin + and origin1 <= last_complete_origin + ) + if cutlass.const_expr(apply_causal_mask): + last_causal_origin = causal_end - fragment_size + can_skip = Boolean( + can_skip and origin0 <= last_causal_origin and origin1 <= last_causal_origin + ) + return can_skip + + +@cute.jit +def _sparse_k32_effective_keep_word( + q_row_is_valid: Boolean, + fragment_origin: Int32, + fragment_valid: Int32, + token_word: Uint32, + seq_len_kv: Int32, + causal_end: Int32, + *, + apply_causal_mask: cutlass.Constexpr[bool], + apply_token_mask: cutlass.Constexpr[bool], +) -> Uint32: + """Fold route, KV-tail, causal, and token predicates for one K32 fragment.""" + + keep_word = Uint32(0) + if q_row_is_valid and fragment_valid != Int32(0): + visible_end = seq_len_kv + if cutlass.const_expr(apply_causal_mask): + visible_end = cute.math.min(visible_end, causal_end) + visible_tokens = visible_end - fragment_origin + if visible_tokens >= Int32(32): + keep_word = Uint32(0xFFFFFFFF) + if cutlass.const_expr(apply_token_mask): + keep_word = token_word + elif visible_tokens > Int32(0): + # Keep the shift strictly below 32; shifting a 32-bit value by its + # width is undefined in PTX and LLVM. + keep_word = (Uint32(1) << visible_tokens) - Uint32(1) + if cutlass.const_expr(apply_token_mask): + keep_word = keep_word & token_word + return keep_word + + +def _qk_mma_operand_contract_for_config( + cfg: FmhaDecodeConfig, +) -> tuple[bool, int, int]: + """Return ``(Q-is-A, M, N)`` for the selected BMM1 orientation.""" + if cfg.use_keeps_mma_ab: + return True, cfg.tile_size_q, cfg.tile_size_kv + return False, cfg.tile_size_kv, cfg.tile_size_q + + +@dataclass(kw_only=True) +class TmemSResource(DecodeGenResourceBase): + """TMEM score resource for BMM1 and softmax. + + Producers run K x Q^T MMA into TMEM S. Consumers load S into registers, + compute running row maxima, and carry the softmax state used by P and + correction. + """ + + _rts_internal_consumer_var_names: ClassVar[tuple[str, ...]] = ( + "old_max_arr", + "sum_arr", + "new_max_arr", + "s_arr", + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "old_max_arr", + cutlass.Array, + None, + "Previous softmax anchor (normally the running row maximum).", + ), + ("sum_arr", cutlass.Array, None, "Running softmax denominator."), + ( + "new_max_arr", + cutlass.Array, + None, + "Current softmax anchor (normally the running row maximum).", + ), + ("s_arr", cutlass.Array, None, "Loaded S scores for the current tile."), + ) + inst_id: Constexpr[int] = 0 + cfg: Constexpr[FmhaDecodeConfig] = None + scale_softmax_log2: Float32 = None + seqlens_kv: cute.Pointer | None = None + max_seq_len_kv: Int32 = None + seq_len_q: Int32 = None + h_r: Int32 | None = None + q_group_idx: Int32 | None = None + q_ref: Constexpr[MemoryResource | None] = None + _p_local_sum_arr: cutlass.Array | None = None + _global_sum_arr: cutlass.Array | None = None + _alloc: Constexpr[TmemAllocation | None] = None + sync_barrier_id: Constexpr[int] = 0 + _scratch_alloc: Constexpr[SmemAllocation | None] = None + _softmax_scratch_u32: cutlass.Array = None + old_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + sum_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + new_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + s_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def _init_placeholder_state(self) -> None: + """Create placeholder register and scratch state for softmax.""" + num_scale_groups = self.cfg.num_softmax_scale_groups + num_s_regs = self.cfg.softmax_score_fragment_regs + self.old_max_arr.default = _placeholder_local_array(Float32, num_scale_groups) + self.sum_arr.default = _placeholder_local_array(Float32, num_scale_groups) + self.new_max_arr.default = _placeholder_local_array(Float32, num_scale_groups) + self.s_arr.default = _placeholder_local_array(Float32, num_s_regs) + self._p_local_sum_arr = _placeholder_local_array(Float32, num_scale_groups) + self._global_sum_arr = _placeholder_local_array(Float32, num_scale_groups) + scratch_entries = ( + 4 * num_scale_groups if self.cfg.use_keeps_mma_ab else self.cfg.tile_size_q + ) + self._softmax_scratch_u32 = _placeholder_smem_array(Uint32, scratch_entries) + + @cute.jit + def store_p_local_sum(self, scale_idx: int, value: Float32) -> None: + """Publish the P producer's local denominator contribution.""" + self._p_local_sum_arr[scale_idx] = value + + @cute.jit + def load_p_local_sum(self, scale_idx: int) -> Float32: + """Load the local denominator published with the current P tile.""" + return self._p_local_sum_arr[scale_idx] + + @cute.jit + def store_global_sum(self, scale_idx: int, value: Float32) -> None: + """Publish the FP8 cross-warp denominator correction.""" + self._global_sum_arr[scale_idx] = value + + @cute.jit + def load_global_sum(self, scale_idx: int) -> Float32: + """Load the FP8 denominator after cross-warp correction.""" + return self._global_sum_arr[scale_idx] + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate softmax scratch used for CTA-wide max reductions.""" + if self._scratch_alloc is None: + scratch_entries = ( + 4 * self.cfg.num_softmax_scale_groups + if self.cfg.use_keeps_mma_ab + else self.cfg.tile_size_q + ) + self._scratch_alloc = SmemAllocation( + name=f"{self.name}_softmaxScratch", + size_bytes=scratch_entries * 4, + alignment=16, + ) + return [self._scratch_alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Allocate TMEM S score columns for QK MMA output.""" + if self._alloc is None: + num_stages = ( + self.pipeline_config.num_stages + if ( + self.pipeline_config is not None + and self.cfg.use_keeps_mma_ab + and self.cfg.num_insts_kv == 1 + ) + else 1 + ) + self._alloc = TmemAllocation( + name=f"{self.name}", + num_columns=self.cfg.tmem_s_cols * num_stages, + ) + return [self._alloc] + + @cute.jit + def _q_desc_for_head_dim_stage( + self, + q_desc: prims.Tcgen05SmemDesc, + head_dim_stage_idx: Constexpr[int], + ) -> tuple[prims.Tcgen05SmemDesc, Constexpr[int]]: + """Select the staged-Q descriptor slice consumed by this BMM1 call.""" + cfg = self.cfg + if cutlass.const_expr(cfg.head_dim_per_stage_kv != 0): + q_stage_offset = Int32( + head_dim_stage_idx + * cfg.head_dim_kv_stage + * cfg.tile_size_q + * cfg.q_dtype_bytes + // 16 + ) + q_desc = q_desc + q_stage_offset + return q_desc, head_dim_stage_idx + + @cute.jit + def _advance_qk_descs_after_mma_k( + self, + k_desc: prims.Tcgen05SmemDesc, + q_desc: prims.Tcgen05SmemDesc, + *, + crosses_64b_chunk: Constexpr[bool], + ) -> tuple[prims.Tcgen05SmemDesc, prims.Tcgen05SmemDesc]: + """Advance K/Q descriptors to the next 16-wide MMA-K slice. + + 16-bit layouts are staged as 64-column chunks. Crossing that chunk + boundary uses the large descriptor jump; all other steps advance by one + MMA-K slice. + """ + cfg = self.cfg + if cutlass.const_expr(not cfg.use_fp8_qkv and crosses_64b_chunk): + k_desc = k_desc + Int32(8 * cfg.tile_size_kv - 6) + if cutlass.const_expr(cfg.tile_size_q >= 16): + q_desc = q_desc + Int32(8 * cfg.tile_size_q - 6) + else: + q_desc = q_desc + Int32(58) + else: + k_desc = k_desc + Int32(2) + q_desc = q_desc + Int32(2) + return k_desc, q_desc + + @cute.jit + def _stage_slot_offset_from_slot(self, slot: Int32) -> Int32: + """Map a logical S pipeline slot to a TMEM column offset.""" + cfg = self.cfg + if cutlass.const_expr(not cfg.use_keeps_mma_ab or cfg.num_insts_kv != 1): + return Int32(0) + return slot * Int32(cfg.tmem_s_cols) + + @cute.jit + def _qk_head_stage_slot_offset(self, stage_info: StageInfo) -> Int32: + """Return the TMEM S slot used by HEAD QK MMA.""" + if cutlass.const_expr(stage_info.stage_idx is not None): + return self._stage_slot_offset_from_slot(Int32(stage_info.stage_idx)) + return self._stage_slot_offset_from_slot(Int32(0)) + + @cute.jit + def _qk_loop_stage_slot_offset(self, stage_info: StageInfo) -> Int32: + """Return the producer TMEM S slot for a LOOP QK MMA wave.""" + if cutlass.const_expr(stage_info.stage_idx is not None): + return self._stage_slot_offset_from_slot(Int32(stage_info.stage_idx)) + return self._stage_slot_offset_from_slot( + (stage_info.loop_offset + Int32(1)) % Int32(2) + ) + + @cute.jit + def _softmax_loop_stage_slot_offset(self, stage_info: StageInfo) -> Int32: + """Return the consumer TMEM S slot read by LOOP softmax.""" + if cutlass.const_expr(stage_info.stage_idx is not None): + return self._stage_slot_offset_from_slot(Int32(stage_info.stage_idx)) + return self._stage_slot_offset_from_slot(stage_info.loop_offset % Int32(2)) + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind softmax scratch and initialize running max/sum state.""" + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self._scratch_alloc is not None + ): + # Shared scratch holds encoded per-scale-group maxima for the + # four softmax warps before they are reduced back to registers. + scratch_ptr = context.smem_base.data_ptr() + self._scratch_alloc.offset + self._softmax_scratch_u32 = cutlass.Array( + scratch_ptr, + dtype=Uint32, + shape=( + ( + 4 * self.cfg.num_softmax_scale_groups + if self.cfg.use_keeps_mma_ab + else self.cfg.tile_size_q + ), + ), + addrspace=3, + ) + + num_scale_groups = self.cfg.num_softmax_scale_groups + num_s_regs = self.cfg.softmax_score_fragment_regs + # Cross-resource mutable arrays are stored as instance attributes, not + # consumer vars, so SmemP and TmemSoftmaxGlobal can update them in + # place between schedule steps. + self._p_local_sum_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + self._global_sum_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + result = { + "old_max_arr": cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ), + "sum_arr": cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ), + "new_max_arr": cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ), + "s_arr": cutlass.Array( + Float32, num_s_regs, space=cutlass.AddressSpace.rmem + ), + } + for idx in cutlass.range_constexpr(num_scale_groups): + # Initialize running max/sum state for the first K/V tile. + result["old_max_arr"][idx] = _neg_max_f32() + result["sum_arr"][idx] = Float32(0.0) + result["new_max_arr"][idx] = _neg_max_f32() + self._p_local_sum_arr[idx] = Float32(0.0) + self._global_sum_arr[idx] = Float32(0.0) + for idx in cutlass.range_constexpr(num_s_regs): + # Invalid lanes start at -inf so masks and empty tiles naturally + # contribute zero probability. + result["s_arr"][idx] = _neg_max_f32() + return result + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Create per-work-tile softmax state for persistent scheduling.""" + _ = context + # Reinitialize the softmax state for each persistent-scheduler work + # tile while preserving the resource-level scratch allocation. + num_scale_groups = self.cfg.num_softmax_scale_groups + num_s_regs = self.cfg.softmax_score_fragment_regs + result = { + "old_max_arr": cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ), + "sum_arr": cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ), + "new_max_arr": cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ), + "s_arr": cutlass.Array( + Float32, num_s_regs, space=cutlass.AddressSpace.rmem + ), + } + for idx in cutlass.range_constexpr(num_scale_groups): + result["old_max_arr"][idx] = _neg_max_f32() + result["sum_arr"][idx] = Float32(0.0) + result["new_max_arr"][idx] = _neg_max_f32() + self._p_local_sum_arr[idx] = Float32(0.0) + self._global_sum_arr[idx] = Float32(0.0) + for idx in cutlass.range_constexpr(num_s_regs): + result["s_arr"][idx] = _neg_max_f32() + return result + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(old_max_arr, sum_arr, new_max_arr, s_arr), + ) + @cute.jit + def init_softmax_state( + self, stage_info: StageInfo + ) -> tuple[cutlass.Array, cutlass.Array, cutlass.Array, cutlass.Array]: + """Initialize and return the softmax task-local state tuple.""" + # ConsAuxWork: seed the running max/sum state and local S buffers before + # the softmax task consumes any QK score tile. + result = self._create_initial_task_locals(stage_info.context) + return ( + result["old_max_arr"], + result["sum_arr"], + result["new_max_arr"], + result["s_arr"], + ) + + @producer_work + @cute.jit + def qk_mma_head( + self, + stage_info: StageInfo, + *, + q_desc: prims.Tcgen05SmemDesc, + kv_desc: prims.Tcgen05SmemDesc, + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue HEAD QK MMA into the initial S slot.""" + # ProdWork: HEAD produces the first score tile, overwriting the initial + # S stage before any loop softmax work has consumed it. + self._qk_mma( + stage_info, + q_desc=q_desc, + kv_desc=kv_desc, + stage_slot_offset=self._qk_head_stage_slot_offset(stage_info), + head_dim_stage_idx=head_dim_stage_idx, + ) + + @producer_work + @cute.jit + def qk_mma_loop( + self, + stage_info: StageInfo, + *, + q_desc: prims.Tcgen05SmemDesc, + kv_desc: prims.Tcgen05SmemDesc, + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue LOOP QK MMA into the next producer S slot.""" + # ProdWork: LOOP produces the next score tile into the stage that + # softmax will consume for this steady-state iteration. + self._qk_mma( + stage_info, + q_desc=q_desc, + kv_desc=kv_desc, + stage_slot_offset=self._qk_loop_stage_slot_offset(stage_info), + head_dim_stage_idx=head_dim_stage_idx, + ) + + @producer_work + @cute.jit + def qk_mma_head_from_q_ref( + self, + stage_info: StageInfo, + *, + kv_desc: prims.Tcgen05SmemDesc, + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue guarded persistent HEAD QK without a routed Q descriptor.""" + assert self.q_ref is not None + self._qk_mma( + stage_info, + q_desc=self.q_ref.current_consumer_q_desc(), + kv_desc=kv_desc, + stage_slot_offset=self._qk_head_stage_slot_offset(stage_info), + head_dim_stage_idx=head_dim_stage_idx, + ) + + @producer_work + @cute.jit + def qk_mma_loop_from_q_ref( + self, + stage_info: StageInfo, + *, + kv_desc: prims.Tcgen05SmemDesc, + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue guarded persistent LOOP QK without a routed Q descriptor.""" + assert self.q_ref is not None + self._qk_mma( + stage_info, + q_desc=self.q_ref.current_consumer_q_desc(), + kv_desc=kv_desc, + stage_slot_offset=self._qk_loop_stage_slot_offset(stage_info), + head_dim_stage_idx=head_dim_stage_idx, + ) + + @cute.jit + def _qk_mma( + self, + stage_info: StageInfo, + *, + q_desc: prims.Tcgen05SmemDesc, + kv_desc: prims.Tcgen05SmemDesc, + stage_slot_offset: Int32, + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Issue BMM1 with the selected Keeps/Swaps MMA orientation. + + Issues all 16-wide MMA-K slices for the current staged head-dim tile. + Descriptor stage selection and per-slice jumps are centralized in the + local descriptor helpers below. + """ + cfg = self.cfg + k_desc = _freeze_smem_descriptor(kv_desc) + q_desc = _freeze_smem_descriptor(q_desc) + q_desc, head_dim_stage_idx = self._q_desc_for_head_dim_stage( + q_desc, head_dim_stage_idx + ) + + # TMEM destination: addrspace-6 pointer from base + alloc offset. + # cutlass's tcgen05_alloc returns the base in tmem_ptr_i32, so add the + # per-resource column offset before issuing MMA. + task_cache = _decode_gen_task_cache(stage_info) + tmem_col = prims.make_tmem_ptr( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + stage_slot_offset, + Float32, + ) + + q_is_a, mma_m, mma_n = _qk_mma_operand_contract_for_config(cfg) + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=cfg.q_dtype, + b_dtype=cfg.q_dtype, + n_dim=mma_n, + m_dim=mma_m, + ) + + if cutlass.const_expr(cfg.head_dim_per_stage_kv == 0): + if prims.elect_sync(): + scale_d = False + for ki in cutlass.range_constexpr(cfg.headdim // _mma_k_step(cfg)): + # Keeps computes Q x K^T (A=Q, B=K); Swaps computes the + # transposed K x Q^T tile (A=K, B=Q). The first + # instruction overwrites S and later slices accumulate. + if cutlass.const_expr(q_is_a): + a_desc, b_desc = q_desc, k_desc + else: + a_desc, b_desc = k_desc, q_desc + if cutlass.const_expr(cfg.tile_size_kv == 256): + tcgen05_mma_ws( + _mma_kind_for_qkv(cfg), + tmem_col, + a_desc, + b_desc, + idesc, + scale_d, + ) + else: + prims.tcgen05_mma( + _mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + a_desc, + b_desc, + idesc, + scale_d, + ) + scale_d = True + if cutlass.const_expr(ki + 1 < cfg.headdim // _mma_k_step(cfg)): + k_desc, q_desc = self._advance_qk_descs_after_mma_k( + k_desc, + q_desc, + crosses_64b_chunk=cfg.headdim == 128 and ki == 3, + ) + else: + mma_k_steps = cfg.head_dim_kv_stage // _mma_k_step(cfg) + if prims.elect_sync(): + # Peel the first MMA so overwrite-vs-accumulate remains a + # compile-time value rather than loop-carried state. + if cutlass.const_expr(q_is_a): + first_a_desc, first_b_desc = q_desc, k_desc + else: + first_a_desc, first_b_desc = k_desc, q_desc + prims.tcgen05_mma( + _mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + first_a_desc, + first_b_desc, + idesc, + cutlass.Boolean(head_dim_stage_idx != 0), + ) + + # Derive every remaining descriptor from the immutable roots. At + # each 64-column boundary the recurrence replaces its ordinary +2 + # step with +1018 for K and +(8 * TileQ - 6), or +58, for Q; the + # closed form therefore adds each boundary jump minus that +2. + # Keeping descriptors out of iter_args avoids staged-D256 spills. + for ki in cutlass.range(1, mma_k_steps, 1, unroll=1): + if cutlass.const_expr(cfg.use_fp8_qkv): + k_desc_offset = ki * Int32(2) + q_desc_offset = ki * Int32(2) + else: + chunk_idx = (ki * Int32(_mma_k_step(cfg))) // Int32(64) + k_desc_offset = ki * Int32(2) + chunk_idx * Int32(1016) + q_chunk_extra = ( + 8 * cfg.tile_size_q - 8 + if cutlass.const_expr(cfg.tile_size_q >= 16) + else 56 + ) + q_desc_offset = ki * Int32(2) + chunk_idx * Int32(q_chunk_extra) + iter_k_desc = k_desc + k_desc_offset + iter_q_desc = q_desc + q_desc_offset + if prims.elect_sync(): + if cutlass.const_expr(q_is_a): + iter_a_desc, iter_b_desc = iter_q_desc, iter_k_desc + else: + iter_a_desc, iter_b_desc = iter_k_desc, iter_q_desc + prims.tcgen05_mma( + _mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + iter_a_desc, + iter_b_desc, + idesc, + cutlass.Boolean(True), + ) + + @cute.jit + def _resolve_keeps_tile_context(self, stage_info: StageInfo): + """Resolve one score tile's logical position and boundary-mask state.""" + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + if cutlass.const_expr(self.seqlens_kv is None): + seq_len_kv = Int32(self.max_seq_len_kv) + else: + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + Int32(0), + Int32(0), + ) + + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + q_token_base = _q_group_token_base(cfg, logical_q_group_idx) + element_mask_end_idx = seq_len_kv + if cutlass.const_expr(cfg.uses_uniform_causal_mask): + element_mask_end_idx = seq_len_kv - self.seq_len_q + q_token_base + Int32(1) + + use_runtime_kv_domain = ( + self.seqlens_kv is not None or cfg.uses_runtime_q_kv_union + ) + local_tile_idx = _softmax_tile_idx(cfg, stage_info, self.inst_id) + if cutlass.const_expr(not use_runtime_kv_domain): + effective_tile_idx = _static_split_kv_global_tile_idx( + cfg, stage_info, local_tile_idx + ) + effective_total_kv_tiles = Int32(cfg.total_kv_tiles) + tile_idx = _clamp_valid_tile_idx(cfg, effective_tile_idx) + tile_idx = tile_idx + Int32(cfg.static_num_skipped_kv_tiles) + window_start_idx = Int32(cfg.static_window_start_idx) + elif cutlass.const_expr(cfg.use_paged_kv and not cfg.use_split_kv): + effective_tile_idx = ( + Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + local_tile_idx + ) + effective_total_kv_tiles = Int32(task_cache[_TASK_CACHE_KV_VALID_TILE_END]) + tile_idx = effective_tile_idx + window_start_idx = Int32(task_cache[_TASK_CACHE_KV_WINDOW_START]) + else: + effective_tile_idx = _runtime_split_kv_global_tile_idx( + cfg, + stage_info, + local_tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + effective_total_kv_tiles = _runtime_total_kv_tiles( + cfg, seq_len_kv, self.seq_len_q, q_token_base + ) + tile_idx = _runtime_clamp_valid_tile_idx( + cfg, + effective_tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + tile_idx = tile_idx + _num_skipped_kv_tiles( + cfg, seq_len_kv, self.seq_len_q, q_token_base + ) + window_start_idx = _sliding_window_start_idx( + cfg, seq_len_kv, self.seq_len_q, q_token_base + ) + + tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) + is_valid_effective_tile = effective_tile_idx < effective_total_kv_tiles + is_masked_final_wave = False + if cutlass.const_expr(not use_runtime_kv_domain and not cfg.use_split_kv): + if cutlass.const_expr(cfg.has_odd_kv_tail and self.inst_id == 1): + is_masked_final_wave = _is_last_loop_iteration(stage_info) + tile_has_valid_scores = ( + is_valid_effective_tile + and (tile_offset_k < seq_len_kv) + and not is_masked_final_wave + ) + tile_is_unmasked = _kv_tile_is_fully_unmasked_for_q_group( + cfg, + tile_offset_k, + seq_len_kv, + self.seq_len_q, + q_token_base, + tile_has_valid_scores, + ) + return ( + seq_len_kv, + logical_q_group_idx, + element_mask_end_idx, + tile_offset_k, + window_start_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + ) + + @cute.jit + def _reduce_keeps_row_max(self, s_vals: cutlass.Array) -> Float32: + """Reduce one Keeps score row while preserving its lane ownership.""" + + cfg = self.cfg + max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for chain_idx in cutlass.range_constexpr(4): + max_chains[chain_idx] = _neg_max_f32() + for reg_base in cutlass.range_constexpr(0, cfg.num_s_regs_per_thread, 4): + for chain_idx in cutlass.range_constexpr(4): + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], + s_vals[reg_base + chain_idx], + ftz=True, + ) + tile_max = cute.math.max( + cute.math.max(max_chains[0], max_chains[1], ftz=True), + cute.math.max(max_chains[2], max_chains[3], ftz=True), + ftz=True, + ) + if cutlass.const_expr(cfg.tile_size_q == 64): + # A Q64 row is split across lanes xor 16; Q128 already owns the + # complete row locally and therefore needs no cross-lane combine. + tile_max = cute.math.max( + tile_max, + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=tile_max, + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + return tile_max + + @cute.jit + def _publish_keeps_softmax_state( + self, + s_vals: cutlass.Array, + tile_max: Float32, + old_max: Float32, + running_sum: Float32, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> None: + """Publish a masked Keeps row and its updated softmax anchor.""" + + new_anchor = cute.math.max(old_max, tile_max, ftz=True) + if cutlass.const_expr( + self.cfg.use_block_sparse and _BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2 > 0.0 + ): + # Online softmax only requires a common finite reference for P, + # sum, and O; it does not require the exact row maximum. Defer a + # small anchor increase so correction can skip a TMEM O rescale. + rescale_log2 = (old_max - new_anchor) * self.scale_softmax_log2 + if (old_max != _neg_max_f32()) and ( + rescale_log2 >= Float32(-_BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2) + ): + new_anchor = old_max + old_max_arr[0] = old_max + sum_arr[0] = running_sum + new_max_arr[0] = new_anchor + for reg_idx in cutlass.range_constexpr(self.cfg.num_s_regs_per_thread): + s_arr[reg_idx] = s_vals[reg_idx] + + @cute.jit + def _mask_and_store_sparse_keeps_atom( + self, + s_vals: cutlass.Array, + loaded: cutlass.Vector, + token_word: Uint32, + *, + atom_col: Constexpr[int], + token_mask_is_required: cutlass.Boolean, + ) -> None: + """Store one 32-score atom, applying its token word when required.""" + + if token_mask_is_required: + for atom_reg_idx in cutlass.range_constexpr(32): + score_idx = atom_col + atom_reg_idx + s_vals[score_idx] = loaded[atom_reg_idx] + token_bit_is_valid = ( + (token_word >> Int32(atom_reg_idx)) & Uint32(1) + ) != Uint32(0) + if not token_bit_is_valid: + s_vals[score_idx] = _neg_max_f32() + else: + for atom_reg_idx in cutlass.range_constexpr(32): + score_idx = atom_col + atom_reg_idx + s_vals[score_idx] = loaded[atom_reg_idx] + + @cute.jit + def _load_keeps_fragment_impl( + self, + stage_info: StageInfo, + s_vals: cutlass.Array, + tile_offset_k: Int32, + element_mask_end_idx: Int32, + window_start_idx: Int32, + seq_len_kv: Int32, + logical_q_group_idx: Int32, + is_valid_effective_tile: cutlass.Boolean, + is_masked_final_wave: cutlass.Boolean, + *, + apply_boundary_mask: Constexpr[bool], + fragment_idx: Constexpr[int] = 0, + ) -> None: + """Load one Keeps score fragment with a compile-time mask policy. + + The caller chooses the masked/unmasked path before TMEM load. Keeping the + score fragment out of the branch condition avoids carrying 64/128 live + S registers through a post-load control-flow edge. Max reduction is a + separate operation because the later P-materialization reload only + needs the masked scores. + """ + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + num_s_regs = cfg.softmax_score_fragment_regs + fragment_reg_base = fragment_idx * num_s_regs + base_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + self._softmax_loop_stage_slot_offset(stage_info) + ) + for load_atom_idx in cutlass.range_constexpr(num_s_regs // 32): + atom_col = load_atom_idx * 32 + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr( + base_addr + Int32(fragment_reg_base + atom_col), Float32 + ), + num=32, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for atom_reg_idx in cutlass.range_constexpr(32): + s_vals[atom_col + atom_reg_idx] = loaded[atom_reg_idx] + + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + tile_row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) + col_base = _keeps_col_base(cfg, lane_idx, num_s_regs) + + if cutlass.const_expr(apply_boundary_mask): + # Runtime native no-split paging uses the absolute effective tile + # index. For active rows, an invalid tile therefore begins at or + # beyond the CTA's causal union; each row's upper mask suppresses + # the complete tile. Inactive rows are safe only when absent or + # guaranteed row-independent and discarded at publication. + per_row_paged_upper_mask_covers_invalid_tile = ( + self.seqlens_kv is not None + and cfg.use_paged_kv + and not cfg.use_split_kv + and cfg.uses_per_row_causal_mask + and not cfg.use_sliding_window_causal + and (cfg.q_tiles_are_full or cfg.uses_guarded_grouped_keeps_output_rows) + ) + if cutlass.const_expr(not per_row_paged_upper_mask_covers_invalid_tile): + if not ( + is_valid_effective_tile + and tile_offset_k < seq_len_kv + and not is_masked_final_wave + ): + for reg_idx in cutlass.range_constexpr(num_s_regs): + s_vals[reg_idx] = _neg_max_f32() + + # A per-row causal endpoint is always <= seq_len_kv and is + # applied by the loop below. Avoid emitting a second, + # mathematically redundant upper-bound pass for grouped Q. + if cutlass.const_expr(not cfg.uses_per_row_causal_mask): + for reg_idx in cutlass.range_constexpr(num_s_regs): + token_idx = tile_offset_k + _keeps_score_col( + cfg, + warp_grp_thread_idx, + fragment_reg_base + reg_idx, + col_base, + ) + if token_idx >= element_mask_end_idx: + s_vals[reg_idx] = _neg_max_f32() + if cutlass.const_expr(cfg.use_sliding_window_causal): + if token_idx < window_start_idx: + s_vals[reg_idx] = _neg_max_f32() + + if cutlass.const_expr(cfg.uses_per_row_causal_mask): + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + causal_start = Int32(0) + if cutlass.const_expr(cfg.use_sliding_window_causal): + causal_start = cute.math.max( + causal_end - Int32(cfg.attention_window_size), Int32(0) + ) + causal_start_rel = causal_start - tile_offset_k + causal_end_rel = causal_end - tile_offset_k + for reg_idx in cutlass.range_constexpr(num_s_regs): + score_col = _keeps_score_col( + cfg, + warp_grp_thread_idx, + fragment_reg_base + reg_idx, + col_base, + ) + if cutlass.const_expr(cfg.use_sliding_window_causal): + if score_col < causal_start_rel: + s_vals[reg_idx] = _neg_max_f32() + if score_col >= causal_end_rel: + s_vals[reg_idx] = _neg_max_f32() + + if cutlass.const_expr(cfg.q_score_rows_need_mask): + if not _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ): + for reg_idx in cutlass.range_constexpr(num_s_regs): + s_vals[reg_idx] = _neg_max_f32() + + @cute.jit + def _load_keeps_fragment( + self, + stage_info: StageInfo, + s_vals: cutlass.Array, + tile_offset_k: Int32, + element_mask_end_idx: Int32, + window_start_idx: Int32, + seq_len_kv: Int32, + logical_q_group_idx: Int32, + is_valid_effective_tile: cutlass.Boolean, + is_masked_final_wave: cutlass.Boolean, + tile_is_unmasked: cutlass.Boolean, + *, + fragment_idx: Constexpr[int] = 0, + ) -> None: + """Select the masked or unmasked fragment loader before LDTM. + + ``tile_is_unmasked`` is runtime state, whereas the implementation's + mask policy remains constexpr. Keeping the branch outside the loader + lets the unmasked specialization erase boundary-mask instructions and + avoids carrying the loaded score registers through a post-LDTM branch. + """ + if tile_is_unmasked: + self._load_keeps_fragment_impl( + stage_info, + s_vals, + tile_offset_k, + element_mask_end_idx, + window_start_idx, + seq_len_kv, + logical_q_group_idx, + is_valid_effective_tile, + is_masked_final_wave, + apply_boundary_mask=False, + fragment_idx=fragment_idx, + ) + else: + self._load_keeps_fragment_impl( + stage_info, + s_vals, + tile_offset_k, + element_mask_end_idx, + window_start_idx, + seq_len_kv, + logical_q_group_idx, + is_valid_effective_tile, + is_masked_final_wave, + apply_boundary_mask=True, + fragment_idx=fragment_idx, + ) + + @cute.jit + def _reduce_keeps_fragment_max(self, s_vals: cutlass.Array) -> Float32: + """Reduce the row maximum of a previously loaded Keeps fragment.""" + cfg = self.cfg + num_s_regs = cfg.softmax_score_fragment_regs + + max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for chain_idx in cutlass.range_constexpr(4): + max_chains[chain_idx] = _neg_max_f32() + for reg_base in cutlass.range_constexpr(0, num_s_regs, 4): + for chain_idx in cutlass.range_constexpr(4): + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], + s_vals[reg_base + chain_idx], + ftz=True, + ) + tile_max = cute.math.max( + cute.math.max(max_chains[0], max_chains[1], ftz=True), + cute.math.max(max_chains[2], max_chains[3], ftz=True), + ftz=True, + ) + if cutlass.const_expr(cfg.tile_size_q == 64 and cfg.tile_size_kv != 256): + return cute.math.max( + tile_max, + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=tile_max, + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + return tile_max + + @cute.jit + def _decode_sparse_mask_metadata( + self, + routed_origin0: Int32, + routed_origin1: Int32, + routed_route_flags: Int32, + routed_token_word0: Uint32, + routed_token_word1: Uint32, + routed_token_word2: Uint32, + routed_token_word3: Uint32, + ) -> tuple[Int32, Int32, Int32, Int32, cutlass.Array, cutlass.Boolean]: + """Decode one prepared, register-routed mask payload.""" + + origin0 = Int32(routed_origin0) + origin1 = Int32(routed_origin1) + route_flags = Int32(routed_route_flags) + valid0 = route_flags & Int32(1) + valid1 = (route_flags >> Int32(1)) & Int32(1) + route_token_mask_is_full = cutlass.Boolean(False) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + route_token_mask_is_full = cutlass.Boolean( + (route_flags & Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG)) != Int32(0) + ) + + num_local_words = 4 if self.cfg.tile_size_q == 128 else 2 + local_token_words = cutlass.Array( + Uint32, + num_local_words, + space=cutlass.AddressSpace.rmem, + ) + for word_idx in cutlass.range_constexpr(num_local_words): + local_token_words[word_idx] = Uint32(0xFFFFFFFF) + if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if not route_token_mask_is_full: + if cutlass.const_expr(self.cfg.tile_size_q == 128): + local_token_words[0] = Uint32(routed_token_word0) + local_token_words[1] = Uint32(routed_token_word1) + local_token_words[2] = Uint32(routed_token_word2) + local_token_words[3] = Uint32(routed_token_word3) + else: + local_word0 = Uint32(routed_token_word0) + local_word1 = Uint32(routed_token_word1) + local_token_words[0] = local_word0 + local_token_words[1] = local_word1 + return ( + origin0, + origin1, + valid0, + valid1, + local_token_words, + route_token_mask_is_full, + ) + + @cute.jit + def _compute_softmax_loop_sparse_keeps( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + routed_origin0: Int32, + routed_origin1: Int32, + routed_route_flags: Int32, + routed_token_word0: Uint32, + routed_token_word1: Uint32, + routed_token_word2: Uint32, + routed_token_word3: Uint32, + ) -> tuple[object, object, object, object]: + """Load Keeps scores and mask them in logical KV coordinates.""" + cfg = self.cfg + num_s_regs = cfg.num_s_regs_per_thread + old_max = new_max_arr[0] + running_sum = sum_arr[0] + s_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) + task_cache = _decode_gen_task_cache(stage_info) + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + Int32(0), + Int32(0), + ) + warp_grp_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) + lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) + tile_row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) + col_base = _keeps_col_base(cfg, lane_idx, num_s_regs) + ( + origin0, + origin1, + valid0, + valid1, + local_token_words, + route_token_mask_is_full, + ) = self._decode_sparse_mask_metadata( + routed_origin0=routed_origin0, + routed_origin1=routed_origin1, + routed_route_flags=routed_route_flags, + routed_token_word0=routed_token_word0, + routed_token_word1=routed_token_word1, + routed_token_word2=routed_token_word2, + routed_token_word3=routed_token_word3, + ) + + base_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + self._softmax_loop_stage_slot_offset(stage_info) + ) + num_load_atoms = num_s_regs // 32 + if cutlass.const_expr(cfg.tile_size_q == 64 and cfg.use_kv_valid_bits): + token_mask_is_required = not route_token_mask_is_full + + # Keep each Q64 atom's load, wait, and mask together. A/B testing + # showed that hoisting both loads extends live fragment ranges and + # regresses the Q64 code generated by ptxas. + for load_atom_idx in cutlass.range_constexpr(2): + atom_col = load_atom_idx * 32 + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), + num=32, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + self._mask_and_store_sparse_keeps_atom( + s_vals, + loaded, + local_token_words[load_atom_idx], + atom_col=atom_col, + token_mask_is_required=token_mask_is_required, + ) + else: + for load_atom_idx in cutlass.range_constexpr(num_load_atoms): + atom_col = load_atom_idx * 32 + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), + num=32, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for atom_reg_idx in cutlass.range_constexpr(32): + score_idx = atom_col + atom_reg_idx + s_vals[score_idx] = loaded[atom_reg_idx] + + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + q_row_is_valid = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ) + causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + can_skip_structural_mask = _can_skip_sparse_keeps_structural_mask( + q_row_is_valid, + origin0, + origin1, + valid0, + valid1, + seq_len_kv, + causal_end, + apply_causal_mask=cfg.mask_type == CAUSAL, + ) + # This guard covers only route/Q/tail/causal structure. Q64 token + # holes were applied while materializing its two LDTM atoms; Q128 + # applies them in the post-pass below. + if not can_skip_structural_mask: + for reg_idx in cutlass.range_constexpr(num_s_regs): + fragment_offset = Int32(reg_idx) + logical_k = origin0 + fragment_offset + fragment_valid = valid0 + if cutlass.const_expr(cfg.tile_size_q == 128 and reg_idx >= 64): + fragment_offset = Int32(reg_idx - 64) + logical_k = origin1 + fragment_offset + fragment_valid = valid1 + elif cutlass.const_expr(cfg.tile_size_q == 64): + if col_base >= Int32(64): + logical_k = origin1 + fragment_offset + fragment_valid = valid1 + + score_is_valid = ( + q_row_is_valid + and fragment_valid != Int32(0) + and logical_k < seq_len_kv + ) + if cutlass.const_expr(cfg.mask_type == CAUSAL): + score_is_valid = score_is_valid and logical_k < causal_end + if not score_is_valid: + s_vals[reg_idx] = _neg_max_f32() + + # Q128 deliberately keeps all four LDTM atoms adjacent: unlike Q64, + # interleaving each load with mask control flow regresses its codegen. + # The post-pass follows structural masking; the producer's runtime + # route flag skips it only when all four current token words are full. + if cutlass.const_expr(cfg.tile_size_q == 128 and cfg.use_kv_valid_bits): + token_mask_is_required = not route_token_mask_is_full + if token_mask_is_required: + for word_idx in cutlass.range_constexpr(4): + token_word = local_token_words[word_idx] + for bit_idx in cutlass.range_constexpr(32): + reg_idx = word_idx * 32 + bit_idx + token_bit_is_valid = ( + (token_word >> Int32(bit_idx)) & Uint32(1) + ) != Uint32(0) + if not token_bit_is_valid: + s_vals[reg_idx] = _neg_max_f32() + + tile_max = self._reduce_keeps_row_max(s_vals) + self._publish_keeps_softmax_state( + s_vals, + tile_max, + old_max, + running_sum, + old_max_arr, + sum_arr, + new_max_arr, + s_arr, + ) + return old_max_arr, sum_arr, new_max_arr, s_arr + + @cute.jit + def _compute_softmax_loop_keeps( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> tuple[object, object, object, object]: + """Load and reduce the row-major Keeps S fragment. + + TQ128 assigns one complete Q row to each warp-group thread. TQ64 + assigns one row to a lane pair: lanes ``xor 16`` own the low/high + 64-column halves. This path deliberately avoids the Swaps scratch + reduction, whose 16x256b register mapping is unrelated to Keeps. + """ + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + num_s_regs = cfg.softmax_score_fragment_regs + old_max = new_max_arr[0] + running_sum = sum_arr[0] + s_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) + use_runtime_paged_dense_load = ( + self.seqlens_kv is not None + and cfg.use_paged_kv + and not cfg.use_sliding_window_causal + and cfg.tile_size_q in (64, 128) + ) + use_preload_mask_split = ( + use_runtime_paged_dense_load or cfg.uses_per_row_causal_mask + ) + if cutlass.const_expr(not use_preload_mask_split): + for reg_idx in cutlass.range_constexpr(num_s_regs): + s_vals[reg_idx] = _neg_max_f32() + + ( + seq_len_kv, + logical_q_group_idx, + element_mask_end_idx, + tile_offset_k, + window_start_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + ) = self._resolve_keeps_tile_context(stage_info) + + if cutlass.const_expr(cfg.tile_size_kv == 256): + # KV256 owns four physical K32 fragments per lane. Reduce the max + # one fragment at a time so only one native LDTM atom is live; the + # P pass reloads the same fragments after the reference max is + # known. + tile_max = _neg_max_f32() + for fragment_idx in cutlass.range_constexpr( + cfg.num_softmax_score_fragments + ): + self._load_keeps_fragment( + stage_info, + s_vals, + tile_offset_k, + element_mask_end_idx, + window_start_idx, + seq_len_kv, + logical_q_group_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + fragment_idx=fragment_idx, + ) + fragment_max = self._reduce_keeps_fragment_max(s_vals) + tile_max = cute.math.max(tile_max, fragment_max, ftz=True) + + new_max = cute.math.max(old_max, tile_max, ftz=True) + if old_max != _neg_max_f32(): + # Keeping the previous reference max avoids an in-place O + # rescale when the new tile raises it only modestly. The + # 16-bit P path can represent the bounded values above one; the + # numerator and denominator remain in the same scale frame. + # Large jumps still rebase to keep P comfortably in range. + max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) + if max_delta_log2 >= Float32(-KV_TILE_256_RESCALE_THRESHOLD_LOG2): + new_max = old_max + old_max_arr[0] = old_max + sum_arr[0] = running_sum + new_max_arr[0] = new_max + for reg_idx in cutlass.range_constexpr(num_s_regs): + s_arr[reg_idx] = s_vals[reg_idx] + return old_max_arr, sum_arr, new_max_arr, s_arr + + if cutlass.const_expr(use_preload_mask_split): + # Select the complete unmasked/masked TMEM load+max path before any S + # registers are materialized. The shared predicate covers the + # intersection of all active grouped-Q causal/window intervals. + tile_max = _neg_max_f32() + self._load_keeps_fragment( + stage_info, + s_vals, + tile_offset_k, + element_mask_end_idx, + window_start_idx, + seq_len_kv, + logical_q_group_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + ) + tile_max = self._reduce_keeps_fragment_max(s_vals) + + self._publish_keeps_softmax_state( + s_vals, + tile_max, + old_max, + running_sum, + old_max_arr, + sum_arr, + new_max_arr, + s_arr, + ) + return old_max_arr, sum_arr, new_max_arr, s_arr + + should_load_s = ( + is_valid_effective_tile + and (tile_offset_k < seq_len_kv) + and not is_masked_final_wave + ) + if should_load_s: + base_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + self._softmax_loop_stage_slot_offset(stage_info) + ) + # Keep each intrinsic result at the native 32-register atom. TQ128 + # uses four consecutive atoms and TQ64 uses two atoms with the same + # half-split offset, avoiding a monolithic x64/x128 LLVM intrinsic + # result. + load_atom_regs = 32 + for load_atom_idx in cutlass.range_constexpr(num_s_regs // load_atom_regs): + atom_col = load_atom_idx * load_atom_regs + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), + num=load_atom_regs, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for atom_reg_idx in cutlass.range_constexpr(load_atom_regs): + s_vals[atom_col + atom_reg_idx] = loaded[atom_reg_idx] + + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + tile_row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) + col_base = _keeps_col_base(cfg, lane_idx, num_s_regs) + + # Tail/uniform-causal/window masks use each register's true logical K + # column. For q64, the paired lane owns the complementary 64-column half. + for reg_idx in cutlass.range_constexpr(num_s_regs): + token_idx = tile_offset_k + _keeps_score_col( + cfg, warp_grp_thread_idx, reg_idx, col_base + ) + # The per-row causal pass below subsumes seq_len_kv's upper bound. + if cutlass.const_expr(not cfg.uses_per_row_causal_mask): + if token_idx >= element_mask_end_idx: + s_vals[reg_idx] = _neg_max_f32() + if cutlass.const_expr(cfg.use_sliding_window_causal): + if token_idx < window_start_idx: + s_vals[reg_idx] = _neg_max_f32() + + if cutlass.const_expr(cfg.uses_per_row_causal_mask): + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + causal_start = Int32(0) + if cutlass.const_expr(cfg.use_sliding_window_causal): + causal_start = cute.math.max( + causal_end - Int32(cfg.attention_window_size), Int32(0) + ) + for reg_idx in cutlass.range_constexpr(num_s_regs): + token_idx = tile_offset_k + _keeps_score_col( + cfg, warp_grp_thread_idx, reg_idx, col_base + ) + if token_idx < causal_start or token_idx >= causal_end: + s_vals[reg_idx] = _neg_max_f32() + + if cutlass.const_expr(cfg.q_score_rows_need_mask): + if not _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ): + for reg_idx in cutlass.range_constexpr(num_s_regs): + s_vals[reg_idx] = _neg_max_f32() + + tile_max = self._reduce_keeps_row_max(s_vals) + self._publish_keeps_softmax_state( + s_vals, + tile_max, + old_max, + running_sum, + old_max_arr, + sum_arr, + new_max_arr, + s_arr, + ) + return old_max_arr, sum_arr, new_max_arr, s_arr + + @consumer_work(returns=s_arr, work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def load_softmax_p_fragment( + self, + stage_info: StageInfo, + *, + fragment_idx: Constexpr[int], + s_arr: cutlass.Array, + ) -> cutlass.Array: + """Reload and mask one KV256 K32 fragment for P materialization.""" + if cutlass.const_expr(self.cfg.use_block_sparse): + return self._load_block_sparse_softmax_p_fragment( + stage_info, + fragment_idx=fragment_idx, + s_arr=s_arr, + ) + ( + seq_len_kv, + logical_q_group_idx, + element_mask_end_idx, + tile_offset_k, + window_start_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + ) = self._resolve_keeps_tile_context(stage_info) + self._load_keeps_fragment( + stage_info, + s_arr, + tile_offset_k, + element_mask_end_idx, + window_start_idx, + seq_len_kv, + logical_q_group_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + fragment_idx=fragment_idx, + ) + return s_arr + + @cute.jit + def _sparse_swaps_logical_k( + self, + lane_k_offset: Int32, + sparse_origin0: Int32, + sparse_origin1: Int32, + sparse_origin2: Int32, + sparse_origin3: Int32, + *, + token_group_idx: Constexpr[int], + ) -> tuple[Int32, Int32]: + """Map one SWAP register group to its routed logical K position.""" + + atom_size = min(self.cfg.kv_block_size, 32) + groups_per_atom = atom_size // 8 + origin_idx = token_group_idx // groups_per_atom + atom_origin = sparse_origin0 + if cutlass.const_expr(origin_idx == 1): + atom_origin = sparse_origin1 + elif cutlass.const_expr(origin_idx == 2): + atom_origin = sparse_origin2 + elif cutlass.const_expr(origin_idx == 3): + atom_origin = sparse_origin3 + token_offset = (token_group_idx % groups_per_atom) * 8 + return atom_origin, atom_origin + Int32(token_offset) + lane_k_offset + + @cute.jit + def _compute_softmax_loop_swaps( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + sparse_origin0: Int32, + sparse_origin1: Int32, + sparse_origin2: Int32, + sparse_origin3: Int32, + sparse_token_word: Uint32, + sparse_route_flags: Uint32, + use_sparse: Constexpr[bool], + ) -> tuple[object, object, object, object]: + """Load SWAP S from TMEM and materialize the running softmax state. + + Operation order: load BMM1 scores, apply tail/window masks, reduce the + row max through shared scratch, and return the old/new max payload that + correction consumes. + """ + cfg = self.cfg + assert not cfg.use_keeps_mma_ab + # ConsWork: consume the committed S tile, update the running max state, + # and forward masked S registers to the P producer. + # Start from the previously published running max/sum and a fresh + # local S buffer for this tile. + num_scale_groups = cfg.num_softmax_scale_groups + q_repeats = max(cfg.tile_size_q // 8, 1) + old_max_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + sum_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + new_max_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + local_max_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + s_vals = cutlass.Array( + Float32, cfg.num_s_regs_per_thread, space=cutlass.AddressSpace.rmem + ) + use_runtime_paged_dense_load = ( + self.seqlens_kv is not None + and cfg.use_paged_kv + and not cfg.use_split_kv + and cfg.max_seq_len_q == 1 + and not cfg.use_sliding_window_causal + ) + for idx in cutlass.range_constexpr(num_scale_groups): + old_max_vals[idx] = new_max_arr[idx] + sum_vals[idx] = sum_arr[idx] + new_max_vals[idx] = new_max_arr[idx] + local_max_vals[idx] = _neg_max_f32() + if cutlass.const_expr(not use_runtime_paged_dense_load): + for idx in cutlass.range_constexpr(cfg.num_s_regs_per_thread): + s_vals[idx] = _neg_max_f32() + task_cache = _decode_gen_task_cache(stage_info) + if cutlass.const_expr(self.seqlens_kv is None): + seq_len_kv = Int32(self.max_seq_len_kv) + else: + # Variable-seqlen kernels read the active sequence length for + # the logical batch carried by the work tile. + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + Int32(0), + Int32(0), + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + q_token_base = _q_group_token_base(cfg, logical_q_group_idx) + element_mask_end_idx = seq_len_kv + if cutlass.const_expr(cfg.uses_uniform_causal_mask): + element_mask_end_idx = seq_len_kv - self.seq_len_q + q_token_base + Int32(1) + should_load_s = True + if cutlass.const_expr(not use_sparse): + use_runtime_kv_domain = ( + self.seqlens_kv is not None or cfg.uses_runtime_q_kv_union + ) + local_tile_idx = _softmax_tile_idx(cfg, stage_info, self.inst_id) + if cutlass.const_expr(not use_runtime_kv_domain): + # Static path: compute the effective global tile and any + # sliding-window prefix skip at compile time. + effective_tile_idx = _static_split_kv_global_tile_idx( + cfg, stage_info, local_tile_idx + ) + effective_total_kv_tiles = Int32(cfg.total_kv_tiles) + tile_idx = _clamp_valid_tile_idx(cfg, effective_tile_idx) + tile_idx = tile_idx + Int32(cfg.static_num_skipped_kv_tiles) + window_start_idx = Int32(cfg.static_window_start_idx) + elif cutlass.const_expr(cfg.use_paged_kv and not cfg.use_split_kv): + # Non-split native paging can consume the task's affine raw tile + # geometry directly. Split-KV retains its existing resolver because + # that path benchmarks faster with the general softmax mapping. + effective_tile_idx = ( + Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + local_tile_idx + ) + effective_total_kv_tiles = Int32( + task_cache[_TASK_CACHE_KV_VALID_TILE_END] + ) + tile_idx = effective_tile_idx + window_start_idx = Int32(task_cache[_TASK_CACHE_KV_WINDOW_START]) + else: + # Runtime path: compute the same values from the batch-specific + # sequence length. + effective_tile_idx = _runtime_split_kv_global_tile_idx( + cfg, + stage_info, + local_tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + effective_total_kv_tiles = _runtime_total_kv_tiles( + cfg, seq_len_kv, self.seq_len_q, q_token_base + ) + tile_idx = _runtime_clamp_valid_tile_idx( + cfg, + effective_tile_idx, + seq_len_kv, + self.seq_len_q, + q_token_base, + ) + tile_idx = tile_idx + _num_skipped_kv_tiles( + cfg, seq_len_kv, self.seq_len_q, q_token_base + ) + window_start_idx = _sliding_window_start_idx( + cfg, seq_len_kv, self.seq_len_q, q_token_base + ) + tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) + is_valid_effective_tile = effective_tile_idx < effective_total_kv_tiles + is_masked_final_wave = False + if cutlass.const_expr(not use_runtime_kv_domain and not cfg.use_split_kv): + if cutlass.const_expr(cfg.has_odd_kv_tail and self.inst_id == 1): + # The second instance in an odd tail is a prefetch duplicate + # and must not contribute to softmax. + is_masked_final_wave = _is_last_loop_iteration(stage_info) + + if cutlass.const_expr(not use_runtime_paged_dense_load): + should_load_s = ( + is_valid_effective_tile + and (tile_offset_k < seq_len_kv) + and not is_masked_final_wave + ) + else: + # Sparse routes always have a committed S tile. Invalid atoms were + # zero-filled by TMA and are suppressed below by either the staged + # origin predicate or the prepared token word. + use_runtime_kv_domain = False + effective_tile_idx = Int32(0) + effective_total_kv_tiles = Int32(1) + tile_offset_k = Int32(0) + window_start_idx = Int32(0) + is_masked_final_wave = cutlass.Boolean(False) + if should_load_s: + # ConsWork: load the S tile produced by BMM1 from TMEM into + # registers. Two TMEM rows cover the two K subtiles. Invalid + # odd-tail waves intentionally skip the load and leave S at + # -inf so the later P path contributes zero probability. + base_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + self._softmax_loop_stage_slot_offset(stage_info) + ) + + shape = "16x256b" + loaded0 = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr(base_addr, Float32), + num=q_repeats, + ) + loaded1 = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr(base_addr + Int32(16 << 16), Float32), + num=q_repeats, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for repeat_idx in cutlass.range_constexpr(q_repeats): + ld_base = repeat_idx * 4 + s_vals[ld_base + 0] = loaded0[ld_base + 0] + s_vals[ld_base + 1] = loaded0[ld_base + 1] + s_vals[ld_base + 2] = loaded0[ld_base + 2] + s_vals[ld_base + 3] = loaded0[ld_base + 3] + s_vals[q_repeats * 4 + ld_base + 0] = loaded1[ld_base + 0] + s_vals[q_repeats * 4 + ld_base + 1] = loaded1[ld_base + 1] + s_vals[q_repeats * 4 + ld_base + 2] = loaded1[ld_base + 2] + s_vals[q_repeats * 4 + ld_base + 3] = loaded1[ld_base + 3] + + if cutlass.const_expr(use_sparse): + # Route, KV-tail, uniform-causal, and token validity depend only on + # K, so one predicate masks the adjacent pair of Q-row registers. + should_apply_sparse_mask = cutlass.Boolean(True) + if cutlass.const_expr( + _swaps_forwards_packed_route_full(cfg) + or _swaps_uses_origin0_k32_full_guard(cfg) + ): + if cutlass.const_expr(_swaps_forwards_packed_route_full(cfg)): + # Prepare already proved structural fullness for the + # complete KV128 route and staging replicated the summary + # for this Softmax warp's logical K32 slice. + k32_is_full = cute.arch.make_warp_uniform( + cutlass.Boolean( + (sparse_route_flags & Uint32(_PREPARED_ROUTE_IS_FULL_FLAG)) + != Uint32(0) + ) + ) + else: + # One origin covers this warp's K32 slice, so it can + # bypass all four lane-local K8 predicates. B16 stays on + # the straight-line path: its two-origin guard is not + # cheaper after code generation. + k32_is_full = cute.arch.make_warp_uniform( + cutlass.Boolean( + sparse_origin0 >= Int32(0) + and sparse_origin0 <= seq_len_kv - Int32(32) + ) + ) + should_apply_sparse_mask = cutlass.Boolean(not k32_is_full) + if should_apply_sparse_mask: + lane_k_offset = Int32(task_cache[_TASK_CACHE_LANE_IDX]) >> Int32(2) + token_word_covers_kv_tail = _swaps_token_word_covers_kv_tail(cfg) + for token_group_idx in cutlass.range_constexpr(4): + atom_origin, logical_k = self._sparse_swaps_logical_k( + lane_k_offset, + sparse_origin0, + sparse_origin1, + sparse_origin2, + sparse_origin3, + token_group_idx=token_group_idx, + ) + # Prepared words zero absent atoms and the logical KV + # tail. Qualified profiles can therefore omit the local + # atom-origin guard, independently of the K/V issuer warp. + score_is_valid = cutlass.Boolean(True) + if cutlass.const_expr( + not _swaps_uses_token_only_score_validity(cfg) + ): + score_is_valid = cutlass.Boolean(atom_origin >= Int32(0)) + if cutlass.const_expr(not token_word_covers_kv_tail): + score_is_valid = cutlass.Boolean( + score_is_valid and logical_k < seq_len_kv + ) + if cutlass.const_expr(cfg.uses_uniform_causal_mask): + score_is_valid = cutlass.Boolean( + score_is_valid and logical_k < element_mask_end_idx + ) + if cutlass.const_expr(cfg.use_kv_valid_bits): + token_bit_idx = Int32(token_group_idx * 8) + lane_k_offset + token_is_valid = ( + (sparse_token_word >> token_bit_idx) & Uint32(1) + ) != Uint32(0) + score_is_valid = cutlass.Boolean( + score_is_valid and token_is_valid + ) + if not score_is_valid: + for repeat_idx in cutlass.range_constexpr(q_repeats): + if cutlass.const_expr(token_group_idx < 2): + s_base = repeat_idx * 4 + token_group_idx * 2 + else: + s_base = ( + q_repeats * 4 + + repeat_idx * 4 + + (token_group_idx - 2) * 2 + ) + s_vals[s_base + 0] = _neg_max_f32() + s_vals[s_base + 1] = _neg_max_f32() + + if cutlass.const_expr(use_runtime_paged_dense_load): + if not ( + is_valid_effective_tile + and (tile_offset_k < seq_len_kv) + and not is_masked_final_wave + ): + for idx in cutlass.range_constexpr(cfg.num_s_regs_per_thread): + s_vals[idx] = _neg_max_f32() + + next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) + # Determine whether this tile crosses the active right endpoint or the + # start of the causal sliding window. Dense full tiles skip per-element + # masking. + if cutlass.const_expr(use_sparse): + should_apply_dense_mask = False + elif cutlass.const_expr( + not use_runtime_kv_domain and not cfg.uses_uniform_causal_mask + ): + has_static_tail_mask = (cfg.static_seq_len_kv % cfg.tile_size_kv) != 0 + has_static_window_prefix_mask = ( + cfg.use_sliding_window_causal + and (cfg.static_window_start_idx % cfg.tile_size_kv) != 0 + ) + if cutlass.const_expr( + not has_static_tail_mask and not has_static_window_prefix_mask + ): + should_apply_dense_mask = False + else: + should_apply_dense_mask = next_tile_offset_k > seq_len_kv + if cutlass.const_expr(has_static_window_prefix_mask): + should_apply_dense_mask = should_apply_dense_mask or ( + (tile_offset_k <= window_start_idx) + and (next_tile_offset_k > window_start_idx) + ) + else: + # Runtime tails and the one-endpoint ungrouped causal path need + # element masking only on the tile that crosses the right bound. + should_apply_dense_mask = next_tile_offset_k > element_mask_end_idx + if cutlass.const_expr(cfg.use_sliding_window_causal): + window_start_remainder = window_start_idx % Int32(cfg.tile_size_kv) + should_apply_dense_mask = should_apply_dense_mask or ( + (window_start_remainder != Int32(0)) + and (tile_offset_k <= window_start_idx) + and (next_tile_offset_k > window_start_idx) + ) + if should_apply_dense_mask: + # Mask invalid S registers to -inf so they produce zero P and + # do not affect row max or row sum. This keeps the schedule + # shape fixed even when only part of the K/V tile is valid. + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + local_idx_k0 = warp_idx * Int32(32) + (lane_idx >> Int32(2)) + for repeat_idx in cutlass.range_constexpr(q_repeats): + for token_group_idx in cutlass.range_constexpr(4): + token_idx = ( + tile_offset_k + local_idx_k0 + Int32(token_group_idx * 8) + ) + if cutlass.const_expr(token_group_idx < 2): + s_base = repeat_idx * 4 + token_group_idx * 2 + else: + s_base = ( + q_repeats * 4 + repeat_idx * 4 + (token_group_idx - 2) * 2 + ) + if token_idx >= element_mask_end_idx: + s_vals[s_base + 0] = _neg_max_f32() + s_vals[s_base + 1] = _neg_max_f32() + if cutlass.const_expr(cfg.use_sliding_window_causal): + if token_idx < window_start_idx: + s_vals[s_base + 0] = _neg_max_f32() + s_vals[s_base + 1] = _neg_max_f32() + + if cutlass.const_expr(cfg.uses_per_row_causal_mask): + apply_per_row_causal_mask = cutlass.Boolean(True) + if cutlass.const_expr(not use_sparse): + tile_has_valid_scores = ( + is_valid_effective_tile + and (tile_offset_k < seq_len_kv) + and not is_masked_final_wave + ) + apply_per_row_causal_mask = cutlass.Boolean( + not _kv_tile_is_fully_unmasked_for_q_group( + cfg, + tile_offset_k, + seq_len_kv, + self.seq_len_q, + q_token_base, + tile_has_valid_scores, + ) + ) + if apply_per_row_causal_mask: + # Grouped causal decode has a distinct causal/window bound for + # every Q token. Sparse routes always use their logical K; + # dense routes retain the boundary-tile fast path above. + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + col_group_idx = lane_idx & Int32(0x3) + local_idx_k0 = warp_idx * Int32(32) + (lane_idx >> Int32(2)) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + repeat_idx = scale_idx // 2 + pair_idx = scale_idx % 2 + tile_row_idx = ( + Int32(repeat_idx * 8) + + col_group_idx * Int32(2) + + Int32(pair_idx) + ) + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + causal_start = Int32(0) + if cutlass.const_expr(cfg.use_sliding_window_causal): + causal_start = cute.math.max( + causal_end - Int32(cfg.attention_window_size), Int32(0) + ) + for token_group_idx in cutlass.range_constexpr(4): + token_idx = ( + tile_offset_k + local_idx_k0 + Int32(token_group_idx * 8) + ) + if cutlass.const_expr(use_sparse): + _, token_idx = self._sparse_swaps_logical_k( + lane_idx >> Int32(2), + sparse_origin0, + sparse_origin1, + sparse_origin2, + sparse_origin3, + token_group_idx=token_group_idx, + ) + s_idx = ( + repeat_idx * 4 + + pair_idx + + (token_group_idx & 1) * 2 + + (token_group_idx >> 1) * q_repeats * 4 + ) + if token_idx < causal_start or token_idx >= causal_end: + s_vals[s_idx] = _neg_max_f32() + + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + if cutlass.const_expr(cfg.q_score_rows_need_mask): + # Structural grouped padding and the final partial token/head band + # must not enter max/sum or P. The Swaps TMEM layout assigns one + # logical Q row to each (column-group, scale-group) pair. + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + col_group_idx = lane_idx & Int32(0x3) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + repeat_idx = scale_idx // 2 + pair_idx = scale_idx % 2 + tile_row_idx = ( + col_group_idx * Int32(2) + Int32(pair_idx) + Int32(repeat_idx * 8) + ) + if not _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ): + s_base = repeat_idx * 4 + pair_idx + s_base_hi = q_repeats * 4 + s_base + s_vals[s_base + 0] = _neg_max_f32() + s_vals[s_base + 2] = _neg_max_f32() + s_vals[s_base_hi + 0] = _neg_max_f32() + s_vals[s_base_hi + 2] = _neg_max_f32() + + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + for scale_idx in cutlass.range_constexpr(num_scale_groups): + # Reduce this lane's S registers to one candidate per scale group. + repeat_idx = scale_idx // 2 + pair_idx = scale_idx % 2 + s_base = repeat_idx * 4 + pair_idx + s_base_hi = q_repeats * 4 + s_base + local_max = cute.math.max( + cute.math.max(s_vals[s_base + 0], s_vals[s_base + 2], ftz=True), + cute.math.max(s_vals[s_base_hi + 0], s_vals[s_base_hi + 2], ftz=True), + ftz=True, + ) + local_max = cute.math.max(local_max, old_max_vals[scale_idx], ftz=True) + local_max_vals[scale_idx] = local_max + + local_row_idx = (lane_idx >> Int32(2)) & Int32(0x3) + if cutlass.const_expr(num_scale_groups > 2): + # Transpose each four-scale block across four strided warp rows. + # Every lane then owns one reduced scale group per block. + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 4): + local_max_vals[scale_base] = _wspro_reduce_max4( + local_max_vals[scale_base], + local_max_vals[scale_base + 1], + local_max_vals[scale_base + 2], + local_max_vals[scale_base + 3], + local_row_idx, + ) + else: + # Two scale groups use the compact partial reduction. + for scale_idx in cutlass.range_constexpr(num_scale_groups): + local_max = cute.math.max( + local_max_vals[scale_idx], + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=local_max_vals[scale_idx], + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + local_max_vals[scale_idx] = cute.math.max( + local_max, + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=local_max, + offset=8, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + + if stage_info.loop_offset == stage_info.loop_start: + # Softmax consumes S in the loop stage. Persistent CTAs reuse this + # scratch across work tiles, so reset it at the first loop iteration + # before the running max atomics. The barrier prevents a lane from + # atomically updating a slot another lane is still reinitializing. + _init_softmax_scratch_u32( + self._softmax_scratch_u32, warp_grp_thread_idx, cfg.tile_size_q + ) + prims.barrier_cta_sync(self.sync_barrier_id, thread_count=128) + + col_group_idx = lane_idx & Int32(0x3) + atomic_reduce_base = col_group_idx * Int32(num_scale_groups) + if cutlass.const_expr(num_scale_groups > 2): + # Two row groups publish one partial per distributed scale group. + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 4): + scale_idx = local_row_idx + Int32(scale_base) + _smem_atomic_max_u32( + self._softmax_scratch_u32.data_ptr() + + atomic_reduce_base + + scale_idx, + _float_to_u32_for_atomic_max(local_max_vals[scale_base]), + ) + elif lane_idx < Int32(8): + # The compact fallback publishes both scale groups from eight lanes. + for scale_idx in cutlass.range_constexpr(num_scale_groups): + _smem_atomic_max_u32( + self._softmax_scratch_u32.data_ptr() + + atomic_reduce_base + + Int32(scale_idx), + _float_to_u32_for_atomic_max(local_max_vals[scale_idx]), + ) + # Wait for every SMEM atomic max to finish before reloading the + # reduced maxima. Without this barrier, the vector reload below can + # race a late writer from another softmax warp. + prims.barrier_cta_sync(self.sync_barrier_id, thread_count=128) + + reduce_base = col_group_idx * Int32(num_scale_groups) + reduced_max_ptr = self._softmax_scratch_u32.data_ptr() + reduce_base + # Reload the reduced max as one aligned vector, then decode back to + # float. Keeping this reload vectorized avoids the scalar LDS shape. + reduced_max = reduced_max_ptr.load( + count=num_scale_groups, + alignment=16 if num_scale_groups >= 4 else 8, + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + # Decode the CTA-wide maxima back into the running softmax + # state carried by this resource. + new_max_vals[scale_idx] = _u32_to_float_for_atomic_max( + reduced_max[scale_idx] + ) + + for scale_idx in cutlass.range_constexpr(num_scale_groups): + old_max_arr[scale_idx] = old_max_vals[scale_idx] + sum_arr[scale_idx] = sum_vals[scale_idx] + new_max_arr[scale_idx] = new_max_vals[scale_idx] + for idx in cutlass.range_constexpr(cfg.num_s_regs_per_thread): + # Forward the loaded/masked S registers to SmemP.compute_p. + s_arr[idx] = s_vals[idx] + return old_max_arr, sum_arr, new_max_arr, s_arr + + @consumer_work(returns=(old_max_arr, sum_arr, new_max_arr, s_arr)) + @cute.jit + def compute_softmax_loop( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> tuple[object, object, object, object]: + """Load S from TMEM and materialize the running softmax state.""" + + if cutlass.const_expr(self.cfg.use_keeps_mma_ab): + return self._compute_softmax_loop_keeps( + stage_info, + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + ) + return self._compute_softmax_loop_swaps( + stage_info, + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + sparse_origin0=Int32(-1), + sparse_origin1=Int32(-1), + sparse_origin2=Int32(-1), + sparse_origin3=Int32(-1), + sparse_token_word=Uint32(0xFFFFFFFF), + sparse_route_flags=Uint32(0), + use_sparse=False, + ) + + @consumer_work( + returns=sum_arr, + work_attrs=WorkAttr.AUXILIARY, + ) + @cute.jit + def reduce_sums( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + ) -> ResourceVars: + """Fold local P sums into the running online-softmax denominators.""" + cfg = self.cfg + # ConsTailWork: denominator update runs after P has been materialized, + # so the resource-owned local sum matches the P payload consumed by BMM2. + if cutlass.const_expr(cfg.use_fp8_qkv): + # FP8 uses TmemSoftmaxGlobal to update sums after P + # quantization, so this stage only copies the corrected sums + # back into the running state. This keeps the denominator + # consistent with the quantized P actually consumed by BMM2. + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + sum_arr[scale_idx] = self.load_global_sum(scale_idx) + return sum_arr + num_scale_groups = cfg.num_softmax_scale_groups + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 2): + # Running sum recurrence: + # sum_new = sum_old * exp(old_max - new_max) + local_sum. + # local_sum comes from SmemP, after P has been produced, so the + # denominator update stays ordered after P materialization. + # Gather one pair of scale groups so the rescale and sum update can + # use paired arithmetic and publish both groups together. + old_max = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + new_max = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + local_sum = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + sum_vals = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + exp_scale = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + pair_width = _softmax_scale_pair_width(num_scale_groups, scale_base) + # KeepsMmaAb has one scale group. Initialize the unused packed-FMA + # lane explicitly, then load and publish only the live lane so no + # out-of-bounds task-local access can become LLVM poison. + for pair_idx in cutlass.range_constexpr(2): + old_max[pair_idx] = _neg_max_f32() + new_max[pair_idx] = _neg_max_f32() + local_sum[pair_idx] = Float32(0.0) + sum_vals[pair_idx] = Float32(0.0) + exp_scale[pair_idx] = Float32(0.0) + for pair_idx in cutlass.range_constexpr(pair_width): + scale_idx = scale_base + pair_idx + old_max[pair_idx] = old_max_arr[scale_idx] + new_max[pair_idx] = new_max_arr[scale_idx] + local_sum[pair_idx] = self.load_p_local_sum(scale_idx) + sum_vals[pair_idx] = sum_arr[scale_idx] + + # Dense full-tile FP16/BF16 paths never see -inf max sentinels, so + # they can compute exp(old-new) directly. General paths guard the + # sentinel to keep empty/masked groups at zero contribution. + if cutlass.const_expr( + cfg.has_static_dense_full_kv_tiles + and cfg.tile_size_q in (16, 32) + and not cfg.use_keeps_mma_ab + and not cfg.use_fp8_qkv + and cfg.q_tiles_are_full + ): + for pair_idx in cutlass.range_constexpr(pair_width): + exp_scale[pair_idx] = cute.math.exp2( + self.scale_softmax_log2 + * (old_max[pair_idx] - new_max[pair_idx]), + fastmath=True, + ) + else: + for pair_idx in cutlass.range_constexpr(pair_width): + if (old_max[pair_idx] != _neg_max_f32()) and ( + new_max[pair_idx] != _neg_max_f32() + ): + exp_scale[pair_idx] = cute.math.exp2( + self.scale_softmax_log2 + * (old_max[pair_idx] - new_max[pair_idx]), + fastmath=True, + ) + updated_sums = ffma2( + (exp_scale[0], exp_scale[1]), + (sum_vals[0], sum_vals[1]), + (local_sum[0], local_sum[1]), + ) + # Publish the updated running denominator for the next softmax tile + # and for tail correction normalization. + for pair_idx in cutlass.range_constexpr(pair_width): + sum_arr[scale_base + pair_idx] = updated_sums[pair_idx] + return sum_arr + + @cute.jit + def _compute_softmax_loop_sparse_keeps_kv256( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + sparse_origin0: Int32, + sparse_origin1: Int32, + sparse_route_flags: Int32, + sparse_token_word0: Uint32, + sparse_token_word1: Uint32, + sparse_token_word2: Uint32, + sparse_token_word3: Uint32, + ) -> tuple[object, object, object, object]: + """Reduce one sparse KV256 route as four bounded K32 fragments. + + The full route path only loads and reduces scores. A partial route + predicates one native 32-score fragment at a time and writes it back + to TMEM, so the later P pass can replay masked scores without keeping + the logical 128-score tile live in registers. + """ + + cfg = self.cfg + assert cfg.tile_size_kv == 256 + task_cache = _decode_gen_task_cache(stage_info) + token_words = ( + sparse_token_word0, + sparse_token_word1, + sparse_token_word2, + sparse_token_word3, + ) + keep_words = cutlass.Array(Uint32, 4, space=cutlass.AddressSpace.rmem) + warp_group_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) + tile_row_idx = _keeps_row_idx(cfg, warp_group_thread_idx) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + q_row_is_valid = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ) + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + Int32(0), + Int32(0), + ) + causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + origin0 = Int32(sparse_origin0) + origin1 = Int32(sparse_origin1) + valid0 = sparse_route_flags & Int32(1) + valid1 = (sparse_route_flags >> Int32(1)) & Int32(1) + for fragment_idx in cutlass.range_constexpr(4): + fragment_origin = origin0 + Int32((fragment_idx % 2) * 32) + fragment_valid = valid0 + if cutlass.const_expr(fragment_idx >= 2): + fragment_origin = origin1 + Int32((fragment_idx % 2) * 32) + fragment_valid = valid1 + keep_words[fragment_idx] = _sparse_k32_effective_keep_word( + q_row_is_valid, + fragment_origin, + fragment_valid, + Uint32(token_words[fragment_idx]), + seq_len_kv, + causal_end, + apply_causal_mask=cfg.mask_type == CAUSAL, + apply_token_mask=cfg.use_kv_valid_bits, + ) + + warp_scores_are_unmasked = cutlass.Boolean(True) + for fragment_idx in cutlass.range_constexpr(4): + warp_scores_are_unmasked = cutlass.Boolean( + warp_scores_are_unmasked + and keep_words[fragment_idx] == Uint32(0xFFFFFFFF) + ) + # The load/store branch must be uniform for each participating warp. + warp_scores_are_unmasked = cute.arch.vote_all_sync(warp_scores_are_unmasked) + + score_tmem_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + self._softmax_loop_stage_slot_offset(stage_info) + ) + max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for chain_idx in cutlass.range_constexpr(4): + max_chains[chain_idx] = _neg_max_f32() + + if warp_scores_are_unmasked: + for fragment_idx in cutlass.range_constexpr(4): + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr( + score_tmem_addr + Int32(fragment_idx * 32), Float32 + ), + num=32, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for score_idx in cutlass.range_constexpr(32): + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], + Float32(loaded[score_idx]), + ftz=True, + ) + else: + for fragment_idx in cutlass.range_constexpr(4): + fragment_addr = score_tmem_addr + Int32(fragment_idx * 32) + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(fragment_addr, Float32), + num=32, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + masked_scores = cutlass.Array( + Float32, 32, space=cutlass.AddressSpace.rmem + ) + for score_idx in cutlass.range_constexpr(32): + score = Float32(loaded[score_idx]) + score_is_kept = ( + (keep_words[fragment_idx] >> Int32(score_idx)) & Uint32(1) + ) != Uint32(0) + if not score_is_kept: + score = _neg_max_f32() + masked_scores[score_idx] = score + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], score, ftz=True + ) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(fragment_addr, Float32), + masked_scores.data_ptr().load(count=32, alignment=4), + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + + tile_max = cute.math.max( + cute.math.max(max_chains[0], max_chains[1], ftz=True), + cute.math.max(max_chains[2], max_chains[3], ftz=True), + ftz=True, + ) + old_max = new_max_arr[0] + new_max = cute.math.max(old_max, tile_max, ftz=True) + if old_max != _neg_max_f32(): + max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) + if max_delta_log2 >= Float32(-KV_TILE_256_RESCALE_THRESHOLD_LOG2): + new_max = old_max + old_max_arr[0] = old_max + new_max_arr[0] = new_max + return old_max_arr, sum_arr, new_max_arr, s_arr + + @cute.jit + def _load_block_sparse_softmax_p_fragment( + self, + stage_info: StageInfo, + *, + fragment_idx: Constexpr[int], + s_arr: cutlass.Array, + ) -> cutlass.Array: + """Reload one full or already-predicated sparse KV256 fragment for P.""" + + assert self.cfg.tile_size_kv == 256 + task_cache = _decode_gen_task_cache(stage_info) + score_tmem_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + self._softmax_loop_stage_slot_offset(stage_info) + + Int32(fragment_idx * 32) + ) + loaded = _keeps_tcgen05_ld( + self.cfg, + prims.make_tmem_ptr(score_tmem_addr, Float32), + num=32, + offset=self.cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for score_idx in cutlass.range_constexpr(32): + s_arr[score_idx] = loaded[score_idx] + return s_arr + + @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) + @cute.jit + def compute_block_sparse_softmax_loop( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + sparse_origin0: Int32, + sparse_origin1: Int32, + sparse_route_flags: Int32, + sparse_token_word0: Uint32, + sparse_token_word1: Uint32, + sparse_token_word2: Uint32, + sparse_token_word3: Uint32, + ) -> tuple[object, object, object, object]: + """Consume S plus one explicitly routed, register-resident payload.""" + + assert self.cfg.use_block_sparse + if cutlass.const_expr(self.cfg.use_keeps_mma_ab): + if cutlass.const_expr(self.cfg.tile_size_kv == 256): + return self._compute_softmax_loop_sparse_keeps_kv256( + stage_info, + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + sparse_origin0=sparse_origin0, + sparse_origin1=sparse_origin1, + sparse_route_flags=sparse_route_flags, + sparse_token_word0=sparse_token_word0, + sparse_token_word1=sparse_token_word1, + sparse_token_word2=sparse_token_word2, + sparse_token_word3=sparse_token_word3, + ) + return self._compute_softmax_loop_sparse_keeps( + stage_info, + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + routed_origin0=sparse_origin0, + routed_origin1=sparse_origin1, + routed_route_flags=sparse_route_flags, + routed_token_word0=sparse_token_word0, + routed_token_word1=sparse_token_word1, + routed_token_word2=sparse_token_word2, + routed_token_word3=sparse_token_word3, + ) + # SWAP reuses the Keeps seven-slot task ABI: all four origins remain + # logical KV atom bases, but origin2 occupies the flags slot and + # origin3 is bit-preserved in word0. Word1 carries the logical K32 token + # mask and word2 optionally carries the prepared route-full summary. + return self._compute_softmax_loop_swaps( + stage_info, + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + sparse_origin0=Int32(sparse_origin0), + sparse_origin1=Int32(sparse_origin1), + sparse_origin2=Int32(sparse_route_flags), + sparse_origin3=sparse_token_word0.bitcast(Int32), + sparse_token_word=Uint32(sparse_token_word1), + sparse_route_flags=Uint32(sparse_token_word2), + use_sparse=True, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_softmax_stats.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_softmax_stats.py new file mode 100644 index 000000000000..426cbdd7c260 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_softmax_stats.py @@ -0,0 +1,798 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Softmax stats resources for FMHA decode TS kernel. + +Holds ``TmemSoftmaxLocalResource`` (per-loop ``old_max``/``new_max``/``sum`` +arrays exchanged with the correction warps) and ``TmemSoftmaxGlobalResource`` +(FP8 sum-correction helper that reapplies running-max correction after P +quantization). ``TmemStatsDoneResource`` carries the overwrite credit for TMEM +columns shared by S and the local stats payload. +""" + +from dataclasses import dataclass +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 + +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ..fmha_decode_config import FmhaDecodeConfig +from ...placeholder_helpers import _placeholder_local_array +from .helpers_common import ( + Constexpr, + DecodeGenResourceBase, + ResourceVars, + fadd2, + ffma2, + fmul2, + _named_barrier_arrive, + _named_barrier_sync, + _TASK_CACHE_TMEM_BASE_OFFSET, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _decode_gen_task_cache, + _neg_max_f32, + _softmax_scale_pair_width, +) +from .tmem_s import TmemSResource + + +@dataclass(kw_only=True) +class TmemStatsDoneResource(MemoryResource): + """Barrier returned after Correction loads stats from aliased TMEM.""" + + is_barrier: Constexpr[bool] = True + + +@dataclass(kw_only=True) +class TmemSoftmaxOrderResource(MemoryResource): + """CTA named-barrier baton for deterministic P0/P1 ordering. + + The barrier is only enabled for profiles where softmax0 and softmax1 share + producer-side hazards. It serializes the P0/P1 publication order without + changing the TS resource graph. + """ + + cfg: Constexpr[FmhaDecodeConfig] = None + is_barrier: Constexpr[bool] = True + + @consumer_work + @cute.jit + def prime_softmax1(self, stage_info: StageInfo) -> None: + """Prime the softmax1 side of the ordered P0/P1 baton.""" + _ = stage_info + if cutlass.const_expr(self.cfg.uses_ordered_softmax_barrier): + # ConsWork: seed the first named barrier so softmax0 is allowed to + # publish before softmax1 enters its ordered slot. + _named_barrier_arrive( + self.cfg.resolved_softmax_order_barrier_threads, + barrier_id=self.cfg.softmax_order_barrier_id, + ) + + @producer_work + @cute.jit + def wait_softmax0(self, stage_info: StageInfo) -> None: + """Wait until softmax0 is allowed to publish before softmax1.""" + _ = stage_info + if cutlass.const_expr(self.cfg.uses_ordered_softmax_barrier): + # ProdWork: softmax0 waits on the baton before publishing P/stats, + # preserving the expected P0 -> P1 order for shared resources. + _named_barrier_sync( + self.cfg.resolved_softmax_order_barrier_threads, + barrier_id=self.cfg.softmax_order_barrier_id, + ) + + @producer_work + @cute.jit + def release_softmax1(self, stage_info: StageInfo) -> None: + """Signal that softmax1 may publish its P payload.""" + _ = stage_info + if cutlass.const_expr( + self.cfg.uses_ordered_softmax_barrier + and not self.cfg.ordered_softmax_early_release + ): + # ProdWork: softmax0 hands the baton to softmax1 after its P/stats + # payload is visible to downstream tasks. + _named_barrier_arrive( + self.cfg.resolved_softmax_order_barrier_threads, + barrier_id=self.cfg.softmax_order_barrier_id + 1, + ) + + @consumer_work + @cute.jit + def wait_softmax1(self, stage_info: StageInfo) -> None: + """Wait until softmax1 has completed its ordered publication slot.""" + _ = stage_info + if cutlass.const_expr(self.cfg.uses_ordered_softmax_barrier): + # ConsWork: softmax1 waits until softmax0's publication for this + # iteration is visible, preserving the P0 -> P1 order. The baton + # is a two-party softmax0/softmax1 protocol; correction does not + # participate. + _named_barrier_sync( + self.cfg.resolved_softmax_order_barrier_threads, + barrier_id=self.cfg.softmax_order_barrier_id + 1, + ) + + @consumer_work + @cute.jit + def release_softmax0(self, stage_info: StageInfo) -> None: + """Release the baton so the next softmax0 publication can proceed.""" + _ = stage_info + if cutlass.const_expr( + self.cfg.uses_ordered_softmax_barrier + and not self.cfg.ordered_softmax_early_release + ): + # ConsWork: softmax1 completes the baton cycle after its own + # publication and allows the next softmax0 slot to enter the + # ordered region. + _named_barrier_arrive( + self.cfg.resolved_softmax_order_barrier_threads, + barrier_id=self.cfg.softmax_order_barrier_id, + ) + + +@dataclass(kw_only=True) +class TmemSoftmaxLocalResource(DecodeGenResourceBase): + """TMEM-local softmax statistics exchanged with correction. + + Loop stats carry old/new maxima for in-place O correction. Tail stats carry + final sums and maxima for output normalization. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "old_max_arr", + cutlass.Array, + None, + "Previous softmax maximum read from TMEM.", + ), + ("new_max_arr", cutlass.Array, None, "Current softmax maximum read from TMEM."), + ("sum_arr", cutlass.Array, None, "Softmax denominator read from TMEM."), + ( + "inst_old_max_arr", + cutlass.Array, + None, + "Per-instruction previous softmax maximum.", + ), + ( + "inst_new_max_arr", + cutlass.Array, + None, + "Per-instruction current softmax maximum.", + ), + ("inst_sum_arr", cutlass.Array, None, "Per-instruction softmax denominator."), + ) + inst_id: Constexpr[int] = 0 + cfg: Constexpr[FmhaDecodeConfig] = None + _alloc: Constexpr[TmemAllocation | None] = None + _smem_alloc: Constexpr[SmemAllocation | None] = None + _inst_new_max_arr: cutlass.Array | None = None + _inst_sum_arr: cutlass.Array | None = None + old_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + new_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + sum_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + inst_old_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + inst_new_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + inst_sum_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def _init_placeholder_state(self) -> None: + """Create placeholder arrays for softmax-local stat handoff.""" + num_sg = self.cfg.num_softmax_scale_groups + self.old_max_arr.default = _placeholder_local_array(Float32, num_sg) + self.new_max_arr.default = _placeholder_local_array(Float32, num_sg) + self.sum_arr.default = _placeholder_local_array(Float32, num_sg) + self.inst_old_max_arr.default = _placeholder_local_array(Float32, num_sg) + self.inst_new_max_arr.default = _placeholder_local_array(Float32, num_sg) + self.inst_sum_arr.default = _placeholder_local_array(Float32, num_sg) + self._inst_new_max_arr = _placeholder_local_array(Float32, num_sg) + self._inst_sum_arr = _placeholder_local_array(Float32, num_sg) + + def _stats_smem_rows(self) -> int: + """One SMEM stats row per softmax warp-group thread.""" + num_warps = ( + self.cfg.softmax0_num_warps + if self.inst_id == 0 + else self.cfg.softmax1_num_warps + ) + return num_warps * 32 + + def _stats_smem_row_elems(self) -> int: + """FP32 payload elements per row: two stat halves per scale group.""" + return self.cfg.num_softmax_scale_groups * 2 + + def _stats_smem_alignment(self) -> int: + return min(16, self._stats_smem_row_elems() * 4) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate the SMEM stats ring when stats cannot own TMEM columns.""" + if not self.cfg.keeps_stats_via_smem: + return [] + if self._smem_alloc is None: + stage_bytes = self._stats_smem_rows() * self._stats_smem_row_elems() * 4 + self._smem_alloc = SmemAllocation( + name=f"{self.name}_smem", + size_bytes=self.pipeline_config.num_stages * stage_bytes, + alignment=16, + ) + return [self._smem_alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Allocate TMEM columns for softmax stat handoff records.""" + if self.cfg.keeps_stats_via_smem: + # Do not expose a fictitious TMEM access to the allocator or the + # exhaustive dependency checker. The constexpr store/load paths + # below use only the SMEM ring for this profile. + return [] + if self._alloc is None: + self._alloc = TmemAllocation( + name=f"{self.name}", + num_columns=self.cfg.tmem_stats_cols, + ) + return [self._alloc] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Initialize loop and tail-visible softmax stat arrays.""" + num_sg = self.cfg.num_softmax_scale_groups + # Retain final per-instance stats alongside the loop-carried values; + # the correction task receives them explicitly for the tail merge. + self._inst_new_max_arr = cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ) + self._inst_sum_arr = cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ) + result = { + "old_max_arr": cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ), + "new_max_arr": cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ), + "sum_arr": cutlass.Array(Float32, num_sg, space=cutlass.AddressSpace.rmem), + "inst_old_max_arr": cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ), + "inst_new_max_arr": self._inst_new_max_arr, + "inst_sum_arr": self._inst_sum_arr, + } + for idx in cutlass.range_constexpr(num_sg): + # Initialize both loop stats and final instance stats to the + # neutral softmax state. + result["old_max_arr"][idx] = _neg_max_f32() + result["new_max_arr"][idx] = _neg_max_f32() + result["sum_arr"][idx] = Float32(0.0) + result["inst_old_max_arr"][idx] = _neg_max_f32() + self._inst_new_max_arr[idx] = _neg_max_f32() + self._inst_sum_arr[idx] = Float32(0.0) + return result + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Create per-work-tile softmax stat arrays for persistent scheduling.""" + _ = context + num_sg = self.cfg.num_softmax_scale_groups + # Tail-visible instance stats are per work tile, so persistent kernels + # must allocate fresh local arrays when the scheduler advances. + self._inst_new_max_arr = cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ) + self._inst_sum_arr = cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ) + result = { + "old_max_arr": cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ), + "new_max_arr": cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ), + "sum_arr": cutlass.Array(Float32, num_sg, space=cutlass.AddressSpace.rmem), + "inst_old_max_arr": cutlass.Array( + Float32, num_sg, space=cutlass.AddressSpace.rmem + ), + "inst_new_max_arr": self._inst_new_max_arr, + "inst_sum_arr": self._inst_sum_arr, + } + for idx in cutlass.range_constexpr(num_sg): + result["old_max_arr"][idx] = _neg_max_f32() + result["new_max_arr"][idx] = _neg_max_f32() + result["sum_arr"][idx] = Float32(0.0) + result["inst_old_max_arr"][idx] = _neg_max_f32() + self._inst_new_max_arr[idx] = _neg_max_f32() + self._inst_sum_arr[idx] = Float32(0.0) + return result + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ), + ) + @cute.jit + def init_stats_state( + self, stage_info: StageInfo + ) -> tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ]: + """Initialize and return correction-side softmax stat state.""" + # ConsAuxWork: allocate correction-visible stat arrays that carry + # loop old/new maxima and tail per-instance sum/max payloads. + result = self._create_initial_task_locals(stage_info.context) + return ( + result["old_max_arr"], + result["new_max_arr"], + result["sum_arr"], + result["inst_old_max_arr"], + result["inst_new_max_arr"], + result["inst_sum_arr"], + ) + + @cute.jit + def _stats_ptr(self, stage_info: StageInfo): + """Return the TMEM pointer for this resource's stats handoff slot.""" + task_cache = _decode_gen_task_cache(stage_info) + stats_base = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + stage_info.stage_idx * self.cfg.tmem_s_cols + ) + return prims.make_tmem_ptr(stats_base, Float32) + + @cute.jit + def _stats_vector(self, first_arr: cutlass.Array, second_arr: cutlass.Array): + """Pack the two stats halves into the tcgen05 vector store layout.""" + cfg = self.cfg + if cutlass.const_expr(cfg.num_softmax_scale_groups == 8): + return cutlass.Vector.from_elements( + ( + first_arr[0], + first_arr[1], + first_arr[2], + first_arr[3], + first_arr[4], + first_arr[5], + first_arr[6], + first_arr[7], + second_arr[0], + second_arr[1], + second_arr[2], + second_arr[3], + second_arr[4], + second_arr[5], + second_arr[6], + second_arr[7], + ), + Float32, + ) + elif cutlass.const_expr(cfg.num_softmax_scale_groups == 4): + return cutlass.Vector.from_elements( + ( + first_arr[0], + first_arr[1], + first_arr[2], + first_arr[3], + second_arr[0], + second_arr[1], + second_arr[2], + second_arr[3], + ), + Float32, + ) + elif cutlass.const_expr(cfg.num_softmax_scale_groups == 1): + return cutlass.Vector.from_elements( + (first_arr[0], second_arr[0]), + Float32, + ) + else: + return cutlass.Vector.from_elements( + ( + first_arr[0], + first_arr[1], + second_arr[0], + second_arr[1], + ), + Float32, + ) + + @cute.jit + def _stats_smem_ptr(self, stage_info: StageInfo): + """Return this thread's SMEM stats slot for the current stage. + + Softmax thread ``i`` and correction thread ``i`` address the same + row, matching the identity lane mapping of the TMEM 32x32b path. + """ + context = stage_info.context + assert context is not None and context.smem_base is not None + task_cache = _decode_gen_task_cache(stage_info) + row_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + row_elems = self._stats_smem_row_elems() + stage_elems = self._stats_smem_rows() * row_elems + base_ptr = context.smem_base.data_ptr() + self._smem_alloc.offset + view = cutlass.Array( + base_ptr, + dtype=Float32, + shape=(self.pipeline_config.num_stages * stage_elems,), + addrspace=3, + ) + elem_offset = stage_info.stage_idx * stage_elems + row_idx * row_elems + return view.subview(elem_offset).data_ptr() + + @cute.jit + def _store_stats_vector( + self, + stage_info: StageInfo, + first_arr: cutlass.Array, + second_arr: cutlass.Array, + ) -> None: + """Store one two-part softmax-stats payload for correction.""" + stats = self._stats_vector(first_arr, second_arr) + if cutlass.const_expr(self.cfg.keeps_stats_via_smem): + # SMEM handoff: the softmax-local mbarrier pipeline already + # orders this generic-proxy store against correction's load, so + # the payload needs no TMEM traffic, waits, or stats-done credit. + self._stats_smem_ptr(stage_info).store( + stats, alignment=self._stats_smem_alignment() + ) + else: + stats_ptr = self._stats_ptr(stage_info) + # Store the two stat halves as one vector payload so correction + # observes old/new or sum/new pairs from the same producer fire. + prims.tcgen05_st( + "32x32b", + stats_ptr, + stats, + ) + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def _load_stats_vector(self, stage_info: StageInfo): + """Load one two-part softmax-stats payload written by softmax.""" + if cutlass.const_expr(self.cfg.keeps_stats_via_smem): + return self._stats_smem_ptr(stage_info).load( + count=self.cfg.num_softmax_scale_groups * 2, + alignment=self._stats_smem_alignment(), + ) + stats_ptr = self._stats_ptr(stage_info) + # Reload exactly the payload shape written by _store_stats_vector; the + # TMEM view fence makes the read visible to scalar consumers. + loaded = prims.tcgen05_ld( + "32x32b", + stats_ptr, + num=self.cfg.num_softmax_scale_groups * 2, + ) + cute.arch.fence_view_async_tmem_load() + return loaded + + @cute.jit + def _return_stats_state( + self, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + inst_old_max_arr: cutlass.Array, + inst_new_max_arr: cutlass.Array, + inst_sum_arr: cutlass.Array, + ) -> tuple[object, object, object, object, object, object]: + """Return the TS task-local stats tuple in scheduler order.""" + return ( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ) + + @producer_work + @cute.jit + def store_loop_old_new_stats( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + ) -> None: + """Loop producer handoff for old/new maxima used by O rescale.""" + _ = sum_arr + # ProdWork: loop old/new max handoff is skipped for a single K/V tile because + # correction's HEAD phase only drains the resource token; there is no + # previous O accumulator to rescale yet. + skip_initial_handoff = stage_info.loop_end == Int32(1) + if not skip_initial_handoff: + # Publish old/new maxima so correction can rescale the live O stage + # before the next PV wave accumulates into it. + self._store_stats_vector(stage_info, old_max_arr, new_max_arr) + + @producer_work + @cute.jit + def store_loop_sum_new_stats( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + ) -> None: + """Loop LastIter producer handoff for final denominator and max.""" + _ = old_max_arr + # ProdWork: LastIter loop payload carries the final denominator and max for + # correction's TAIL normalization. + # Pair sum with new max because tail correction needs both to normalize + # the final O stages. + self._store_stats_vector(stage_info, sum_arr, new_max_arr) + + @producer_work + @cute.jit + def store_tail_sum_new_stats( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + ) -> None: + """Tail producer handoff for final denominator and max.""" + _ = old_max_arr + # ProdWork: Keeps-MMA-AB publishes final sums after the loop because its P path + # avoids peeling the full P body into a LastIter loop guard. + # Use the same sum/new-max payload shape as the loop LastIter handoff + # so tail correction has one consumer path for both schedule variants. + self._store_stats_vector(stage_info, sum_arr, new_max_arr) + + @consumer_work( + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ), + ) + @cute.jit + def load_head_stats( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + inst_old_max_arr: cutlass.Array, + inst_new_max_arr: cutlass.Array, + inst_sum_arr: cutlass.Array, + ) -> tuple[object, object, object, object, object, object]: + """Drain the initial stats token without reading TMEM.""" + _ = stage_info + # ConsWork: HEAD only drains the initial resource token. Reading TMEM + # here would observe stale data on the single-tile path. + return self._return_stats_state( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ) + + @consumer_work( + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ), + ) + @cute.jit + def load_loop_stats( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + inst_old_max_arr: cutlass.Array, + inst_new_max_arr: cutlass.Array, + inst_sum_arr: cutlass.Array, + ) -> tuple[object, object, object, object, object, object]: + """Load loop old/new maxima for correction's in-place O rescale.""" + cfg = self.cfg + # ConsWork: load the old/new max payload that correction uses to + # rescale the live O accumulator before the next PV MMA wave. + loaded = self._load_stats_vector(stage_info) + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + loaded_new_max = loaded[cfg.num_softmax_scale_groups + scale_idx] + new_max_arr[scale_idx] = loaded_new_max + # Loop stats carry old max for in-place O rescaling. + loaded_old_max = loaded[scale_idx] + old_max_arr[scale_idx] = loaded_old_max + inst_old_max_arr[scale_idx] = loaded_old_max + inst_new_max_arr[scale_idx] = loaded_new_max + return self._return_stats_state( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ) + + @consumer_work( + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ), + ) + @cute.jit + def load_tail_stats( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + inst_old_max_arr: cutlass.Array, + inst_new_max_arr: cutlass.Array, + inst_sum_arr: cutlass.Array, + ) -> tuple[object, object, object, object, object, object]: + """Load tail final sums/maxima and mirror them into instance stats.""" + cfg = self.cfg + # ConsWork: load the final denominator/max payload and copy it into the + # per-instance arrays consumed by correction_tail_epilogue. + loaded = self._load_stats_vector(stage_info) + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + loaded_new_max = loaded[cfg.num_softmax_scale_groups + scale_idx] + loaded_sum = loaded[scale_idx] + new_max_arr[scale_idx] = loaded_new_max + sum_arr[scale_idx] = loaded_sum + # Tail stats are mirrored into instance-local arrays so the final + # reducer can combine the two K/V instruction streams. + inst_new_max_arr[scale_idx] = loaded_new_max + inst_sum_arr[scale_idx] = loaded_sum + return self._return_stats_state( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ) + + +@dataclass(kw_only=True) +class TmemSoftmaxGlobalResource(DecodeGenResourceBase): + """FP8 softmax sum correction helper. + + FP8 P is quantized before the running sum is finalized. This resource + applies the max correction after P production and publishes the corrected + sums back through TmemS. + """ + + inst_id: Constexpr[int] = 0 + cfg: Constexpr[FmhaDecodeConfig] = None + scale_softmax_log2: Float32 = None + sum_barrier_id: Constexpr[int] = 2 + local_ref: Constexpr[MemoryResource | None] = None + p_ref: Constexpr[MemoryResource | None] = None + tmem_s_ref: Constexpr[TmemSResource] = None + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Global softmax correction uses no private SMEM allocation.""" + return [] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Global softmax correction reuses peer resource state.""" + return [] + + @producer_work + @cute.jit + def global_correction( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + new_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + ) -> None: + """Apply FP8 P-quantization denominator correction through TmemS.""" + cfg = self.cfg + if cutlass.const_expr(not cfg.use_fp8_qkv): + return + + num_scale_groups = cfg.num_softmax_scale_groups + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 2): + # ProdWork: FP8 P is quantized before the running denominator is finalized. + # Apply the same max correction as reduce_sums and publish the + # corrected sums through TmemS for the next stage. + # Gather old/new max, previous sum, and quantized-P local sum for + # one pair of scale groups. + old_max = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + new_max = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + sum_vals = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + local_sum = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + exp_scale = cutlass.Array(Float32, 2, space=cutlass.AddressSpace.rmem) + pair_width = _softmax_scale_pair_width(num_scale_groups, scale_base) + # KeepsMmaAb tracks one scale group. Pad the inactive packed-FMA + # lane with neutral values, then touch only live task-local lanes. + for pair_idx in cutlass.range_constexpr(2): + old_max[pair_idx] = Float32(0.0) + new_max[pair_idx] = Float32(0.0) + sum_vals[pair_idx] = Float32(0.0) + local_sum[pair_idx] = Float32(0.0) + for pair_idx in cutlass.range_constexpr(pair_width): + scale_idx = scale_base + pair_idx + old_max[pair_idx] = old_max_arr[scale_idx] + new_max[pair_idx] = new_max_arr[scale_idx] + sum_vals[pair_idx] = sum_arr[scale_idx] + if self.p_ref is not None: + local_sum[pair_idx] = self.tmem_s_ref.load_p_local_sum(scale_idx) + + # Convert the max delta into the online-softmax rescale factor that + # brings the previous denominator into the new max frame. + max_diff_pair = fadd2((old_max[0], old_max[1]), (-new_max[0], -new_max[1])) + scale_pair = fmul2( + (self.scale_softmax_log2, self.scale_softmax_log2), max_diff_pair + ) + for pair_idx in cutlass.range_constexpr(2): + exp_scale[pair_idx] = cute.math.exp2( + scale_pair[pair_idx], fastmath=True + ) + updated_sums = ffma2( + (exp_scale[0], exp_scale[1]), + (sum_vals[0], sum_vals[1]), + (local_sum[0], local_sum[1]), + ) + + # Publish corrected sums both through this resource return value and + # through TmemS, which reduce_sums copies for FP8. + for pair_idx in cutlass.range_constexpr(pair_width): + scale_idx = scale_base + pair_idx + sum_arr[scale_idx] = updated_sums[pair_idx] + self.tmem_s_ref.store_global_sum(scale_idx, updated_sums[pair_idx]) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py new file mode 100644 index 000000000000..4321a0d0d46e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py @@ -0,0 +1,4035 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task definitions for the FMHA decode TS kernel. + +The schedule follows the SwapsMmaAb decode pipeline shape: + +- HEAD: initial Q / K setup and first BMM1 wave +- LOOP: staggered K/V loads with Qk0 -> Pv0 -> Qk1 -> Pv1 +- TAIL: final V loads, final BMM2, final correction/output + +Each schedule is written as explicit TS resource transitions. Producer +resources use acquire/work/commit; consumer resources use wait/work/release. +The comments below name the logical pipeline step so the ordering can be read +without expanding the decorators on each resource method. +""" + +import functools +from dataclasses import dataclass, field +from typing import Any, Callable + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + TaskLocalVariable, + WorkQueue, + consumer_work, + producer_work, +) +from cutlass.experimental.task_scheduling.schedule_builder import ( + Schedule, + domain_loop, + schedule, + work_tile_loop, +) +from cutlass.experimental.task_scheduling.task import Task + +from ..stage import FmhaStage +from .fmha_decode_config import FmhaDecodeConfig +from .fmha_decode_constants import ( + KV_INST0, + KV_INST1, + KV_KIND_K, + KV_KIND_V, + KV_TILE_256_SHARED_FIFO_STAGES, +) +from .fmha_decode_resources.helpers_common import ( + ResourceVars, + _q_group_token_base, + _q_seq_bounds, + _warp_broadcast_i32, +) +from .fmha_decode_resources.helpers_kv_tile_idx import ( + _runtime_last_valid_page_idx, + _runtime_total_kv_tiles, + _sliding_window_start_idx, +) + +QBoundBinding = tuple[MemoryResource, bool] +TaskKwarg = ( + int + | bool + | cutlass.Int32 + | cute.Pointer + | FmhaDecodeConfig + | tuple[QBoundBinding, ...] + | None +) + + +def _schedule_with_optional_resources( + fn: Callable[..., None], +) -> Callable[..., Schedule]: + """Capture one schedule while omitting absent resources from its graph. + + CUTLASS ``@schedule`` intentionally accepts only concrete resources. Some + FMHA task variants have orthogonal optional resources, such as block-sparse + metadata and a persistent work queue. Filter absent slots before tracing, + then restore their named positions for the schedule body. This adapter is + entirely host-side: each captured schedule contains only present resources + and the ``None`` branches disappear while tracing. + """ + + @functools.wraps(fn) + def traced(*resource_slots: object) -> Schedule: + present_slots = tuple(slot is not None for slot in resource_slots) + resources = tuple( + slot + for slot, is_present in zip(resource_slots, present_slots, strict=True) + if is_present + ) + + @functools.wraps(fn) + def restore_slots(*resource_proxies: object) -> None: + proxy_iter = iter(resource_proxies) + restored_slots = tuple( + next(proxy_iter) if is_present else None for is_present in present_slots + ) + fn(*restored_slots) + + return schedule(restore_slots)(*resources) + + return traced + + +def _block_sparse_route_loop_domain( + route_count: cutlass.Int32, + *, + num_insts_kv: int, +) -> cutlass.Int32: + """Return LOOP iterations after HEAD reserves one candidate per instance.""" + + remaining = route_count - cutlass.Int32(num_insts_kv) + remaining = cute.math.max(remaining, cutlass.Int32(0)) + insts = cutlass.Int32(num_insts_kv) + return (remaining + insts - cutlass.Int32(1)) // insts + + +@dataclass(kw_only=True) +class ScheduleTokenThrottleResource(MemoryResource): + """Order persistent load consumption before the scheduler reuses a slot.""" + + @producer_work + @cute.jit + def publish_schedule_token(self, stage_info: StageInfo) -> None: + """Publish that the load task has consumed the current schedule token.""" + del stage_info + + @consumer_work + @cute.jit + def consume_schedule_token(self, stage_info: StageInfo) -> None: + """Wait until the load task no longer needs the current schedule token.""" + del stage_info + + +@dataclass(kw_only=True) +class SmemKvReuseCreditResource(MemoryResource): + """One-slot credit carrying the rotating KV256 exchange-stage index. + + Load publishes which drained 64-KiB physical K/V stage Correction may use + as tail scratch. The one-stage pipeline couples that payload to the same + ownership epoch: the following Load may use the other two physical stages, + but cannot publish a new alias until Correction releases this credit after + all output work completes. + """ + + cfg: cutlass.Constexpr[FmhaDecodeConfig] = None + _alloc: cutlass.Constexpr[SmemAllocation | None] = None + scratch_stage_slot: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def __post_init__(self) -> None: + """Create the routed consumer slot for one physical K/V stage.""" + assert KV_TILE_256_SHARED_FIFO_STAGES == 3, ( + "KV256 reuse-credit rotation requires exactly three shared FIFO stages" + ) + if not self.cfg.uses_rotating_kv256_exchange: + raise ValueError( + "rotating KV scratch requires persistent direct Q64/KV256 " + "with two KV instructions, one head-dimension stage, and " + "one load warp" + ) + self.scratch_stage_slot = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Physical shared-K/V stage reserved for KV256 tail exchange.", + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate the one-word stage payload guarded by this pipeline.""" + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}_scratchStage", + size_bytes=4, + alignment=4, + ) + return [self._alloc] + + @cute.jit + def _payload(self, stage_info: StageInfo) -> cutlass.Array: + """Return the natural next-stage cursor owned by this credit.""" + return cutlass.Array( + stage_info.context.smem_base.data_ptr() + self._alloc.offset, + dtype=Int32, + shape=(1,), + addrspace=3, + ) + + @cute.jit + def create_function_variables( + self, + context: ResourceContext | None = None, + ) -> ResourceVars: + """Initialize the persistent ring cursor before TS tasks start.""" + if cutlass.const_expr(context is not None and context.smem_base is not None): + payload = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=Int32, + shape=(1,), + addrspace=3, + ) + thread_idx, _, _ = cute.arch.thread_idx() + if thread_idx == Int32(0): + payload[0] = Int32(0) + return {} + + @producer_work + @cute.jit + def publish_scratch_stage(self, stage_info: StageInfo) -> None: + """Advance the persistent ring cursor and publish the drained stage.""" + num_stages = Int32(KV_TILE_256_SHARED_FIFO_STAGES) + if prims.elect_sync(): + payload = self._payload(stage_info) + # Each work commits T = 4 * (loop_end + 1) K/V transactions. + # Since 4 == 1 (mod 3), the cursor advances by loop_end + 1. + # loop_end is the resolved per-work domain, so heterogeneous + # runtime sequence lengths do not inherit a captured host bound. + payload[0] = ( + Int32(payload[0]) + stage_info.loop_end + Int32(1) + ) % num_stages + + @consumer_work(returns=scratch_stage_slot) + @cute.jit + def read_scratch_stage(self, stage_info: StageInfo) -> Int32: + """Read the alias only after the matching credit wait completes.""" + num_stages = Int32(KV_TILE_256_SHARED_FIFO_STAGES) + next_stage = Int32(self._payload(stage_info)[0]) + return (next_stage + num_stages - Int32(1)) % num_stages + + +@dataclass(kw_only=True) +class PackedDecodeWorkQueue(WorkQueue): + """CLC work queue that drops packed-Q tiles beyond a batch's Q length.""" + + cfg: cutlass.Constexpr[FmhaDecodeConfig] = field(init=False, default=None) + cu_seqlens_q: Any = field(init=False, default=None) + + def __init__( + self, + cfg: FmhaDecodeConfig, + cu_seqlens_q: cute.Pointer, + **kwargs: Any, + ) -> None: + """Attach packed-Q metadata to the shared persistent work queue.""" + super().__init__(**kwargs) + self.cfg = cfg + self.cu_seqlens_q = cu_seqlens_q + + @cute.jit + def skip_work_tile_if(self, work_tile: Any) -> cutlass.Boolean: + """Skip a fetched Q group when its first token is outside this batch.""" + q_group_cta_idx, _, b_idx = work_tile.tile_idx + q_group_idx = cutlass.Int32(q_group_cta_idx) + if cutlass.const_expr(self.cfg.use_split_kv): + q_group_idx = q_group_idx // cutlass.Int32(self.cfg.splits_kv) + _, seq_len_q = _q_seq_bounds( + self.cfg, + self.cu_seqlens_q, + cutlass.Int32(b_idx), + ) + return _q_group_token_base(self.cfg, q_group_idx) >= seq_len_q + + +def _page_offsets_consume( + smem_page_offsets: MemoryResource | None, label: str = "read_offsets" +) -> None: + """Consume the paged-KV offsets that match the next K/V TMA load.""" + if smem_page_offsets is None: + return + # ConsWait: wait until the page-offset producer staged the page IDs. + smem_page_offsets.wait() + # ConsWork: expose the page IDs to the K/V load resource. + getattr(smem_page_offsets, label)() + + +def _page_offsets_release(smem_page_offsets: MemoryResource | None) -> None: + """Release the page-offset stage after the matching K/V load is issued.""" + if smem_page_offsets is None: + return + # ConsRelease: the load resource no longer needs these page IDs. + smem_page_offsets.release() + + +def _page_offsets_produce( + smem_page_offsets: MemoryResource, label: str, section: FmhaStage +) -> None: + """Produce page IDs for one K/V load label in a schedule section.""" + # ProdAcquire: reserve a page-offset SMEM pipeline stage. + smem_page_offsets.acquire() + # ProdWork: load the page-table entries for K0/K1/V0/V1. + getattr(smem_page_offsets, label)(section=section) + # ProdCommit: publish the staged offsets to LoadTask. + smem_page_offsets.commit() + + +def _can_hold_native_split_page_window( + cfg: FmhaDecodeConfig, + smem_page_offsets: MemoryResource | None, +) -> bool: + """Return whether one native page-ID stage covers every runtime KV range.""" + if ( + smem_page_offsets is None + or not smem_page_offsets.use_native_paged_kv + or not cfg.use_split_kv + or cfg.use_sliding_window_causal + ): + return False + pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page + # Runtime ragged lengths can reduce the split-local span in + # ``num_insts_kv`` increments and therefore change every split rank's + # aligned starting page. Hold one stage only when every possible span + # both fits and evenly partitions a 32-ID window; otherwise a shorter row + # could straddle the next window even when the static maximum does not. + for local_tiles in range( + cfg.num_insts_kv, + cfg.static_local_kv_tiles + 1, + cfg.num_insts_kv, + ): + window_pages = local_tiles * pages_per_tile + if window_pages <= 0 or window_pages > 32 or 32 % window_pages != 0: + return False + return True + + +def _staged_kv_load( + resource: MemoryResource, + label: str, + section: FmhaStage, + cfg: FmhaDecodeConfig, + *, + smem_page_offsets: MemoryResource | None = None, + page_offsets_label: str = "read_offsets", + manage_page_offsets: bool = True, + cached_page_ids: Any = None, +) -> Any: + """Issue all head-dim stages for one logical K/V load. + + H256 SwapsMmaAb uses two 128-wide K/V stages, but both slices address the + same token tile and therefore the same page IDs. Keep one page-offset + consumer stage live across the complete logical K/V load. + """ + reuse_page_ids = smem_page_offsets is not None and cfg.num_head_dim_stages_kv > 1 + # Optional ConsWait/ConsWork: fetch page IDs for this logical K/V tile. + # The cached D256 path performs its ConsumerWork below while materializing + # the register array; single-stage kernels keep the original no-cache path. + if manage_page_offsets: + if reuse_page_ids: + smem_page_offsets.wait() + else: + _page_offsets_consume(smem_page_offsets, page_offsets_label) + load_label = label + if reuse_page_ids: + inst_id = KV_INST1 if label.endswith("1") else KV_INST0 + kv_kind = KV_KIND_V if "_v" in label else KV_KIND_K + cached_page_ids = smem_page_offsets.cache_page_ids( + cached_page_ids=cached_page_ids, + inst_id=inst_id, + kv_kind=kv_kind, + section=section, + ) + load_label = f"{label}_cached" + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + # ProdAcquire/ProdWork/ProdCommit: issue one K or V slice into the + # matching SMEM stage. + resource.acquire() + if cached_page_ids is None: + getattr(resource, load_label)( + section=section, head_dim_stage_idx=head_dim_stage_idx + ) + else: + getattr(resource, load_label)( + cached_page_ids=cached_page_ids, + section=section, + head_dim_stage_idx=head_dim_stage_idx, + ) + resource.commit() + # ConsRelease: every head-dimension slice consumed the staged page IDs. + if manage_page_offsets: + _page_offsets_release(smem_page_offsets) + return cached_page_ids + + +def _produce_staged_page_offsets( + smem_page_offsets: MemoryResource, + label: str, + section: FmhaStage, + cfg: FmhaDecodeConfig, +) -> None: + """Produce the page-offset stage shared by one logical K/V load.""" + _ = cfg + _page_offsets_produce(smem_page_offsets, label, section) + + +def _consume_staged_qk_mma( + smem_kv: MemoryResource, + tmem_s: MemoryResource, + aliased_p: MemoryResource, + q_desc: Any, + k_desc_label: str, + qk_mma_label: str, + section: FmhaStage, + cfg: FmhaDecodeConfig, +) -> None: + """Consume all K head-dim stages for one QK MMA wave.""" + tmem_s.acquire() + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + smem_kv.wait() + kv_desc = getattr(smem_kv, k_desc_label)() + if cutlass.const_expr( + cfg.streams_tmem_p_fragments + and head_dim_stage_idx == 0 + and (section == FmhaStage.Loop or cfg.use_persistent_scheduler) + ): + # Wait as late as possible: K staging overlaps the previous PV, + # but QK cannot overwrite the matching S/P alias until PV is done. + # Static HEAD has no previous tile; persistent HEAD may follow the + # same CTA's tail from another logical work tile and must wait. + aliased_p.wait_until_reusable_before_qk() + if cfg.uses_q_desc_ref: + getattr(tmem_s, f"{qk_mma_label}_from_q_ref")( + kv_desc=kv_desc, + head_dim_stage_idx=head_dim_stage_idx, + ) + else: + getattr(tmem_s, qk_mma_label)( + q_desc=q_desc, + kv_desc=kv_desc, + head_dim_stage_idx=head_dim_stage_idx, + ) + smem_kv.release() + tmem_s.commit() + + +def _consume_staged_pv_mma( + smem_kv: MemoryResource, + tmem_p: MemoryResource, + tmem_o: MemoryResource, + v_desc_label: str, + vp_mma_label: str, + p_desc_idx: int, + section: FmhaStage, + cfg: FmhaDecodeConfig, +) -> None: + """Consume all V head-dim stages for one PV MMA wave.""" + _ = section + if cutlass.const_expr(cfg.streams_tmem_p_fragments): + assert cfg.num_head_dim_stages_kv == 1 + fragment_label = f"{vp_mma_label}_fragment" + + # P fragment 0 is the earliest dependency. Wait for it and for the + # correction credit before holding the shared V FIFO stage. + p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=0) + tmem_o.acquire() + smem_kv.wait() + v_desc = getattr(smem_kv, v_desc_label)() + getattr(tmem_o, fragment_label)( + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=0, + ) + + # Later P fragments may become ready while the previous PV fragment is + # already executing. Keep every slot live through the complete async + # UMMA wave so the producer cannot overwrite an operand prematurely. + for fragment_idx in range(1, cfg.num_softmax_score_fragments): + p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=fragment_idx) + getattr(tmem_o, fragment_label)( + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=fragment_idx, + ) + smem_kv.release() + tmem_o.commit() + return + + tmem_p.wait() + p_desc_0, p_desc_1, p_tmem_addr_0, p_tmem_addr_1 = tmem_p.p_operands() + p_desc = p_desc_0 if p_desc_idx == KV_INST0 else p_desc_1 + p_tmem_addr = p_tmem_addr_0 if p_desc_idx == KV_INST0 else p_tmem_addr_1 + tmem_o.acquire() + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + smem_kv.wait() + v_desc = getattr(smem_kv, v_desc_label)() + getattr(tmem_o, vp_mma_label)( + v_desc_0=v_desc, + v_desc_1=v_desc, + p_desc_0=p_desc, + p_desc_1=p_desc, + p_tmem_addr_0=p_tmem_addr, + p_tmem_addr_1=p_tmem_addr, + inst_idx=p_desc_idx, + head_dim_stage_idx=head_dim_stage_idx, + ) + smem_kv.release() + tmem_o.commit() + tmem_p.release() + + +def _work_queue_tail(work_queue: WorkQueue | None, work_tile=None): + """Finish one persistent-scheduler tile after task-local work drains.""" + _ = work_tile + if work_queue is None: + return None + # ConsWait: wait for the scheduler token for this task. + work_queue.wait() + # ConsWork: advance the task-local tile cursor. + work_queue.get_and_advance_work_tile() + # ConsRelease: allow the scheduler to hand out the next tile. + work_queue.release() + return None + + +def _schedule_token_throttle_head( + schedule_token_throttle: MemoryResource | None, +) -> None: + """Let the scheduler recycle a schedule token slot once Load owns its tile state.""" + if schedule_token_throttle is None: + return + schedule_token_throttle.acquire() + schedule_token_throttle.publish_schedule_token() + schedule_token_throttle.commit() + + +def _schedule_token_throttle_tail( + schedule_token_throttle: MemoryResource | None, +) -> None: + """Pace scheduler schedule token reuse against the persistent Load task.""" + if schedule_token_throttle is None: + return + schedule_token_throttle.wait() + schedule_token_throttle.consume_schedule_token() + schedule_token_throttle.release() + + +def _decode_work_tile_schedule( + cfg: FmhaDecodeConfig, + work_queue: WorkQueue | None, + body: Callable[[], None], + non_skippable_head: Callable[[], None] | None = None, +) -> None: + """Trace one worker schedule, making only packed persistent data skippable.""" + if cfg.use_persistent_scheduler: + assert work_queue is not None + if not cfg.use_variable_seqlens_q: + with work_tile_loop(work_queue): + if non_skippable_head is not None: + non_skippable_head() + body() + _work_queue_tail(work_queue) + return + with work_tile_loop( + work_queue, + skip_if=PackedDecodeWorkQueue.skip_work_tile_if, + ) as work_tiles: + if non_skippable_head is not None: + non_skippable_head() + with work_tiles.skippable(): + body() + # WorkQueue operations stay inside the persistent loop but outside + # the skippable data region. DecodeGenTask therefore advances the + # queue exactly once for both active and inactive packed tiles. + _work_queue_tail(work_queue) + return + + # Non-persistent kernels execute one unguarded data region. + body() + _work_queue_tail(work_queue) + + +def _decode_work_tile_schedule_with_invariant_bridge( + cfg: FmhaDecodeConfig, + work_queue: WorkQueue | None, + invariant_setup: Callable[[], None], + bridge: Callable[[], Any], + active_prelude: Callable[[], None], + body: Callable[[Any], None], +) -> None: + """Trace data work around an invariant value needed across schedule phases. + + The bridge is reserved for descriptor construction that reads only + task-local descriptor metadata. It must not issue a memory operation, + barrier, or pipeline-state transition. + """ + if cfg.use_persistent_scheduler: + assert work_queue is not None + # Descriptor backing storage depends only on the task's SMEM + # allocation, so initialize it once before the persistent loop. + invariant_setup() + if not cfg.use_variable_seqlens_q: + with work_tile_loop(work_queue): + active_prelude() + invariant = bridge() + body(invariant) + _work_queue_tail(work_queue) + return + with work_tile_loop( + work_queue, + skip_if=PackedDecodeWorkQueue.skip_work_tile_if, + ) as work_tiles: + with work_tiles.skippable(): + active_prelude() + # Packed persistent QK reads Q's consumer-work stage directly, + # so it needs no descriptor task local crossing this guard. + body(None) + _work_queue_tail(work_queue) + return + + # Fixed paths initialize descriptor state, wait for Q, construct the + # descriptor, and then issue the remaining MMA work. + invariant_setup() + active_prelude() + invariant = bridge() + body(invariant) + _work_queue_tail(work_queue) + + +@cute.jit +def _load_prepared_sparse_row_warp( + row_route_offsets: cute.Pointer, + row_route_counts: cute.Pointer, + row_address: cutlass.Int32, + lane_idx: cutlass.Int32, +) -> tuple[cutlass.Int32, cutlass.Int32]: + """Load one prepared row header once per warp and broadcast it. + + DecodeGenTask.get_domain is a regular Task override, so staged control flow + lives in this JIT helper rather than in the Python method itself. + """ + + loaded_row_route_begin = cutlass.Int32(0) + loaded_route_count = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0): + loaded_row_route_begin = cutlass.Int32(row_route_offsets[row_address]) + loaded_route_count = cutlass.Int32(row_route_counts[row_address]) + row_route_begin = _warp_broadcast_i32(loaded_row_route_begin, 0) + route_count = _warp_broadcast_i32(loaded_route_count, 0) + return row_route_begin, route_count + + +class DecodeGenTask(Task): + """Decode-gen task with task-cached values used by hot resource paths.""" + + def __init__(self, **kwargs: TaskKwarg) -> None: + """Capture decode-specific task config and initialize cache slots.""" + self.cfg = kwargs.pop("cfg", None) + self.seqlens_kv = kwargs.pop("seqlens_kv", None) + self.block_table_capacity = kwargs.pop("block_table_capacity", None) + self.sparse_row_route_offsets = kwargs.pop("sparse_row_route_offsets", None) + self.sparse_row_route_counts = kwargs.pop("sparse_row_route_counts", None) + self.num_heads_kv = kwargs.pop("num_heads_kv", None) + self.max_seq_len_kv = kwargs.pop("max_seq_len_kv", cutlass.Int32(0)) + self.seq_len_q = kwargs.pop("seq_len_q", None) + self.domain_bias = kwargs.pop("domain_bias", 0) + self.q_bound_resources = kwargs.pop("q_bound_resources", ()) + super().__init__(**kwargs) + self._tmem_base_offset = cutlass.Int32(0) + self._warp_grp_thread_idx = cutlass.Int32(0) + self._local_warp_idx = cutlass.Int32(0) + self._lane_idx = cutlass.Int32(0) + self._seq_len_kv = cutlass.Int32(0) + self._kv_request_begin = cutlass.Int32(0) + self._kv_page_idx_ub = cutlass.Int32(0) + self._kv_raw_tile_base = cutlass.Int32(0) + self._kv_valid_tile_end = cutlass.Int32(0) + self._kv_window_start = cutlass.Int32(0) + # Keep the DSL loop-carried task structure stable. Persistent loops + # assign this liveness marker on every path, so it must exist before + # the first dynamic ``while`` is lowered by the stock compiler. + self.dummy = cutlass.Boolean(False) + + def init_variables(self, context: cute.Pointer | None = None) -> None: + """Initialize per-task thread and TMEM cached values.""" + super().init_variables(context) + # Cache thread identity inside the 4-warp task group for TMEM and SMEM + # resource operations. + tidx, _, _ = cute.arch.thread_idx() + warp_grp_start = cutlass.Int32((self.warp_idx // 4) * 4 * 32) + self._warp_grp_thread_idx = tidx - warp_grp_start + self._local_warp_idx = self._warp_grp_thread_idx >> cutlass.Int32(5) + self._lane_idx = self._warp_grp_thread_idx & cutlass.Int32(0x1F) + + if context is not None and context.tmem_ptr_i32 is not None: + # tcgen05_alloc provides the CTA's TMEM base through shared + # context. Broadcast it so all lanes issue TMEM operations from the + # same base column. + loaded = cutlass.Int32(context.tmem_ptr_i32.load()) + self._tmem_base_offset = _warp_broadcast_i32(loaded, 0) + + @cute.jit + def make_task_cache( + self, + ) -> tuple[ + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + cutlass.Int32, + ]: + """Return cached task-local values passed through StageInfo.""" + # The cache is threaded through StageInfo so resource methods can read + # task-local constants without recomputing them or adding more + # resource variables to every schedule edge. + return ( + self._tmem_base_offset, + self._warp_grp_thread_idx, + self._local_warp_idx, + self._lane_idx, + self._seq_len_kv, + self._kv_request_begin, + self._kv_page_idx_ub, + self._kv_raw_tile_base, + self._kv_valid_tile_end, + self._kv_window_start, + ) + + @cute.jit + def _refresh_packed_q_bounds(self, work_tile: Any) -> None: + """Attach one active CLC tile's packed Q bounds to its data resources.""" + assert self.cfg is not None + assert isinstance(self.work_queue, PackedDecodeWorkQueue) + _, _, b_idx = work_tile.tile_idx + q_token_offset, seq_len_q = _q_seq_bounds( + self.cfg, + self.work_queue.cu_seqlens_q, + cutlass.Int32(b_idx), + ) + if cutlass.const_expr(self.seq_len_q is not None): + self.seq_len_q = seq_len_q + for resource, updates_q_token_offset in self.q_bound_resources: + if cutlass.const_expr(updates_q_token_offset): + resource.q_token_offset = q_token_offset + resource.seq_len_q = seq_len_q + + @cute.jit + def _run_packed_skip_iteration( + self, + work_tile: Any, + context: ResourceContext | None = None, + ) -> None: + """Advance one inactive tile through WorkQueue bookkeeping only.""" + # Packed schedules place every data-path entry inside ``skippable()``; + # only the WorkQueue wait/advance/release tail remains outside it. Use + # a unit domain solely to populate that tail's StageInfo. In particular, + # do not call get_domain(), which would read per-batch KV metadata for a + # tile whose Q sequence is empty or overlaunched. + bookkeeping_domain = cutlass.Int32(1) + for is_skippable_head, head_entries in self._head_exec_groups: + if cutlass.const_expr(not is_skippable_head): + self._run_head_entry_group( + head_entries, + work_tile, + bookkeeping_domain, + context, + ) + for is_skippable_tail, tail_entries in self._tail_exec_groups: + if cutlass.const_expr(not is_skippable_tail): + self._run_tail_entry_group( + tail_entries, + work_tile, + bookkeeping_domain, + context, + ) + + @cute.jit + def _run_task_body_impl( + self, + work_tile: cute.Coord, + skip_work_tile: Any = None, + context: ResourceContext | None = None, + ) -> None: + """Run one ordinary task tile and synchronize attention-sink tails.""" + Task._run_task_body_impl( + self, + work_tile, + skip_work_tile, + context=context, + ) + if cutlass.const_expr( + self.cfg is not None + and self.cfg.use_persistent_scheduler + and self.cfg.use_attention_sinks + ): + # Attention sinks extend correction's tail beyond the ordinary + # task graph. Keep all persistent tasks on the same logical tile + # until that tail has drained. KV256's shared-KV alias instead + # uses a narrow Load/Correction credit in the captured schedule. + if cutlass.const_expr( + self.cfg.use_variable_seqlens_q and self.cfg.use_persistent_scheduler + ): + assert isinstance(self.work_queue, PackedDecodeWorkQueue) + q_group_cta_idx, _, b_idx = work_tile.tile_idx + _, seq_len_q = _q_seq_bounds( + self.cfg, + self.work_queue.cu_seqlens_q, + cutlass.Int32(b_idx), + ) + q_group_idx = cutlass.Int32(q_group_cta_idx) + if _q_group_token_base(self.cfg, q_group_idx) < seq_len_q: + prims.barrier_cta_sync(12, thread_count=16 * 32) + else: + prims.barrier_cta_sync(12, thread_count=16 * 32) + + @cute.jit + def _run_task_body_persistent( + self, + context: ResourceContext | None = None, + ) -> None: + """Drain inactive packed tiles before each unconditional active body.""" + use_packed_early_stop = ( + self.cfg is not None + and self.cfg.use_variable_seqlens_q + and self.cfg.use_persistent_scheduler + and self._has_skip_if + ) + if cutlass.const_expr(not use_packed_early_stop): + Task._run_task_body_persistent(self, context) + return + + assert self.work_queue is not None + work_tile = self.work_queue.initial_work_tile_info() + self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) + + self._run_pre_work_loop_entries(work_tile, context) + work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") + for resource in self.dst_resources: + if cutlass.const_expr( + resource.pipeline_config is not None + and resource.pipeline_config.advance_on_acquire + and not self._is_fork_secondary(resource) + ): + self._thread_advance_on_acquire_state(resource) + + # Consume overlaunched Q groups before entering the active loop. The + # inner loop executes only the non-skippable WorkQueue tail, so no TMA, + # descriptor, pipeline, task data, or sink barrier is issued. + while work_tile.is_valid_tile and self._should_skip_work_tile(work_tile): + self._run_packed_skip_iteration(work_tile, context) + work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") + self.dummy = cutlass.Boolean(True) + + while work_tile.is_valid_tile: + self._refresh_packed_q_bounds(work_tile) + # The tile is known active here. Running the complete schedule + # without a dynamic skip guard keeps HEAD-produced pipeline state + # in scope for LOOP and TAIL. + Task._run_task_body_impl(self, work_tile, None, context=context) + if cutlass.const_expr(self.cfg.use_attention_sinks): + prims.barrier_cta_sync(12, thread_count=16 * 32) + work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") + self.dummy = cutlass.Boolean(True) + + while work_tile.is_valid_tile and self._should_skip_work_tile(work_tile): + self._run_packed_skip_iteration(work_tile, context) + work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") + self.dummy = cutlass.Boolean(True) + + self._run_post_work_loop_entries(work_tile, context) + for resource in self.dst_resources: + if cutlass.const_expr( + resource.pipeline_config is not None + and resource is not self.work_queue + and not self._is_fork_secondary(resource) + ): + pipeline_config = resource.pipeline_config + assert pipeline_config is not None + if cutlass.const_expr(pipeline_config.advance_on_acquire): + self._thread_advance_on_acquire_state(resource) + self._producer_tail(resource) + if cutlass.const_expr( + self.work_queue in self.dst_resources + and self.work_queue.pipeline_config is not None + ): + self.work_queue.producer_tail() + self.dummy = cutlass.Boolean(True) + + def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: + """Return this task's loop domain for one static or persistent tile.""" + if self.cfg is None: + return self.domain + + # Sparse rows have a runtime-dependent number of prepared KV routes + # even when sequence lengths are static. Load their compact header + # before the fixed-dense early return below. Persistent workers pass + # the logical WorkQueue tile here, so static and CLC schedules share + # the same (q_group, head, batch) mapping. + if self.cfg.use_block_sparse: + # Validation-only TaskManagers intentionally omit prepared GMEM + # pointers and retain their configured static graph domain. + row_route_offsets = self.sparse_row_route_offsets + row_route_counts = self.sparse_row_route_counts + if row_route_offsets is None or row_route_counts is None: + return self.domain + if self.num_heads_kv is None: + raise ValueError( + "num_heads_kv is required to resolve block-sparse rows" + ) + + q_group_idx = cutlass.Int32(tile_coord[0]) + h_idx = cutlass.Int32(tile_coord[1]) + b_idx = cutlass.Int32(tile_coord[2]) + q_token_base = _q_group_token_base(self.cfg, q_group_idx) + + q_block = q_token_base // self.cfg.q_block_size + num_q_blocks = ( + self.cfg.max_seq_len_q + self.cfg.q_block_size - 1 + ) // self.cfg.q_block_size + row_address = (b_idx * self.num_heads_kv + h_idx) * num_q_blocks + q_block + + row_route_begin, route_count = _load_prepared_sparse_row_warp( + row_route_offsets, + row_route_counts, + cutlass.Int32(row_address), + self._lane_idx, + ) + + # Sparse route-span accessors share two underlying cache words + # with paged KV. Clear dense/paged-only coordinates on every + # logical tile because persistent tasks reuse the same task object. + if self.seqlens_kv is None: + self._seq_len_kv = self.max_seq_len_kv + else: + self._seq_len_kv = cutlass.Int32(self.seqlens_kv[b_idx]) + self._kv_request_begin = row_route_begin + self._kv_page_idx_ub = route_count + self._kv_raw_tile_base = cutlass.Int32(0) + self._kv_valid_tile_end = route_count + self._kv_window_start = cutlass.Int32(0) + + loop_domain = _block_sparse_route_loop_domain( + route_count, + num_insts_kv=self.cfg.num_insts_kv, + ) + return loop_domain + cutlass.Int32(self.domain_bias) + + # Resolve the sequence length for this work tile. Static-seqlen kernels + # can use the configured max length; variable-seqlen kernels read the + # batch-specific length from GMEM. + b_idx = cutlass.Int32(tile_coord[2]) + if cutlass.const_expr(self.block_table_capacity is not None): + self._kv_page_idx_ub = cutlass.Int32( + self.block_table_capacity + ) - cutlass.Int32(1) + if self.seqlens_kv is None: + seq_len_kv = cutlass.Int32(self.max_seq_len_kv) + else: + seq_len_kv = cutlass.Int32(self.seqlens_kv[b_idx]) + self._seq_len_kv = seq_len_kv + if cutlass.const_expr(self.block_table_capacity is not None): + self._kv_page_idx_ub = cute.math.min( + self._kv_page_idx_ub, + _runtime_last_valid_page_idx(self.cfg, seq_len_kv), + ) + if ( + self.seqlens_kv is None + and not self.cfg.use_split_kv + and not self.cfg.uses_runtime_q_kv_union + ): + return self.domain + tile_size_kv = cutlass.Int32(self.cfg.tile_size_kv) + + # Q-independent full-K nonsplit decode has no leading window skip. + # Resolve its runtime loop span directly and avoid repeating the general + # causal/window/split coordinate construction in every task warp. + if cutlass.const_expr( + not self.cfg.use_split_kv + and not self.cfg.uses_runtime_q_kv_union + and not self.cfg.use_sliding_window_causal + ): + total_kv_tiles = ( + seq_len_kv + tile_size_kv - cutlass.Int32(1) + ) // tile_size_kv + self._kv_window_start = cutlass.Int32(0) + self._kv_valid_tile_end = total_kv_tiles + self._kv_raw_tile_base = cutlass.Int32(0) + remaining_kv_tiles = cute.math.max( + total_kv_tiles - cutlass.Int32(self.cfg.num_insts_kv), + cutlass.Int32(0), + ) + num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) + loop_domain = ( + remaining_kv_tiles + num_insts_kv - cutlass.Int32(1) + ) // num_insts_kv + return loop_domain + cutlass.Int32(self.domain_bias) + + # Decode the logical Q tile with the configured physical split fanout, + # then derive its causal/window K union and useful runtime split prefix. + q_group_cta_idx = cutlass.Int32(tile_coord[0]) + q_group_idx = q_group_cta_idx + if self.cfg.use_split_kv: + q_group_idx = q_group_cta_idx // cutlass.Int32(self.cfg.splits_kv) + q_token_base = _q_group_token_base(self.cfg, q_group_idx) + seq_len_q = ( + cutlass.Int32(self.cfg.max_seq_len_q) + if self.seq_len_q is None + else cutlass.Int32(self.seq_len_q) + ) + self._kv_window_start = _sliding_window_start_idx( + self.cfg, + seq_len_kv, + seq_len_q, + q_token_base, + ) + skipped_tiles = self._kv_window_start // tile_size_kv + total_kv_tiles = _runtime_total_kv_tiles( + self.cfg, + seq_len_kv, + seq_len_q, + q_token_base, + ) + self._kv_valid_tile_end = skipped_tiles + total_kv_tiles + + # Split-KV groups CTAs by K/V split. The loop domain is rounded so + # each CTA in the group executes the same number of instruction pairs. + if self.cfg.use_split_kv: + # The useful active prefix is derived from this configured-fanout + # local span at kernel entry. Recomputing it in every scheduled + # task produces the same span but adds a runtime integer division + # to each task warp. + splits_kv = cutlass.Int32(self.cfg.splits_kv) + num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) + tiles_per_cta_group = splits_kv * num_insts_kv + num_groups = ( + total_kv_tiles + tiles_per_cta_group - cutlass.Int32(1) + ) // tiles_per_cta_group + total_kv_tiles = cute.math.max( + num_groups * num_insts_kv, + num_insts_kv, + ) + # Physical split coordinates and page-cache strides are laid out with + # the configured launch fanout. Runtime pruning only shortens the + # useful prefix; it must not renumber the remaining CTAs. + split_idx = cutlass.Int32(tile_coord[0]) % cutlass.Int32(self.cfg.splits_kv) + self._kv_raw_tile_base = skipped_tiles + split_idx * total_kv_tiles + else: + self._kv_raw_tile_base = skipped_tiles + remaining_kv_tiles = cute.math.max( + total_kv_tiles - cutlass.Int32(self.cfg.num_insts_kv), cutlass.Int32(0) + ) + num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) + loop_domain = ( + remaining_kv_tiles + num_insts_kv - cutlass.Int32(1) + ) // num_insts_kv + # All tasks share the MMA-loop domain; tail-only tasks add a bias. + return loop_domain + cutlass.Int32(self.domain_bias) + + +# ====================================================================== +# LoadTask — warp 13 (or warp 15 under CLC persistent), 1 warp +# K and V share a single SmemKv ring; loads alternate K and V tiles. +# HEAD: Q + K0 + K1 +# LOOP[i]: K(i+2) + V(i) +# TAIL: V(last-1) + V(last) +# Dense paged-KV consumes page offsets produced by PageTableTask. Sparse +# paged-KV instead consumes physical page IDs retained with its prepared route. +# ====================================================================== +def _resolve_and_store_sparse_route( + sparse_kv_metadata: MemoryResource | None, + section: FmhaStage, +) -> tuple[Any, Any, Any, Any] | None: + """Resolve one prepared route and retain it for the matching K/V pair.""" + + if sparse_kv_metadata is None: + return None + ( + resolved_origin0, + resolved_origin1, + resolved_atom_validity, + route_record_word_offset, + ) = sparse_kv_metadata.resolve_route(section=section) + sparse_kv_metadata.store_route( + resolved_origin0=resolved_origin0, + resolved_origin1=resolved_origin1, + resolved_atom_validity=resolved_atom_validity, + route_record_word_offset=route_record_word_offset, + ) + return ( + resolved_origin0, + resolved_origin1, + resolved_atom_validity, + route_record_word_offset, + ) + + +def _publish_sparse_softmax_route( + sparse_softmax_metadata: MemoryResource | None, + route: tuple[Any, Any, Any, Any] | None, +) -> None: + """Stage one resolved route for its paired Softmax consumer.""" + + if sparse_softmax_metadata is None: + return + assert route is not None + ( + resolved_origin0, + resolved_origin1, + resolved_atom_validity, + route_record_word_offset, + ) = route + sparse_softmax_metadata.acquire() + sparse_softmax_metadata.store_route( + resolved_origin0=resolved_origin0, + resolved_origin1=resolved_origin1, + resolved_atom_validity=resolved_atom_validity, + route_record_word_offset=route_record_word_offset, + ) + sparse_softmax_metadata.commit() + + +def create_load_task( + smem_q: MemoryResource, + smem_kv: MemoryResource, + work_queue: WorkQueue | None, + schedule_token_throttle: MemoryResource | None, + smem_kv_reuse_credit: MemoryResource | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + smem_page_offsets: MemoryResource | None = None, + sparse_kv_metadata0: MemoryResource | None = None, + sparse_kv_metadata1: MemoryResource | None = None, + sparse_softmax_metadata0: MemoryResource | None = None, + sparse_softmax_metadata1: MemoryResource | None = None, + warp_idx: int | None = None, + num_warps: int | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the shared-KV load task and optional page-offset dependency.""" + hold_page_window = _can_hold_native_split_page_window(cfg, smem_page_offsets) + + def load_schedule_body( + smem_q: MemoryResource, + smem_kv: MemoryResource, + smem_page_offsets: MemoryResource | None, + schedule_token_throttle: MemoryResource | None, + smem_kv_reuse_credit: MemoryResource | None, + sparse_kv_metadata0: MemoryResource | None = None, + sparse_kv_metadata1: MemoryResource | None = None, + sparse_softmax_metadata0: MemoryResource | None = None, + sparse_softmax_metadata1: MemoryResource | None = None, + ) -> None: + """Build the shared-KV load cadence for HEAD, LOOP, and TAIL.""" + smem_q.init_load_state() + smem_kv.init_load_state() + for sparse_resource in ( + sparse_kv_metadata0, + sparse_kv_metadata1, + sparse_softmax_metadata0, + sparse_softmax_metadata1, + ): + if sparse_resource is not None: + sparse_resource.init_load_state() + cached_page_ids = None + if smem_page_offsets is not None: + if cfg.num_head_dim_stages_kv > 1: + cached_page_ids = smem_page_offsets.init_cached_read_state() + else: + smem_page_offsets.init_read_state() + + def _kv_load(label: str, section: FmhaStage) -> None: + """Run one staged shared-KV load with optional page-offset handoff.""" + nonlocal cached_page_ids + cached_page_ids = _staged_kv_load( + smem_kv, + label, + section, + cfg, + smem_page_offsets=smem_page_offsets, + manage_page_offsets=not hold_page_window, + cached_page_ids=cached_page_ids, + ) + + # HEAD: load Q once, then prefetch the first two K tiles. + smem_q.acquire() + smem_q.tma_load() + smem_q.commit() + if hold_page_window: + # The native K/V caches share one page table. Keep its single + # 32-ID consumer stage live across every K/V and head-dimension + # load owned by this split CTA, matching the reference page-window + # lifetime and avoiding redundant pipeline handoffs. + if cfg.num_head_dim_stages_kv > 1: + smem_page_offsets.wait() + else: + _page_offsets_consume(smem_page_offsets) + if sparse_kv_metadata0 is None: + for label in ("load_k0", "load_k1"): + _kv_load(label, FmhaStage.Head) + else: + route0 = _resolve_and_store_sparse_route( + sparse_kv_metadata0, FmhaStage.Head + ) + _kv_load("load_k0", FmhaStage.Head) + route1 = _resolve_and_store_sparse_route( + sparse_kv_metadata1, FmhaStage.Head + ) + _kv_load("load_k1", FmhaStage.Head) + # Issue both K tiles before either metadata FIFO can backpressure + # the load warp, matching the split-resource sparse cadence. + _publish_sparse_softmax_route(sparse_softmax_metadata0, route0) + _publish_sparse_softmax_route(sparse_softmax_metadata1, route1) + if smem_kv_reuse_credit is not None: + # K0/K1 occupy the two stages disjoint from the previous work's + # scratch. Acquire only before issuing the third K/V transaction. + smem_kv_reuse_credit.acquire() + + # LOOP: each iter prefetches the full ``num_insts_kv`` K/V pair set. + # When P aliases the consumed S columns, MMA must consume each V/P pair + # before the following same-instance QK overwrites S. Keep the + # shared-ring producer order identical to that consumer order. + loop_labels = ( + ("load_v0", "load_k0", "load_v1", "load_k1") + if cfg.uses_two_inst_tmem_p + else ("load_k0", "load_v0", "load_k1", "load_v1") + ) + with domain_loop(0, domain, 1, unroll=1): + if sparse_kv_metadata0 is None: + for label in loop_labels: + _kv_load(label, FmhaStage.Loop) + else: + # Follow the dense KV256 stage order exactly. Each V consumes + # its retained route before the matching K label replaces it. + loop_routes = [] + for label in loop_labels: + route = None + sparse_softmax_metadata = None + if label == "load_k0": + route = _resolve_and_store_sparse_route( + sparse_kv_metadata0, FmhaStage.Loop + ) + sparse_softmax_metadata = sparse_softmax_metadata0 + elif label == "load_k1": + route = _resolve_and_store_sparse_route( + sparse_kv_metadata1, FmhaStage.Loop + ) + sparse_softmax_metadata = sparse_softmax_metadata1 + _kv_load(label, FmhaStage.Loop) + if route is not None: + loop_routes.append((sparse_softmax_metadata, route)) + for sparse_softmax_metadata, route in loop_routes: + _publish_sparse_softmax_route(sparse_softmax_metadata, route) + + # TAIL: after no more future K tiles are needed, load the final two V + # tiles consumed by the final BMM2 calls. + for label in ("load_v0", "load_v1"): + _kv_load(label, FmhaStage.Tail) + if smem_kv_reuse_credit is not None: + # Publish the physical stage drained by this work together with + # the ownership token consumed by the correction tail. + smem_kv_reuse_credit.publish_scratch_stage() + smem_kv_reuse_credit.commit() + if hold_page_window: + _page_offsets_release(smem_page_offsets) + + @_schedule_with_optional_resources + def load_schedule( + smem_q: MemoryResource, + smem_kv: MemoryResource, + smem_page_offsets: MemoryResource | None, + sparse_kv_metadata0: MemoryResource | None, + sparse_kv_metadata1: MemoryResource | None, + sparse_softmax_metadata0: MemoryResource | None, + sparse_softmax_metadata1: MemoryResource | None, + work_queue: WorkQueue | None, + schedule_token_throttle: MemoryResource | None, + smem_kv_reuse_credit: MemoryResource | None, + ) -> None: + """Schedule shared-KV loads with only the resources in this profile.""" + + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: load_schedule_body( + smem_q, + smem_kv, + smem_page_offsets, + schedule_token_throttle, + smem_kv_reuse_credit, + sparse_kv_metadata0, + sparse_kv_metadata1, + sparse_softmax_metadata0, + sparse_softmax_metadata1, + ), + lambda: _schedule_token_throttle_head(schedule_token_throttle), + ) + + sparse_resources = ( + sparse_kv_metadata0, + sparse_kv_metadata1, + sparse_softmax_metadata0, + sparse_softmax_metadata1, + ) + sparse_resources_present = tuple( + resource is not None for resource in sparse_resources + ) + has_sparse_metadata = any(sparse_resources_present) + if has_sparse_metadata and not all(sparse_resources_present): + raise ValueError("shared sparse K/V requires both route/Softmax pairs") + if has_sparse_metadata and smem_page_offsets is not None: + raise ValueError("block-sparse and paged-KV cannot share a load task") + + captured_schedule = load_schedule( + smem_q, + smem_kv, + smem_page_offsets, + sparse_kv_metadata0, + sparse_kv_metadata1, + sparse_softmax_metadata0, + sparse_softmax_metadata1, + work_queue, + schedule_token_throttle, + smem_kv_reuse_credit, + ) + src = [] + for sparse_kv_metadata in (sparse_kv_metadata0, sparse_kv_metadata1): + if sparse_kv_metadata is not None: + src.append(sparse_kv_metadata) + if smem_page_offsets is not None: + src.append(smem_page_offsets) + if work_queue is not None: + src.append(work_queue) + dst = [smem_q, smem_kv] + for sparse_resource in sparse_resources: + if sparse_resource is not None and sparse_resource not in dst: + dst.append(sparse_resource) + if schedule_token_throttle is not None: + dst.append(schedule_token_throttle) + if smem_kv_reuse_credit is not None: + dst.append(smem_kv_reuse_credit) + return task_class( + src_resources=src, + dst_resources=dst, + q_bound_resources=((smem_q, True), (smem_kv, False)), + cfg=cfg, + warp_idx=cfg.load_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.load_num_warps if num_warps is None else num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_task_num_registers, + name="LoadTask", + **kw, + ) + + +def create_page_offsets_task( + smem_page_offsets: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + warp_idx: int | None = None, + num_warps: int | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Prefetch page-table entries that LoadTask consumes for paged KV. + + The schedule matches LoadTask's K/V cadence exactly so the page-offsets + ring and the SmemKv ring stay aligned. + """ + hold_page_window = _can_hold_native_split_page_window(cfg, smem_page_offsets) + + def page_offsets_schedule_body( + smem_page_offsets: MemoryResource, + ) -> None: + """Schedule page-offset prefetches for the shared-KV load cadence.""" + smem_page_offsets.init_load_state() + if hold_page_window: + # K0's aligned 32-ID window covers every contiguous tile assigned + # to this split CTA, and the native table uses those IDs for V too. + _page_offsets_produce(smem_page_offsets, "load_k0", FmhaStage.Head) + # Preserve the runtime domain contract even though this fast path + # needs no per-iteration page-window work. + with domain_loop(0, domain, 1, unroll=1): + pass + return + # Page offsets are staged by a separate producer so the load warp can + # issue K/V TMA copies without reading page tables itself. + # HEAD: produce page IDs for the two prefetched K tiles. + _produce_staged_page_offsets(smem_page_offsets, "load_k0", FmhaStage.Head, cfg) + _produce_staged_page_offsets(smem_page_offsets, "load_k1", FmhaStage.Head, cfg) + + # LOOP: mirror LoadTask's K/V production cadence exactly. + loop_labels = ( + ("load_v0", "load_k0", "load_v1", "load_k1") + if cfg.uses_two_inst_tmem_p + else ("load_k0", "load_v0", "load_k1", "load_v1") + ) + with domain_loop(0, domain, 1, unroll=1): + for label in loop_labels: + _produce_staged_page_offsets( + smem_page_offsets, label, FmhaStage.Loop, cfg + ) + + # TAIL: produce page IDs for the final two V loads. + for label in ("load_v0", "load_v1"): + _produce_staged_page_offsets(smem_page_offsets, label, FmhaStage.Tail, cfg) + + @schedule + def page_offsets_schedule( + smem_page_offsets: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Wrap page-offset data work in packed persistent skip handling.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: page_offsets_schedule_body(smem_page_offsets), + ) + + captured_schedule = ( + page_offsets_schedule(smem_page_offsets) + if work_queue is None + else page_offsets_schedule(smem_page_offsets, work_queue) + ) + src = [] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[smem_page_offsets], + q_bound_resources=((smem_page_offsets, False),), + cfg=cfg, + warp_idx=cfg.page_offsets_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.page_offsets_num_warps if num_warps is None else num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_task_num_registers, + name="PageTableTask", + **kw, + ) + + +def create_page_offsets_task_split_kv( + smem_page_offsets_k: MemoryResource, + smem_page_offsets_v: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + warp_idx: int | None = None, + num_warps: int | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Prefetch split-KV K/V page windows with independent consumer state.""" + + def page_offsets_schedule_body( + smem_page_offsets_k: MemoryResource, + smem_page_offsets_v: MemoryResource, + ) -> None: + """Publish one page-offset stage for every paired K/V tile.""" + smem_page_offsets_k.init_load_state() + smem_page_offsets_v.init_load_state() + + # HEAD: publish the initial K0/K1 pair as two independent stages. + _page_offsets_produce(smem_page_offsets_k, "load_k0", FmhaStage.Head) + _page_offsets_produce(smem_page_offsets_k, "load_k1", FmhaStage.Head) + + # LOOP: mirror LoadTask's cross-resource consumption order exactly. + with domain_loop(0, domain, 1, unroll=1): + _page_offsets_produce(smem_page_offsets_v, "load_v0", FmhaStage.Loop) + _page_offsets_produce(smem_page_offsets_k, "load_k0", FmhaStage.Loop) + _page_offsets_produce(smem_page_offsets_v, "load_v1", FmhaStage.Loop) + _page_offsets_produce(smem_page_offsets_k, "load_k1", FmhaStage.Loop) + + # TAIL: drain V0/V1 through distinct page-offset stages. + _page_offsets_produce(smem_page_offsets_v, "load_v0", FmhaStage.Tail) + _page_offsets_produce(smem_page_offsets_v, "load_v1", FmhaStage.Tail) + + @schedule + def page_offsets_schedule( + smem_page_offsets_k: MemoryResource, + smem_page_offsets_v: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Wrap split page-offset work in packed persistent skip handling.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: page_offsets_schedule_body( + smem_page_offsets_k, + smem_page_offsets_v, + ), + ) + + schedule_result = ( + page_offsets_schedule(smem_page_offsets_k, smem_page_offsets_v) + if work_queue is None + else page_offsets_schedule(smem_page_offsets_k, smem_page_offsets_v, work_queue) + ) + src = [] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[smem_page_offsets_k, smem_page_offsets_v], + q_bound_resources=( + (smem_page_offsets_k, False), + (smem_page_offsets_v, False), + ), + cfg=cfg, + warp_idx=cfg.page_offsets_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.page_offsets_num_warps if num_warps is None else num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_task_num_registers, + name="PageTableTask", + **kw, + ) + + +def create_page_offsets_task_one_inst_qkv( + smem_page_offsets: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + warp_idx: int | None = None, + num_warps: int | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Prefetch page-table entries for the one-inst keepsMmaAb QKV path.""" + + def page_offsets_schedule_body( + smem_page_offsets: MemoryResource, + ) -> None: + """Schedule page-offset prefetches for the one-inst QKV load cadence.""" + smem_page_offsets.init_load_state() + + _page_offsets_produce(smem_page_offsets, "load_k0", FmhaStage.Head) + with domain_loop(0, domain, 1, unroll=1): + for label in ("load_k0", "load_v0"): + _page_offsets_produce(smem_page_offsets, label, FmhaStage.Loop) + _page_offsets_produce(smem_page_offsets, "load_v0", FmhaStage.Tail) + + @schedule + def page_offsets_schedule( + smem_page_offsets: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Wrap one-inst page-offset work in packed persistent skip handling.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: page_offsets_schedule_body(smem_page_offsets), + ) + + schedule_result = ( + page_offsets_schedule(smem_page_offsets) + if work_queue is None + else page_offsets_schedule(smem_page_offsets, work_queue) + ) + src = [] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[smem_page_offsets], + q_bound_resources=((smem_page_offsets, False),), + cfg=cfg, + warp_idx=cfg.page_offsets_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.page_offsets_num_warps if num_warps is None else num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_task_num_registers, + name="PageTableTask", + **kw, + ) + + +def create_load_task_split_kv( + smem_q: MemoryResource | None, + smem_k0: MemoryResource | None, + smem_k1: MemoryResource | None, + smem_v0: MemoryResource | None, + smem_v1: MemoryResource | None, + work_queue: WorkQueue | None, + schedule_token_throttle: MemoryResource | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + smem_page_offsets: MemoryResource | None = None, + smem_page_offsets_v: MemoryResource | None = None, + sparse_kv_metadata0: MemoryResource | None = None, + sparse_kv_metadata1: MemoryResource | None = None, + sparse_softmax_metadata0: MemoryResource | None = None, + sparse_softmax_metadata1: MemoryResource | None = None, + warp_idx: int | None = None, + num_warps: int | None = None, + task_name: str = "LoadTask", + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create loads for independent K0/K1/V0/V1 resources. + + Here ``split`` describes those resources, not split-KV reduction. + """ + if schedule_token_throttle is not None and work_queue is None: + raise ValueError("schedule-token throttle requires a work queue") + if smem_page_offsets_v is not None and smem_page_offsets is None: + raise ValueError("V page offsets require K page offsets") + if (smem_k0 is None) != (smem_v0 is None) or (smem_k1 is None) != (smem_v1 is None): + raise ValueError("each split K resource requires its matching V resource") + if smem_k0 is None and smem_k1 is None: + raise ValueError("at least one split K/V instance is required") + + def load_schedule_body( + smem_q: MemoryResource | None, + smem_k0: MemoryResource | None, + smem_k1: MemoryResource | None, + smem_v0: MemoryResource | None, + smem_v1: MemoryResource | None, + sparse_kv_metadata0: MemoryResource | None, + sparse_kv_metadata1: MemoryResource | None, + sparse_softmax_metadata0: MemoryResource | None, + sparse_softmax_metadata1: MemoryResource | None, + smem_page_offsets: MemoryResource | None, + smem_page_offsets_v: MemoryResource | None = None, + schedule_token_throttle: MemoryResource | None = None, + ) -> None: + """Build the split-resource K/V load cadence for all schedule phases.""" + if smem_q is not None: + smem_q.init_load_state() + active_instances = ( + ( + smem_k0, + smem_v0, + sparse_kv_metadata0, + sparse_softmax_metadata0, + "load_k0", + "load_v0", + ), + ( + smem_k1, + smem_v1, + sparse_kv_metadata1, + sparse_softmax_metadata1, + "load_k1", + "load_v1", + ), + ) + # Preserve the original full-resource lowering order while allowing a + # per-instance task to omit the other stream's resources. + for smem_k, _, _, _, _, _ in active_instances: + if smem_k is not None: + smem_k.init_load_state() + for _, smem_v, _, _, _, _ in active_instances: + if smem_v is not None: + smem_v.init_load_state() + for _, _, sparse_kv_metadata, _, _, _ in active_instances: + if sparse_kv_metadata is not None: + sparse_kv_metadata.init_load_state() + for _, _, _, sparse_softmax_metadata, _, _ in active_instances: + if sparse_softmax_metadata is not None: + sparse_softmax_metadata.init_load_state() + if smem_page_offsets is not None: + smem_page_offsets.init_read_state() + if smem_page_offsets_v is not None: + smem_page_offsets_v.init_read_state() + + smem_page_offsets_k = smem_page_offsets + smem_page_offsets_v_local = ( + smem_page_offsets_v + if smem_page_offsets_v is not None + else smem_page_offsets + ) + + def load_tile( + resource: MemoryResource, + label: str, + offsets: MemoryResource | None, + section: FmhaStage, + ) -> None: + """Acquire, load all head-dim stages, and release page offsets.""" + _page_offsets_consume(offsets, label.replace("load", "read_offsets")) + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + resource.acquire() + getattr(resource, label)( + section=section, head_dim_stage_idx=head_dim_stage_idx + ) + resource.commit() + _page_offsets_release(offsets) + + if smem_q is not None: + smem_q.acquire() + smem_q.tma_load() + smem_q.commit() + + head_routes = [] + for ( + smem_k, + _, + sparse_kv_metadata, + sparse_softmax_metadata, + load_k, + _, + ) in active_instances: + if smem_k is None: + continue + route = _resolve_and_store_sparse_route(sparse_kv_metadata, FmhaStage.Head) + load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Head) + head_routes.append((sparse_softmax_metadata, route)) + # In the combined task, preserve both K issues ahead of Softmax + # backpressure. A per-instance task naturally stages its sole route. + for sparse_softmax_metadata, route in head_routes: + _publish_sparse_softmax_route(sparse_softmax_metadata, route) + + with domain_loop(0, domain, 1, unroll=1): + # V consumes the retained route before the next K overwrites it. + loop_routes = [] + for ( + smem_k, + smem_v, + sparse_kv_metadata, + sparse_softmax_metadata, + load_k, + load_v, + ) in active_instances: + if smem_k is None: + continue + assert smem_v is not None + load_tile( + smem_v, + load_v, + smem_page_offsets_v_local, + FmhaStage.Loop, + ) + route = _resolve_and_store_sparse_route( + sparse_kv_metadata, FmhaStage.Loop + ) + load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Loop) + loop_routes.append((sparse_softmax_metadata, route)) + for sparse_softmax_metadata, route in loop_routes: + _publish_sparse_softmax_route(sparse_softmax_metadata, route) + + for _, smem_v, _, _, _, load_v in active_instances: + if smem_v is not None: + load_tile( + smem_v, + load_v, + smem_page_offsets_v_local, + FmhaStage.Tail, + ) + + @_schedule_with_optional_resources + def load_schedule( + smem_q: MemoryResource | None, + smem_k0: MemoryResource | None, + smem_k1: MemoryResource | None, + smem_v0: MemoryResource | None, + smem_v1: MemoryResource | None, + sparse_kv_metadata0: MemoryResource | None, + sparse_kv_metadata1: MemoryResource | None, + sparse_softmax_metadata0: MemoryResource | None, + sparse_softmax_metadata1: MemoryResource | None, + smem_page_offsets: MemoryResource | None, + smem_page_offsets_v: MemoryResource | None, + work_queue: WorkQueue | None, + schedule_token_throttle: MemoryResource | None, + ) -> None: + """Schedule split K/V loads with only the supplied optional resources.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: load_schedule_body( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + sparse_kv_metadata0, + sparse_kv_metadata1, + sparse_softmax_metadata0, + sparse_softmax_metadata1, + smem_page_offsets, + smem_page_offsets_v, + schedule_token_throttle, + ), + lambda: _schedule_token_throttle_head(schedule_token_throttle), + ) + + for smem_k, sparse_kv_metadata, sparse_softmax_metadata in ( + (smem_k0, sparse_kv_metadata0, sparse_softmax_metadata0), + (smem_k1, sparse_kv_metadata1, sparse_softmax_metadata1), + ): + if smem_k is None and ( + sparse_kv_metadata is not None or sparse_softmax_metadata is not None + ): + raise ValueError("inactive K/V instances cannot own sparse metadata") + if sparse_softmax_metadata is not None and sparse_kv_metadata is None: + raise ValueError("Softmax sparse metadata requires retained KV metadata") + if sparse_kv_metadata0 is not None or sparse_kv_metadata1 is not None: + if smem_page_offsets is not None or smem_page_offsets_v is not None: + raise ValueError( + "block-sparse and paged-KV load resources cannot be combined" + ) + + schedule_result = load_schedule( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + sparse_kv_metadata0, + sparse_kv_metadata1, + sparse_softmax_metadata0, + sparse_softmax_metadata1, + smem_page_offsets, + smem_page_offsets_v, + work_queue, + schedule_token_throttle, + ) + src = [] + # Route resolution is ConsumerWork and K/V retention is ProducerWork on + # the same pipeline-free resource, so register both sides of that route. + for sparse_kv_metadata in (sparse_kv_metadata0, sparse_kv_metadata1): + if sparse_kv_metadata is not None: + src.append(sparse_kv_metadata) + if smem_page_offsets is not None: + src.append(smem_page_offsets) + if smem_page_offsets_v is not None: + src.append(smem_page_offsets_v) + if work_queue is not None: + src.append(work_queue) + dst = [ + resource + for resource in (smem_q, smem_k0, smem_k1, smem_v0, smem_v1) + if resource is not None + ] + for sparse_resource in ( + sparse_kv_metadata0, + sparse_kv_metadata1, + sparse_softmax_metadata0, + sparse_softmax_metadata1, + ): + if sparse_resource is not None: + dst.append(sparse_resource) + if schedule_token_throttle is not None: + dst.append(schedule_token_throttle) + q_bound_resources = tuple( + (resource, resource is smem_q) + for resource in (smem_q, smem_k0, smem_k1, smem_v0, smem_v1) + if resource is not None + ) + return task_class( + src_resources=src, + dst_resources=dst, + q_bound_resources=q_bound_resources, + cfg=cfg, + warp_idx=cfg.load_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.load_num_warps if num_warps is None else num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_task_num_registers, + name=task_name, + **kw, + ) + + +def create_block_sparse_load_tasks_per_inst( + smem_q: MemoryResource, + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + work_queue: WorkQueue | None, + schedule_token_throttle: MemoryResource | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + sparse_kv_metadata0: MemoryResource, + sparse_kv_metadata1: MemoryResource, + sparse_softmax_metadata0: MemoryResource, + sparse_softmax_metadata1: MemoryResource, + warp_indices: tuple[int, int], + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> tuple[Task, Task]: + """Assign each sparse K/V instruction stream to an independent load warp. + + Load0 alone owns Q and the persistent schedule-token throttle. Both tasks + consume the same logical work tile, while their K/V and sparse-metadata + pipelines remain disjoint. + """ + + if not cfg.use_block_sparse: + raise ValueError("per-instance load tasks require block-sparse metadata") + + load0 = create_load_task_split_kv( + smem_q, + smem_k0, + None, + smem_v0, + None, + work_queue, + schedule_token_throttle, + cfg, + domain=domain, + sparse_kv_metadata0=sparse_kv_metadata0, + sparse_softmax_metadata0=sparse_softmax_metadata0, + warp_idx=warp_indices[0], + task_name="LoadTask0", + task_class=task_class, + **kw, + ) + load1 = create_load_task_split_kv( + None, + None, + smem_k1, + None, + smem_v1, + work_queue, + None, + cfg, + domain=domain, + sparse_kv_metadata1=sparse_kv_metadata1, + sparse_softmax_metadata1=sparse_softmax_metadata1, + warp_idx=warp_indices[1], + task_name="LoadTask1", + task_class=task_class, + **kw, + ) + return load0, load1 + + +def create_load_task_one_inst_qkv( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + work_queue: WorkQueue | None, + schedule_token_throttle: MemoryResource | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + smem_page_offsets: MemoryResource | None = None, + warp_idx: int | None = None, + num_warps: int | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Load schedule for the one-inst keepsMmaAb QKV path.""" + + def load_schedule_body( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + smem_page_offsets: MemoryResource | None, + schedule_token_throttle: MemoryResource | None, + ) -> None: + """Build the one-inst Q/K/V load cadence.""" + smem_q.init_load_state() + smem_k.init_load_state() + smem_v.init_load_state() + if smem_page_offsets is not None: + smem_page_offsets.init_read_state() + + def load_tile(resource: MemoryResource, label: str, section: FmhaStage) -> None: + """Acquire, load all head-dim stages, and release one page window.""" + _page_offsets_consume( + smem_page_offsets, label.replace("load", "read_offsets") + ) + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + resource.acquire() + getattr(resource, label)( + section=section, head_dim_stage_idx=head_dim_stage_idx + ) + resource.commit() + _page_offsets_release(smem_page_offsets) + + smem_q.acquire() + smem_q.tma_load() + smem_q.commit() + load_tile(smem_k, "load_k0", FmhaStage.Head) + + with domain_loop(0, domain, 1, unroll=1): + load_tile(smem_k, "load_k0", FmhaStage.Loop) + load_tile(smem_v, "load_v0", FmhaStage.Loop) + + load_tile(smem_v, "load_v0", FmhaStage.Tail) + + @schedule + def load_schedule( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + work_queue: WorkQueue | None = None, + schedule_token_throttle: MemoryResource | None = None, + ) -> None: + """Schedule one-inst Q/K/V loads without page-offset resources.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: load_schedule_body( + smem_q, smem_k, smem_v, None, schedule_token_throttle + ), + lambda: _schedule_token_throttle_head(schedule_token_throttle), + ) + + @schedule + def load_page_offsets_schedule( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + smem_page_offsets: MemoryResource, + work_queue: WorkQueue | None = None, + schedule_token_throttle: MemoryResource | None = None, + ) -> None: + """Schedule one-inst Q/K/V loads with paged-KV offsets.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: load_schedule_body( + smem_q, + smem_k, + smem_v, + smem_page_offsets, + schedule_token_throttle, + ), + lambda: _schedule_token_throttle_head(schedule_token_throttle), + ) + + if smem_page_offsets is None: + if work_queue is None: + schedule_result = load_schedule(smem_q, smem_k, smem_v) + elif schedule_token_throttle is None: + schedule_result = load_schedule(smem_q, smem_k, smem_v, work_queue) + else: + schedule_result = load_schedule( + smem_q, smem_k, smem_v, work_queue, schedule_token_throttle + ) + else: + if work_queue is None: + schedule_result = load_page_offsets_schedule( + smem_q, smem_k, smem_v, smem_page_offsets + ) + elif schedule_token_throttle is None: + schedule_result = load_page_offsets_schedule( + smem_q, + smem_k, + smem_v, + smem_page_offsets, + work_queue, + ) + else: + schedule_result = load_page_offsets_schedule( + smem_q, + smem_k, + smem_v, + smem_page_offsets, + work_queue, + schedule_token_throttle, + ) + src = [] + if smem_page_offsets is not None: + src.append(smem_page_offsets) + if work_queue is not None: + src.append(work_queue) + dst = [smem_q, smem_k, smem_v] + if schedule_token_throttle is not None: + dst.append(schedule_token_throttle) + return task_class( + src_resources=src, + dst_resources=dst, + q_bound_resources=((smem_q, True), (smem_k, False), (smem_v, False)), + cfg=cfg, + warp_idx=cfg.load_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.load_num_warps if num_warps is None else num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_task_num_registers, + name="LoadTask", + **kw, + ) + + +def create_mma_task_split_kv( + smem_q: MemoryResource, + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + tmem_stats_done0: MemoryResource | None = None, + tmem_stats_done1: MemoryResource | None = None, + warp_idx: int | None = None, + num_warps: int | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the MMA task for split K/V resources.""" + + def mma_schedule_body( + smem_q: MemoryResource, + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + tmem_stats_done0: MemoryResource | None, + tmem_stats_done1: MemoryResource | None, + q_desc: Any, + ) -> None: + """Schedule split-resource QK and PV waves across HEAD/LOOP/TAIL.""" + + def qk_mma( + smem_kv: MemoryResource, + tmem_s: MemoryResource, + tmem_stats_done: MemoryResource | None, + q_desc, + qk_mma_label: str, + section: FmhaStage, + ) -> None: + """Issue one scheduled QK wave using the selected phase work.""" + _ = section + if tmem_stats_done is not None: + tmem_stats_done.acquire() + tmem_s.acquire() + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + smem_kv.wait() + kv_desc = smem_kv.kv_desc() + if cfg.uses_q_desc_ref: + getattr(tmem_s, f"{qk_mma_label}_from_q_ref")( + kv_desc=kv_desc, + head_dim_stage_idx=head_dim_stage_idx, + ) + else: + getattr(tmem_s, qk_mma_label)( + q_desc=q_desc, + kv_desc=kv_desc, + head_dim_stage_idx=head_dim_stage_idx, + ) + smem_kv.release() + tmem_s.commit() + if tmem_stats_done is not None: + tmem_stats_done.commit() + + def pv_mma( + smem_kv: MemoryResource, + tmem_p: MemoryResource, + vp_mma_label: str, + inst_idx: int, + section: FmhaStage, + ) -> None: + """Issue one scheduled PV wave using the selected phase work.""" + _ = section + tmem_p.wait() + p_desc_0, p_desc_1, p_tmem_addr_0, p_tmem_addr_1 = tmem_p.p_operands() + tmem_o.acquire() + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + smem_kv.wait() + v_desc = smem_kv.v_desc() + getattr(tmem_o, vp_mma_label)( + v_desc_0=v_desc, + v_desc_1=v_desc, + p_desc_0=p_desc_0, + p_desc_1=p_desc_1, + p_tmem_addr_0=p_tmem_addr_0, + p_tmem_addr_1=p_tmem_addr_1, + inst_idx=inst_idx, + head_dim_stage_idx=head_dim_stage_idx, + ) + smem_kv.release() + tmem_o.commit() + tmem_p.release() + + qk_mma( + smem_k0, + tmem_s0, + tmem_stats_done0, + q_desc, + "qk_mma_head", + FmhaStage.Head, + ) + qk_mma( + smem_k1, + tmem_s1, + tmem_stats_done1, + q_desc, + "qk_mma_head", + FmhaStage.Head, + ) + + with domain_loop(0, domain, 1, unroll=1): + pv_mma(smem_v0, smem_p0, "vp_mma_loop", KV_INST0, FmhaStage.Loop) + qk_mma( + smem_k0, + tmem_s0, + tmem_stats_done0, + q_desc, + "qk_mma_loop", + FmhaStage.Loop, + ) + pv_mma(smem_v1, smem_p1, "vp_mma_loop", KV_INST1, FmhaStage.Loop) + qk_mma( + smem_k1, + tmem_s1, + tmem_stats_done1, + q_desc, + "qk_mma_loop", + FmhaStage.Loop, + ) + + pv_mma(smem_v0, smem_p0, "vp_mma_tail", KV_INST0, FmhaStage.Tail) + pv_mma(smem_v1, smem_p1, "vp_mma_tail", KV_INST1, FmhaStage.Tail) + smem_q.release() + + def mma_schedule_prelude( + smem_q: MemoryResource, + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + ) -> None: + """Initialize invariant split-resource descriptor slots.""" + smem_q.init_descriptor_state() + smem_k0.init_descriptor_state() + smem_k1.init_descriptor_state() + smem_v0.init_descriptor_state() + smem_v1.init_descriptor_state() + smem_p0.init_descriptor_state() + smem_p1.init_descriptor_state() + + def run_mma_schedule( + smem_q: MemoryResource, + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + tmem_stats_done0: MemoryResource | None, + tmem_stats_done1: MemoryResource | None, + work_queue: WorkQueue | None, + ) -> None: + """Wrap split-resource MMA work with optional stats lifetime gates.""" + _decode_work_tile_schedule_with_invariant_bridge( + cfg, + work_queue, + lambda: mma_schedule_prelude( + smem_q, smem_k0, smem_k1, smem_v0, smem_v1, smem_p0, smem_p1 + ), + lambda: smem_q.q_desc(), + lambda: smem_q.wait(), + lambda q_desc: mma_schedule_body( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + tmem_stats_done0, + tmem_stats_done1, + q_desc, + ), + ) + + @schedule + def mma_schedule( + smem_q: MemoryResource, + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Capture the Swaps split-resource MMA schedule.""" + run_mma_schedule( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + None, + None, + work_queue, + ) + + @schedule + def mma_keeps_schedule( + smem_q: MemoryResource, + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + tmem_stats_done0: MemoryResource, + tmem_stats_done1: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Capture Keeps MMA with explicit stats lifetime gates.""" + run_mma_schedule( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + tmem_stats_done0, + tmem_stats_done1, + work_queue, + ) + + if tmem_stats_done0 is None or tmem_stats_done1 is None: + schedule_result = ( + mma_schedule( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + ) + if work_queue is None + else mma_schedule( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue, + ) + ) + dst = [tmem_s0, tmem_s1, tmem_o] + else: + schedule_result = ( + mma_keeps_schedule( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + tmem_stats_done0, + tmem_stats_done1, + ) + if work_queue is None + else mma_keeps_schedule( + smem_q, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + tmem_stats_done0, + tmem_stats_done1, + work_queue, + ) + ) + dst = [ + tmem_s0, + tmem_s1, + tmem_o, + tmem_stats_done0, + tmem_stats_done1, + ] + src = [smem_q, smem_k0, smem_k1, smem_v0, smem_v1, smem_p0, smem_p1] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=dst, + cfg=cfg, + warp_idx=cfg.mma_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.mma_num_warps if num_warps is None else num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_task_num_registers, + name="MmaTask", + **kw, + ) + + +def create_mma_task_one_inst_qkv( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + tmem_s: MemoryResource, + smem_p: MemoryResource, + tmem_o: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + tmem_stats_done: MemoryResource, + domain: int | cutlass.Int32, + warp_idx: int | None = None, + num_warps: int | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the one-inst QKV MMA task.""" + + def mma_schedule_body( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + tmem_s: MemoryResource, + smem_p: MemoryResource, + tmem_o: MemoryResource, + tmem_stats_done: MemoryResource, + q_desc: Any, + ) -> None: + """Schedule single-instance QK and PV waves across HEAD/LOOP/TAIL.""" + + def qk_mma(q_desc, qk_mma_label: str, section: FmhaStage) -> None: + """Issue one scheduled single-instance QK wave.""" + _ = section + tmem_stats_done.acquire() + tmem_s.acquire() + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + smem_k.wait() + kv_desc = smem_k.kv_desc() + if cfg.uses_q_desc_ref: + getattr(tmem_s, f"{qk_mma_label}_from_q_ref")( + kv_desc=kv_desc, + head_dim_stage_idx=head_dim_stage_idx, + ) + else: + getattr(tmem_s, qk_mma_label)( + q_desc=q_desc, + kv_desc=kv_desc, + head_dim_stage_idx=head_dim_stage_idx, + ) + smem_k.release() + tmem_s.commit() + tmem_stats_done.commit() + + def pv_mma(vp_mma_label: str, section: FmhaStage) -> None: + """Issue one scheduled single-instance PV wave.""" + _ = section + smem_p.wait() + p_desc_0, p_desc_1, p_tmem_addr_0, p_tmem_addr_1 = smem_p.p_operands() + tmem_o.acquire() + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + smem_v.wait() + v_desc = smem_v.v_desc() + getattr(tmem_o, vp_mma_label)( + v_desc_0=v_desc, + v_desc_1=v_desc, + p_desc_0=p_desc_0, + p_desc_1=p_desc_1, + p_tmem_addr_0=p_tmem_addr_0, + p_tmem_addr_1=p_tmem_addr_1, + inst_idx=KV_INST0, + head_dim_stage_idx=head_dim_stage_idx, + ) + smem_v.release() + tmem_o.commit() + smem_p.release() + + qk_mma(q_desc, "qk_mma_head", FmhaStage.Head) + + with domain_loop(0, domain, 1, unroll=1): + qk_mma(q_desc, "qk_mma_loop", FmhaStage.Loop) + pv_mma("vp_mma_loop", FmhaStage.Loop) + + pv_mma("vp_mma_tail", FmhaStage.Tail) + smem_q.release() + + def mma_schedule_prelude( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + smem_p: MemoryResource, + ) -> None: + """Initialize invariant one-inst descriptor slots.""" + smem_q.init_descriptor_state() + smem_k.init_descriptor_state() + smem_v.init_descriptor_state() + smem_p.init_descriptor_state() + + @schedule + def mma_schedule( + smem_q: MemoryResource, + smem_k: MemoryResource, + smem_v: MemoryResource, + tmem_s: MemoryResource, + smem_p: MemoryResource, + tmem_o: MemoryResource, + tmem_stats_done: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Wrap one-inst MMA work in packed persistent skip handling.""" + _decode_work_tile_schedule_with_invariant_bridge( + cfg, + work_queue, + lambda: mma_schedule_prelude(smem_q, smem_k, smem_v, smem_p), + lambda: smem_q.q_desc(), + lambda: smem_q.wait(), + lambda q_desc: mma_schedule_body( + smem_q, + smem_k, + smem_v, + tmem_s, + smem_p, + tmem_o, + tmem_stats_done, + q_desc, + ), + ) + + schedule_result = ( + mma_schedule( + smem_q, + smem_k, + smem_v, + tmem_s, + smem_p, + tmem_o, + tmem_stats_done, + ) + if work_queue is None + else mma_schedule( + smem_q, + smem_k, + smem_v, + tmem_s, + smem_p, + tmem_o, + tmem_stats_done, + work_queue, + ) + ) + src = [smem_q, smem_k, smem_v, smem_p] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_s, tmem_o, tmem_stats_done], + cfg=cfg, + warp_idx=cfg.mma_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.mma_num_warps if num_warps is None else num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_task_num_registers, + name="MmaTask", + **kw, + ) + + +# ====================================================================== +# MmaTask — warp 12, 1 warp +# K and V share a single SmemKv ring; each MMA loop iter consumes 4 stages +# (K0, V0, K1, V1) of the shared buffer. +# HEAD: wait Q, BMM1(K0), BMM1(K1) +# LOOP[i]: BMM1(nextK0), BMM2(currV0), BMM1(nextK1), BMM2(currV1) +# TAIL: final BMM2(lastV0), final BMM2(lastV1), release Q +# ====================================================================== +def create_mma_task( + smem_q: MemoryResource, + smem_kv: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the MMA task for the shared K/V ring path.""" + + def mma_schedule_body( + smem_q: MemoryResource, + smem_kv: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + q_desc: Any, + ) -> None: + """Schedule shared-KV QK and PV waves across HEAD/LOOP/TAIL.""" + # HEAD: consume Q once and launch the two initial BMM1 waves. + _consume_staged_qk_mma( + smem_kv, + tmem_s0, + smem_p0, + q_desc, + "k_desc_0", + "qk_mma_head", + FmhaStage.Head, + cfg, + ) + _consume_staged_qk_mma( + smem_kv, + tmem_s1, + smem_p1, + q_desc, + "k_desc_1", + "qk_mma_head", + FmhaStage.Head, + cfg, + ) + + # LOOP: consume aliased TMEM P before the next same-instance QK + # overwrites its S columns. SMEM-P profiles retain their established + # K-before-V cadence because P no longer depends on S lifetime. + with domain_loop(0, domain, 1, unroll=1): + if cfg.uses_two_inst_tmem_p: + _consume_staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + "v_desc_0", + "vp_mma_loop", + KV_INST0, + FmhaStage.Loop, + cfg, + ) + _consume_staged_qk_mma( + smem_kv, + tmem_s0, + smem_p0, + q_desc, + "k_desc_0", + "qk_mma_loop", + FmhaStage.Loop, + cfg, + ) + if not cfg.uses_two_inst_tmem_p: + _consume_staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + "v_desc_0", + "vp_mma_loop", + KV_INST0, + FmhaStage.Loop, + cfg, + ) + if cfg.uses_two_inst_tmem_p: + _consume_staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + "v_desc_1", + "vp_mma_loop", + KV_INST1, + FmhaStage.Loop, + cfg, + ) + _consume_staged_qk_mma( + smem_kv, + tmem_s1, + smem_p1, + q_desc, + "k_desc_1", + "qk_mma_loop", + FmhaStage.Loop, + cfg, + ) + if not cfg.uses_two_inst_tmem_p: + _consume_staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + "v_desc_1", + "vp_mma_loop", + KV_INST1, + FmhaStage.Loop, + cfg, + ) + + # TAIL: no future K tiles remain, so only the final two BMM2 waves run. + _consume_staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + "v_desc_0", + "vp_mma_tail", + KV_INST0, + FmhaStage.Tail, + cfg, + ) + _consume_staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + "v_desc_1", + "vp_mma_tail", + KV_INST1, + FmhaStage.Tail, + cfg, + ) + # Q is live for every BMM1 call and can be released only after the loop. + smem_q.release() + + def mma_schedule_prelude( + smem_q: MemoryResource, + smem_kv: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + ) -> None: + """Initialize invariant shared-ring descriptor slots.""" + smem_q.init_descriptor_state() + smem_kv.init_descriptor_state() + smem_p0.init_descriptor_state() + smem_p1.init_descriptor_state() + + @schedule + def mma_schedule( + smem_q: MemoryResource, + smem_kv: MemoryResource, + tmem_s0: MemoryResource, + tmem_s1: MemoryResource, + smem_p0: MemoryResource, + smem_p1: MemoryResource, + tmem_o: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Wrap shared-ring MMA work in packed persistent skip handling.""" + _decode_work_tile_schedule_with_invariant_bridge( + cfg, + work_queue, + lambda: mma_schedule_prelude(smem_q, smem_kv, smem_p0, smem_p1), + lambda: smem_q.q_desc(), + lambda: smem_q.wait(), + lambda q_desc: mma_schedule_body( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + q_desc, + ), + ) + + captured_schedule = ( + mma_schedule( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + ) + if work_queue is None + else mma_schedule( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue, + ) + ) + src = [smem_q, smem_kv, smem_p0, smem_p1] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_s0, tmem_s1, tmem_o], + cfg=cfg, + warp_idx=cfg.mma_warp_idx, + num_warps=cfg.mma_num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_task_num_registers, + name="MmaTask", + **kw, + ) + + +# ====================================================================== +# Softmax0Task — warps 0-3, 4 warps, profile-selected register budget +# LOOP: consume S, produce stats + P + running sum +# LoopLastIter: emit final sum/max for correction tail +# ====================================================================== +def create_softmax0_task( + tmem_s0: MemoryResource, + tmem_softmax_local0: MemoryResource, + smem_p0: MemoryResource, + tmem_softmax_global0: MemoryResource, + tmem_softmax_order: MemoryResource | None, + sparse_softmax_metadata: MemoryResource | None, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the first softmax task, including optional ordered publication.""" + + def softmax0_schedule_body( + tmem_s0: MemoryResource, + tmem_softmax_local0: MemoryResource, + smem_p0: MemoryResource, + tmem_softmax_global0: MemoryResource, + tmem_softmax_order: MemoryResource | None, + sparse_softmax_metadata: MemoryResource | None, + ) -> None: + """Build the softmax0 loop, P publication, and final stats handoff.""" + ( + old_max_arr, + sum_arr, + new_max_arr, + s_arr, + ) = tmem_s0.init_softmax_state() + smem_p0.init_compute_state() + if sparse_softmax_metadata is not None: + sparse_softmax_metadata.init_read_state() + + with domain_loop(0, domain, 1, unroll=1) as d: + # ConsWait/ConsWork: load S from TMEM and compute the tile max. + tmem_s0.wait() + if sparse_softmax_metadata is not None: + sparse_softmax_metadata.wait() + # Copy the complete payload to registers before release, so + # masking cannot race the producer's next SMEM-stage reuse. + ( + sparse_origin0, + sparse_origin1, + sparse_route_flags, + sparse_token_word0, + sparse_token_word1, + sparse_token_word2, + sparse_token_word3, + ) = sparse_softmax_metadata.load_route() + sparse_softmax_metadata.release() + old_max_arr, sum_arr, new_max_arr, s_arr = ( + tmem_s0.compute_block_sparse_softmax_loop( + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + sparse_origin0=sparse_origin0, + sparse_origin1=sparse_origin1, + sparse_route_flags=sparse_route_flags, + sparse_token_word0=sparse_token_word0, + sparse_token_word1=sparse_token_word1, + sparse_token_word2=sparse_token_word2, + sparse_token_word3=sparse_token_word3, + ) + ) + else: + old_max_arr, sum_arr, new_max_arr, s_arr = tmem_s0.compute_softmax_loop( + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + ) + if cutlass.const_expr(not cfg.use_keeps_mma_ab or not cfg.uses_tmem_p): + # ConsRelease: free S once the scores are in registers unless + # a Keeps TMEM-P operand still aliases the consumed columns. + tmem_s0.release() + # Publish old/new max before the P path so correction can observe + # the same stats order as the decode pipeline. + # ProdWork: store old/new max for correction's in-loop O update. + tmem_softmax_local0.acquire() + tmem_softmax_local0.store_loop_old_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local0.commit() + if cutlass.const_expr(cfg.streams_tmem_p_fragments): + # Publish one K32 probability fragment at a time so PV can + # consume early fragments while later scores are processed. + for fragment_idx in range(cfg.num_softmax_score_fragments): + s_arr = tmem_s0.load_softmax_p_fragment( + fragment_idx=fragment_idx, + s_arr=s_arr, + ) + smem_p0.compute_p_fragment( + fragment_idx=fragment_idx, + new_max_arr=new_max_arr, + s_arr=s_arr, + ) + else: + # Wait for a free P stage before entering the ordered window so + # BMM2 backpressure on this group's P pipeline cannot extend the + # baton hold and stall the partner softmax group. + smem_p0.acquire() + if tmem_softmax_order is not None: + tmem_softmax_order.wait_softmax0() + # ProdWork: compute P=exp(S-new_max), store it in the profile's + # SMEM or staged-TMEM operand, and record local sums for the + # running softmax sum update. + smem_p0.compute_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + ) # publishes the local denominator through tmem_s0 + smem_p0.commit() + if tmem_softmax_order is not None: + tmem_softmax_order.release_softmax1() + if cutlass.const_expr(cfg.use_keeps_mma_ab and cfg.uses_tmem_p): + # The TMEM-P store has consumed the aliased S columns, so the + # next QK wave can now overwrite them. + tmem_s0.release() + # ProdWork: FP8 path applies the cross-resource sum correction + # before TmemS.reduce_sums publishes the new running sums. + tmem_softmax_global0.global_correction( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) # publishes the corrected denominator through tmem_s0 + # ConsTailWork: update the running sum after P is available. + sum_arr = tmem_s0.reduce_sums( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + if cutlass.const_expr(not cfg.use_keeps_mma_ab): + with d.last_iter(): + # LastIter ProdWork: store the final sums for correction's + # normalization/output tail. + tmem_softmax_local0.acquire() + tmem_softmax_local0.store_loop_sum_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local0.commit() + if cutlass.const_expr(cfg.use_keeps_mma_ab): + # The final sum payload has no matching QK/StatsDone token. + tmem_softmax_local0.acquire() + tmem_softmax_local0.store_tail_sum_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local0.commit() + if cutlass.const_expr( + cfg.use_persistent_scheduler and cfg.uses_staged_one_inst_tmem_p + ): + # Tail stats add one stage beyond the S cadence. Advance the + # second stats slot so both pipelines start the next work tile + # on the same physical TMEM stage. + tmem_softmax_local0.acquire() + tmem_softmax_local0.commit() + + @_schedule_with_optional_resources + def softmax0_schedule( + tmem_s0: MemoryResource, + tmem_softmax_local0: MemoryResource, + smem_p0: MemoryResource, + tmem_softmax_global0: MemoryResource, + tmem_softmax_order: MemoryResource | None, + sparse_softmax_metadata: MemoryResource | None, + work_queue: WorkQueue | None, + ) -> None: + """Schedule softmax0 with the supplied order and sparse resources.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: softmax0_schedule_body( + tmem_s0, + tmem_softmax_local0, + smem_p0, + tmem_softmax_global0, + tmem_softmax_order, + sparse_softmax_metadata, + ), + ) + + schedule_result = softmax0_schedule( + tmem_s0, + tmem_softmax_local0, + smem_p0, + tmem_softmax_global0, + tmem_softmax_order, + sparse_softmax_metadata, + work_queue, + ) + src = [tmem_s0] + if sparse_softmax_metadata is not None: + src.append(sparse_softmax_metadata) + if work_queue is not None: + src.append(work_queue) + dst = [tmem_softmax_local0, smem_p0, tmem_softmax_global0] + if tmem_softmax_order is not None: + dst.append(tmem_softmax_order) + return task_class( + src_resources=src, + dst_resources=dst, + q_bound_resources=((tmem_s0, False),), + cfg=cfg, + warp_idx=cfg.softmax0_warp_idx, + num_warps=cfg.softmax0_num_warps, + schedule=schedule_result, + num_registers=cfg.softmax_task_num_registers, + name="Softmax0Task", + **kw, + ) + + +# ====================================================================== +# Softmax1Task — warps 4-7, 4 warps, profile-selected register budget +# ====================================================================== +def create_softmax1_task( + tmem_s1: MemoryResource, + tmem_softmax_local1: MemoryResource, + smem_p1: MemoryResource, + tmem_softmax_global1: MemoryResource, + tmem_softmax_order: MemoryResource | None, + sparse_softmax_metadata: MemoryResource | None, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the second softmax task, including optional ordered publication.""" + + def softmax1_schedule_body( + tmem_s1: MemoryResource, + tmem_softmax_local1: MemoryResource, + smem_p1: MemoryResource, + tmem_softmax_global1: MemoryResource, + tmem_softmax_order: MemoryResource | None, + sparse_softmax_metadata: MemoryResource | None, + ) -> None: + """Build the softmax1 loop, P publication, and final stats handoff.""" + ( + old_max_arr, + sum_arr, + new_max_arr, + s_arr, + ) = tmem_s1.init_softmax_state() + smem_p1.init_compute_state() + if sparse_softmax_metadata is not None: + sparse_softmax_metadata.init_read_state() + + with domain_loop(0, domain, 1, unroll=1) as d: + # ConsWait/ConsWork: load the second S instance and compute max. + tmem_s1.wait() + if sparse_softmax_metadata is not None: + sparse_softmax_metadata.wait() + # Copy to registers before release so the producer can reuse + # the SMEM stage while this warp group applies the masks. + ( + sparse_origin0, + sparse_origin1, + sparse_route_flags, + sparse_token_word0, + sparse_token_word1, + sparse_token_word2, + sparse_token_word3, + ) = sparse_softmax_metadata.load_route() + sparse_softmax_metadata.release() + old_max_arr, sum_arr, new_max_arr, s_arr = ( + tmem_s1.compute_block_sparse_softmax_loop( + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + sparse_origin0=sparse_origin0, + sparse_origin1=sparse_origin1, + sparse_route_flags=sparse_route_flags, + sparse_token_word0=sparse_token_word0, + sparse_token_word1=sparse_token_word1, + sparse_token_word2=sparse_token_word2, + sparse_token_word3=sparse_token_word3, + ) + ) + else: + old_max_arr, sum_arr, new_max_arr, s_arr = tmem_s1.compute_softmax_loop( + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + ) + if cutlass.const_expr(not cfg.use_keeps_mma_ab or not cfg.uses_tmem_p): + # ConsRelease: SMEM-P Keeps and Swaps no longer need S after + # the score fragment has been loaded into registers. + tmem_s1.release() + # ProdWork: store old/new max for the correction task. + tmem_softmax_local1.acquire() + tmem_softmax_local1.store_loop_old_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local1.commit() + if cutlass.const_expr(cfg.streams_tmem_p_fragments): + for fragment_idx in range(cfg.num_softmax_score_fragments): + s_arr = tmem_s1.load_softmax_p_fragment( + fragment_idx=fragment_idx, + s_arr=s_arr, + ) + smem_p1.compute_p_fragment( + fragment_idx=fragment_idx, + new_max_arr=new_max_arr, + s_arr=s_arr, + ) + else: + # Wait for a free P stage before entering the ordered window so + # BMM2 backpressure on this group's P pipeline cannot extend the + # baton hold and stall the partner softmax group. + smem_p1.acquire() + if tmem_softmax_order is not None: + tmem_softmax_order.wait_softmax1() + # ProdWork: compute and publish P1 for BMM2. + smem_p1.compute_p(new_max_arr=new_max_arr, s_arr=s_arr) + smem_p1.commit() + if tmem_softmax_order is not None: + tmem_softmax_order.release_softmax0() + if cutlass.const_expr(cfg.use_keeps_mma_ab and cfg.uses_tmem_p): + # The TMEM-P store has consumed the aliased S columns, so the + # next QK wave can now overwrite them. + tmem_s1.release() + # ProdWork: update FP8 global sums if the configuration needs the + # split softmax-sum correction path. + tmem_softmax_global1.global_correction( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + # ConsTailWork: fold local sums into the running sums. + sum_arr = tmem_s1.reduce_sums( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + if cutlass.const_expr(not cfg.use_keeps_mma_ab): + with d.last_iter(): + # LastIter ProdWork: publish final sums for output + # normalization. + tmem_softmax_local1.acquire() + tmem_softmax_local1.store_loop_sum_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local1.commit() + if cutlass.const_expr(cfg.use_keeps_mma_ab): + # The final sum payload has no matching QK/StatsDone token. + tmem_softmax_local1.acquire() + tmem_softmax_local1.store_tail_sum_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local1.commit() + + @_schedule_with_optional_resources + def softmax1_schedule( + tmem_s1: MemoryResource, + tmem_softmax_local1: MemoryResource, + smem_p1: MemoryResource, + tmem_softmax_global1: MemoryResource, + tmem_softmax_order: MemoryResource | None, + sparse_softmax_metadata: MemoryResource | None, + work_queue: WorkQueue | None, + ) -> None: + """Schedule softmax1 with the supplied order and sparse resources.""" + # Prime the first P0 -> P1 baton once per CTA. Each completed P1 + # publication leaves the same barrier half-arrived for the next + # persistent work tile, so re-priming inside the work-tile loop would + # toggle its phase early and deadlock on the second tile. + if tmem_softmax_order is not None: + tmem_softmax_order.prime_softmax1() + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: softmax1_schedule_body( + tmem_s1, + tmem_softmax_local1, + smem_p1, + tmem_softmax_global1, + tmem_softmax_order, + sparse_softmax_metadata, + ), + ) + + schedule_result = softmax1_schedule( + tmem_s1, + tmem_softmax_local1, + smem_p1, + tmem_softmax_global1, + tmem_softmax_order, + sparse_softmax_metadata, + work_queue, + ) + src = [tmem_s1] + if sparse_softmax_metadata is not None: + src.append(sparse_softmax_metadata) + if tmem_softmax_order is not None: + src.append(tmem_softmax_order) + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_softmax_local1, smem_p1, tmem_softmax_global1], + q_bound_resources=((tmem_s1, False),), + cfg=cfg, + warp_idx=cfg.softmax1_warp_idx, + num_warps=cfg.softmax1_num_warps, + schedule=schedule_result, + num_registers=cfg.softmax_task_num_registers, + name="Softmax1Task", + **kw, + ) + + +# ====================================================================== +# CorrectionTask — warps 8-11, 4 warps, profile-selected register budget +# HEAD: drain initial softmax-local stats +# LOOP: correct O0 / O1 in-place before the next BMM2 accumulation +# TAIL: combine final O0/O1, normalize, store to GMEM +# ====================================================================== +def create_correction_task( + tmem_softmax_local0: MemoryResource, + tmem_softmax_local1: MemoryResource, + tmem_o: MemoryResource, + tmem_corr0: MemoryResource, + tmem_corr1: MemoryResource, + work_queue: WorkQueue | None, + smem_kv_reuse_credit: MemoryResource | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + tmem_stats_done0: MemoryResource | None = None, + tmem_stats_done1: MemoryResource | None = None, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the two-instance correction and output task.""" + + if smem_kv_reuse_credit is not None and work_queue is None: + raise ValueError("KV reuse credit requires a work queue") + + def correction_schedule_body( + tmem_softmax_local0: MemoryResource, + tmem_softmax_local1: MemoryResource, + tmem_o: MemoryResource, + tmem_corr0: MemoryResource, + tmem_corr1: MemoryResource, + tmem_stats_done0: MemoryResource | None, + tmem_stats_done1: MemoryResource | None, + smem_kv_reuse_credit: MemoryResource | None, + ) -> None: + """Schedule two-instance O correction and final output normalization.""" + + def consume_local_with_load( + tmem_softmax_local: MemoryResource, + tmem_stats_done: MemoryResource | None, + local_state: tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ], + load_stats_work, + ) -> tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ]: + """Wait, load, and release one phase-specific stats payload.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ) = local_state + # ConsWait/ConsWork: load the phase-specific softmax stats payload + # from TMEM-local storage. + tmem_softmax_local.wait() + local_state = load_stats_work( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst_old_max_arr=inst_old_max_arr, + inst_new_max_arr=inst_new_max_arr, + inst_sum_arr=inst_sum_arr, + ) + if tmem_stats_done is not None: + # Return the matching S-overwrite credit only after the stats + # payload is resident in Correction registers. + tmem_stats_done.wait() + tmem_stats_done.release() + # ConsRelease: the softmax-local payload can now be reused. + tmem_softmax_local.release() + return local_state + + def correct_o( + tmem_softmax_local: MemoryResource, + tmem_stats_done: MemoryResource | None, + tmem_corr: MemoryResource, + local_state: tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ], + tail_0: cutlass.Int32, + tail_1: cutlass.Int32, + ) -> tuple[ + tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ], + cutlass.Int32, + cutlass.Int32, + ]: + """Consume loop stats and rescale the matching O stage.""" + # Load old/new max for the O stage that is about to be corrected. + local_state = consume_local_with_load( + tmem_softmax_local, + tmem_stats_done, + local_state, + tmem_softmax_local.load_loop_stats, + ) + old_max_arr = local_state[0] + new_max_arr = local_state[1] + # ConsWait/ConsWork: wait for the matching TMEM O stage and record + # which O buffer is being consumed. + tmem_o.wait() + o_stage_idx, tail_0, tail_1 = tmem_o.update_o_stage_loop( + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + ) + # ProdWork: rescale O in place before the next BMM2 accumulation. + tmem_corr.correction_loop_epilogue( + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + ) + tmem_o.release() + return local_state, tail_0, tail_1 + + local0_state = tmem_softmax_local0.init_stats_state() + local1_state = tmem_softmax_local1.init_stats_state() + _, tail_0, tail_1 = tmem_o.init_stage_state() + tmem_corr0.init_epilogue_state() + tmem_corr1.init_epilogue_state() + + # HEAD: drain the first softmax-local handoff. No O is corrected here; + # this aligns the stats pipeline before loop work starts. + local0_state = consume_local_with_load( + tmem_softmax_local0, + tmem_stats_done0, + local0_state, + tmem_softmax_local0.load_head_stats, + ) + local1_state = consume_local_with_load( + tmem_softmax_local1, + tmem_stats_done1, + local1_state, + tmem_softmax_local1.load_head_stats, + ) + + # LOOP: each iteration corrects O0 and O1 before later BMM2 waves + # accumulate into the same TMEM columns. + with domain_loop(0, domain, 1, unroll=1): + local0_state, tail_0, tail_1 = correct_o( + tmem_softmax_local0, + tmem_stats_done0, + tmem_corr0, + local0_state, + tail_0, + tail_1, + ) + local1_state, tail_0, tail_1 = correct_o( + tmem_softmax_local1, + tmem_stats_done1, + tmem_corr1, + local1_state, + tail_0, + tail_1, + ) + + # TAIL inst0: consume the final stats and mark the first tail O stage. + local0_state = consume_local_with_load( + tmem_softmax_local0, + None, + local0_state, + tmem_softmax_local0.load_tail_stats, + ) + old_max_arr = local0_state[0] + new_max_arr = local0_state[1] + inst0_new_max_arr = local0_state[4] + inst0_sum_arr = local0_state[5] + tmem_o.wait() + o_stage_idx, tail_0, tail_1 = tmem_o.update_o_stage_tail( + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + inst_idx=KV_INST0, + ) + tmem_corr0.correction_tail_epilogue( + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst0_new_max_arr, + inst1_sum_arr=inst0_sum_arr, + ) + # TAIL inst1: consume final stats for the second instance. This call + # performs the final two-instance normalization and output store. + local1_state = consume_local_with_load( + tmem_softmax_local1, + None, + local1_state, + tmem_softmax_local1.load_tail_stats, + ) + old_max_arr = local1_state[0] + new_max_arr = local1_state[1] + inst1_new_max_arr = local1_state[4] + inst1_sum_arr = local1_state[5] + tmem_o.wait() + o_stage_idx, tail_0, tail_1 = tmem_o.update_o_stage_tail( + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + inst_idx=KV_INST1, + ) + if smem_kv_reuse_credit is None: + tmem_corr1.correction_tail_epilogue( + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + else: + # The stage selector and ownership token share one pipeline epoch. + # Wait before the first aliased access and release immediately + # after correction stops touching the selected KV-ring stage. + smem_kv_reuse_credit.wait() + scratch_stage = smem_kv_reuse_credit.read_scratch_stage() + tmem_corr1.correction_tail_epilogue_rotating_exchange( + scratch_stage=scratch_stage, + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + smem_kv_reuse_credit.release() + # Inst1 final reduction consumes both O0 and O1, so defer O0 release + # until after inst1 has finished reading it. + tmem_o.release() + tmem_o.release() + + def run_correction_schedule( + tmem_softmax_local0: MemoryResource, + tmem_softmax_local1: MemoryResource, + tmem_o: MemoryResource, + tmem_corr0: MemoryResource, + tmem_corr1: MemoryResource, + tmem_stats_done0: MemoryResource | None, + tmem_stats_done1: MemoryResource | None, + smem_kv_reuse_credit: MemoryResource | None, + work_queue: WorkQueue | None, + ) -> None: + """Wrap correction with optional stats lifetime gates.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: correction_schedule_body( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + tmem_stats_done0, + tmem_stats_done1, + smem_kv_reuse_credit, + ), + ) + + @schedule + def correction_schedule( + tmem_softmax_local0: MemoryResource, + tmem_softmax_local1: MemoryResource, + tmem_o: MemoryResource, + tmem_corr0: MemoryResource, + tmem_corr1: MemoryResource, + work_queue: WorkQueue | None = None, + smem_kv_reuse_credit: MemoryResource | None = None, + ) -> None: + """Capture the Swaps correction schedule.""" + run_correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + None, + None, + smem_kv_reuse_credit, + work_queue, + ) + + @schedule + def correction_keeps_schedule( + tmem_softmax_local0: MemoryResource, + tmem_softmax_local1: MemoryResource, + tmem_o: MemoryResource, + tmem_corr0: MemoryResource, + tmem_corr1: MemoryResource, + tmem_stats_done0: MemoryResource, + tmem_stats_done1: MemoryResource, + work_queue: WorkQueue | None = None, + smem_kv_reuse_credit: MemoryResource | None = None, + ) -> None: + """Capture Keeps correction with explicit stats lifetime gates.""" + run_correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + tmem_stats_done0, + tmem_stats_done1, + smem_kv_reuse_credit, + work_queue, + ) + + if tmem_stats_done0 is None or tmem_stats_done1 is None: + if work_queue is None: + captured_schedule = correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + ) + elif smem_kv_reuse_credit is None: + captured_schedule = correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + work_queue, + ) + else: + captured_schedule = correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + work_queue, + smem_kv_reuse_credit, + ) + src = [tmem_softmax_local0, tmem_softmax_local1, tmem_o] + else: + if work_queue is None: + captured_schedule = correction_keeps_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + tmem_stats_done0, + tmem_stats_done1, + ) + elif smem_kv_reuse_credit is None: + captured_schedule = correction_keeps_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + tmem_stats_done0, + tmem_stats_done1, + work_queue, + ) + else: + captured_schedule = correction_keeps_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + tmem_stats_done0, + tmem_stats_done1, + work_queue, + smem_kv_reuse_credit, + ) + src = [ + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_stats_done0, + tmem_stats_done1, + ] + if work_queue is not None: + src.append(work_queue) + if smem_kv_reuse_credit is not None: + src.append(smem_kv_reuse_credit) + return task_class( + src_resources=src, + dst_resources=[tmem_corr0, tmem_corr1], + q_bound_resources=((tmem_corr0, True), (tmem_corr1, True)), + cfg=cfg, + warp_idx=cfg.correction_warp_idx, + num_warps=cfg.correction_num_warps, + schedule=captured_schedule, + num_registers=cfg.correction_task_num_registers, + name="CorrectionTask", + **kw, + ) + + +def create_correction_task_one_inst_qkv( + tmem_softmax_local: MemoryResource, + tmem_o: MemoryResource, + tmem_corr: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + tmem_stats_done: MemoryResource, + domain: int | cutlass.Int32, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the correction task for the one-inst QKV path.""" + + def correction_schedule_body( + tmem_softmax_local: MemoryResource, + tmem_o: MemoryResource, + tmem_corr: MemoryResource, + tmem_stats_done: MemoryResource, + ) -> None: + """Schedule one-inst O correction and final output normalization.""" + + def consume_local_with_load( + local_state: tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ], + load_stats_work, + release_stats_done: bool, + ) -> tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ]: + """Wait, load, and release one one-inst stats payload.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst_old_max_arr, + inst_new_max_arr, + inst_sum_arr, + ) = local_state + # ConsWait/ConsWork: load the stage-specific softmax-local payload + # produced by the one-inst softmax task. + tmem_softmax_local.wait() + local_state = load_stats_work( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst_old_max_arr=inst_old_max_arr, + inst_new_max_arr=inst_new_max_arr, + inst_sum_arr=inst_sum_arr, + ) + if release_stats_done: + # Return one S-overwrite credit for this QK-derived payload. + tmem_stats_done.wait() + tmem_stats_done.release() + # ConsRelease: the softmax-local payload can now be reused. + tmem_softmax_local.release() + return local_state + + def correct_o( + local_state: tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ], + tail_0: cutlass.Int32, + tail_1: cutlass.Int32, + ) -> tuple[ + tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + ], + cutlass.Int32, + cutlass.Int32, + ]: + """Consume loop stats and rescale the one-inst O stage.""" + local_state = consume_local_with_load( + local_state, + tmem_softmax_local.load_loop_stats, + True, + ) + old_max_arr = local_state[0] + new_max_arr = local_state[1] + # ConsWait/ConsWork: consume the O stage produced by BMM2. The + # stage state tracks which TMEM O slot is live and which two slots + # become the final tail pair. + tmem_o.wait() + o_stage_idx, tail_0, tail_1 = tmem_o.update_o_stage_loop( + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + ) + # ProdWork: loop phase rescales O in place before the next BMM2 + # accumulation. + tmem_corr.correction_loop_epilogue( + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + ) + # ConsRelease: the consumed O stage can be reused by later BMM2. + tmem_o.release() + return local_state, tail_0, tail_1 + + local_state = tmem_softmax_local.init_stats_state() + _, tail_0, tail_1 = tmem_o.init_stage_state() + tmem_corr.init_epilogue_state() + + # HEAD: drain the first stats handoff. The first payload only seeds + # the correction state; there is no prior O stage to rescale. + local_state = consume_local_with_load( + local_state, + tmem_softmax_local.load_head_stats, + True, + ) + # LOOP: each payload/O pair corresponds to a completed BMM2 wave whose + # accumulator must be max-corrected before more V work accumulates. + with domain_loop(0, domain, 1, unroll=1): + local_state, tail_0, tail_1 = correct_o(local_state, tail_0, tail_1) + # TAIL: the final payload carries sum/new max and the final O stage is + # normalized/stored instead of being staged for another accumulation. + local_state = consume_local_with_load( + local_state, + tmem_softmax_local.load_tail_stats, + False, + ) + if cutlass.const_expr( + cfg.use_persistent_scheduler and cfg.uses_staged_one_inst_tmem_p + ): + # Match the producer's empty tail stage before the persistent + # work queue advances to the next tile. + tmem_softmax_local.wait() + tmem_softmax_local.release() + old_max_arr = local_state[0] + new_max_arr = local_state[1] + inst_new_max_arr = local_state[4] + inst_sum_arr = local_state[5] + # ConsWait/ConsWork: consume the final O stage produced by BMM2. + tmem_o.wait() + o_stage_idx, tail_0, tail_1 = tmem_o.update_o_stage_tail( + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + inst_idx=KV_INST0, + ) + # ProdWork: tail phase normalizes and stores final O. + tmem_corr.correction_tail_epilogue( + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst_new_max_arr, + inst0_sum_arr=inst_sum_arr, + inst1_new_max_arr=inst_new_max_arr, + inst1_sum_arr=inst_sum_arr, + ) + # ConsRelease: the final O stage is no longer needed. + tmem_o.release() + + @schedule + def correction_schedule( + tmem_softmax_local: MemoryResource, + tmem_o: MemoryResource, + tmem_corr: MemoryResource, + tmem_stats_done: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + """Wrap one-inst correction in packed persistent skip handling.""" + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: correction_schedule_body( + tmem_softmax_local, + tmem_o, + tmem_corr, + tmem_stats_done, + ), + ) + + schedule_result = ( + correction_schedule( + tmem_softmax_local, + tmem_o, + tmem_corr, + tmem_stats_done, + ) + if work_queue is None + else correction_schedule( + tmem_softmax_local, + tmem_o, + tmem_corr, + tmem_stats_done, + work_queue, + ) + ) + src = [tmem_softmax_local, tmem_o, tmem_stats_done] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_corr], + q_bound_resources=((tmem_corr, True),), + cfg=cfg, + warp_idx=cfg.correction_warp_idx, + num_warps=cfg.correction_num_warps, + schedule=schedule_result, + num_registers=cfg.correction_task_num_registers, + name="CorrectionTask", + **kw, + ) + + +# ====================================================================== +# PaddingTask — fills the unused tail warps of one warp group +# ====================================================================== +def create_padding_task( + cfg: FmhaDecodeConfig, + work_queue: WorkQueue | None = None, + *, + warp_idx: int, + num_warps: int, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the padding task used to balance persistent warp groups.""" + + @schedule + def padding_schedule(work_queue: WorkQueue | None = None) -> None: + """Run the padding schedule and advance the work queue if present.""" + + def padding_body() -> None: + # Empty loop body: this warp participates in task/register + # scheduling so the persistent warpgroup layout remains balanced. + with domain_loop(0, 1, 1, unroll=1): + pass + + _decode_work_tile_schedule(cfg, work_queue, padding_body) + + captured_schedule = ( + padding_schedule() if work_queue is None else padding_schedule(work_queue) + ) + src = [] + if work_queue is not None: + src = [work_queue] + return task_class( + src_resources=src, + dst_resources=[], + cfg=cfg, + warp_idx=warp_idx, + num_warps=num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_task_num_registers, + name="PaddingTask", + **kw, + ) + + +def create_scheduler_task( + work_queue: WorkQueue, + schedule_token_throttle: MemoryResource | None, + cfg: FmhaDecodeConfig, + *, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the CLC scheduler task. + + Two-inst profiles place scheduler/load/padding roles in WG3. One-inst Keeps + profiles place scheduler at warp 9 beside MMA, page-or-padding, and load in + WG2; a separate padding task keeps WG3 and the 512-thread CLC contract full. + """ + + @schedule + def scheduler_schedule( + work_queue: WorkQueue, + schedule_token_throttle: MemoryResource | None = None, + ) -> None: + """Fetch and publish the next persistent-scheduler work tile.""" + with work_tile_loop(work_queue): + # The scheduler owns work-tile discovery for persistent kernels. + # Empty domain keeps the generated schedule shape consistent with + # other TS tasks while all real work happens through WorkQueue. + with domain_loop(0, 0, 1, unroll=1): + pass + _schedule_token_throttle_tail(schedule_token_throttle) + # ProdAcquire/ProdWork/ProdCommit: fetch and publish the next work + # tile to all persistent tasks. + work_queue.acquire() + work_queue.fetch_work_tile() + work_queue.commit() + _work_queue_tail(work_queue) + + captured_schedule = ( + scheduler_schedule(work_queue) + if schedule_token_throttle is None + else scheduler_schedule(work_queue, schedule_token_throttle) + ) + src = [work_queue] + if schedule_token_throttle is not None: + src.append(schedule_token_throttle) + return task_class( + src_resources=src, + dst_resources=[work_queue], + cfg=cfg, + warp_idx=cfg.scheduler_warp_idx, + num_warps=cfg.scheduler_num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_task_num_registers, + name="SchedulerTask", + **kw, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py new file mode 100644 index 000000000000..673d3205aa39 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py @@ -0,0 +1,902 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Standalone GMEM reducers for FMHA decode TS split-KV profiles. + +The decode kernel publishes one normalized 16-bit O vector and one FP32 +log2-LSE scalar per ``(batch, kv_head, split_kv, output_row)``. Reducer threads +own 16-byte output fragments and combine those normalized states with the +shared log2-LSE recurrence. + +The production schedule is selected from the split count: + +* S2-S4 use one 512-thread CTA per 8 KiB output slice and fold the exact split + count in registers. +* S5 and larger use 128-thread CTAs over 2 KiB slices. Each cluster rank folds + its split slots locally, publishes ``(LSE, O)`` through SMEM, and rank zero + merges the distributed states. G16 uses a 4x4 two-level merge. + +A 512-thread exact-split kernel remains as the serial reference schedule. Both +schedules merge the optional attention-sink denominator, pack the requested +output dtype, and write final O. Batch and KV head remain grid dimensions; +only split-KV is reduced here. +""" + +import math + +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda_drv +from cutlass import Float32, Int32, Int64 +from cutlass.experimental import primitives as prims + +from .fmha_decode_constants import ( + FP32_BYTES, + FP8_PACKED_OUTPUT_REGS_PER_THREAD, + FP8_VALUES_PER_REG, + FP16_VALUES_PER_REG, + OUTPUT_VALUES_PER_THREAD, + PACKED_OUTPUT_REGS_PER_THREAD, + PACKED_REGISTER_BYTES, + PARALLEL_REDUCTION_BYTES_PER_SLICE, + PARALLEL_REDUCTION_FINAL_REDUCERS, + PARALLEL_REDUCTION_LOAD_BATCH, + PARALLEL_REDUCTION_THREADS_PER_CTA, + PARTIAL_O_ELEMENT_BYTES, + REDUCTION_BYTES_PER_SLICE, + REDUCTION_BYTES_PER_THREAD, + REDUCTION_THREADS_PER_CTA, + SEPARATE_REDUCTION_LSE_VALUES_PER_ROW, +) +from .fmha_decode_config import FmhaDecodeConfig +from .fmha_decode_resources.helpers_common import ( + _attention_sink_head_stride, + _local_head_from_q_output_row, + _pack_float2_to_bf16, + _pack_float2_to_fp16, + _q_group_token_base, + _q_logical_output_row_token_and_local_head, + _q_logical_output_row_is_valid_for_seq, + _q_physical_output_row_from_logical, + _q_seq_bounds, + fmul2, +) +from .fmha_decode_resources.helpers_kv_tile_idx import _runtime_active_splits_kv +from .fmha_decode_resources.helpers_softmax import _pack_float4_to_fp8_e4m3 +from ..separate_reduction import ( + merge_log2_lse, + unpack_normalized_vec8, +) + + +@cute.jit +def _separate_workspace_row_offset( + logical_kv_idx: Int64, + split_idx: Int32, + row_idx: Int32, + rows_per_split: Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], +) -> Int64: + """Return a split/row offset without overflowing 32-bit products.""" + + return (logical_kv_idx * Int64(cfg.max_splits_kv) + Int64(split_idx)) * Int64( + rows_per_split + ) + Int64(row_idx) + + +@cute.jit +def _reduction_q_group_idx( + cfg: cutlass.Constexpr[FmhaDecodeConfig], + h_r: Int32, + logical_output_row_idx: Int32, +) -> Int32: + """Recover the producer Q-group owning one logical scratch row.""" + token_idx, local_head_idx = _q_logical_output_row_token_and_local_head( + cfg, h_r, logical_output_row_idx + ) + if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): + return token_idx // Int32(cfg.q_tokens_per_cta) + if cutlass.const_expr(cfg.max_seq_len_q > 1): + head_ctas_per_token = Int32( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + return token_idx * head_ctas_per_token + local_head_idx // Int32( + cfg.tile_size_q + ) + return local_head_idx // Int32(cfg.tile_size_q) + + +@cute.jit +def _reduction_active_splits_kv( + cfg: cutlass.Constexpr[FmhaDecodeConfig], + g_seqlens_kv: cute.Pointer, + b_idx: Int32, + h_r: Int32, + logical_output_row_idx: Int32, + seq_len_q: Int32, +) -> Int32: + """Recompute the producer's useful split prefix for one output row.""" + q_group_idx = _reduction_q_group_idx(cfg, h_r, logical_output_row_idx) + return _runtime_active_splits_kv( + cfg, + Int32(g_seqlens_kv[b_idx]), + seq_len_q, + _q_group_token_base(cfg, q_group_idx), + ) + + +@cute.jit +def _reduce_exact_splits_body( + o_iter: cute.Pointer, + g_seqlens_kv: cute.Pointer, + g_cu_seqlens_q: cute.Pointer, + g_partial_o: cute.Pointer, + g_partial_stats: cute.Pointer, + g_attention_sinks: cute.Pointer, + g_q_output_rows: Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + static_full_split_prefix: cutlass.Constexpr[bool] = False, +) -> None: + """Fold every split in one 512-thread CTA over one 8 KiB output slice. + + ``g_partial_stats`` stores one log2-LSE scalar for each split and output + row. ``g_partial_o`` stores the corresponding normalized 16-bit O fragment. + This body is shared by the serial reference kernel and the compact S2-S4 + production schedule; PDL ordering remains in the production outer kernel. + """ + + thread_idx, _, _ = cute.arch.thread_idx() + slice_idx, h_k_idx, b_idx = cute.arch.block_idx() + _, grid_h_k, _ = cute.arch.grid_dim() + # Flatten batch and KV-head so the partial buffers use one contiguous + # logical tile index independent of the reducer launch grid layout. + logical_kv_idx = Int64(b_idx) * Int64(grid_h_k) + Int64(h_k_idx) + attention_sink_h_r = _attention_sink_head_stride(cfg, g_q_output_rows) + + # One CTA owns one contiguous byte slice of the final output tile. Within + # the CTA, each thread maps to one 16-byte fragment and derives both the + # output row and the column offset inside that row from the byte offset. + bytes_per_output_row = Int32(cfg.headdim * PARTIAL_O_ELEMENT_BYTES) + bytes_per_thread = Int32(REDUCTION_BYTES_PER_THREAD) + bytes_per_slice = Int32(REDUCTION_BYTES_PER_SLICE) + reduce_base_offset = slice_idx * bytes_per_slice + thread_idx * bytes_per_thread + reduce_row_idx = reduce_base_offset // bytes_per_output_row + reduce_col_idx = (reduce_base_offset % bytes_per_output_row) // Int32( + PARTIAL_O_ELEMENT_BYTES + ) + q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) + active_splits_kv = Int32(cfg.splits_kv) + if cutlass.const_expr(not static_full_split_prefix): + active_splits_kv = _reduction_active_splits_kv( + cfg, + g_seqlens_kv, + b_idx, + g_q_output_rows, + reduce_row_idx, + seq_len_q, + ) + valid_reduce_row = _q_logical_output_row_is_valid_for_seq( + cfg, + g_q_output_rows, + reduce_row_idx, + seq_len_q, + ) + output_row_idx = _q_physical_output_row_from_logical( + cfg, + g_q_output_rows, + grid_h_k, + b_idx, + h_k_idx, + reduce_row_idx, + q_token_offset, + ) + attention_sink_head_idx = _local_head_from_q_output_row( + cfg, + g_q_output_rows, + reduce_row_idx, + ) + + output_vals = cutlass.Array( + Float32, OUTPUT_VALUES_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + output_vals[elem_idx] = Float32(0.0) + + global_lse = Float32(-Float32.inf) + + if valid_reduce_row: + # The workspace retains configured-max strides, but only the runtime + # active prefix was published by producer CTAs. + for split_idx_i in cutlass.range_constexpr(cfg.max_splits_kv): + split_idx = Int32(split_idx_i) + split_is_active = split_idx < active_splits_kv + if cutlass.const_expr(static_full_split_prefix): + split_is_active = cutlass.const_expr(split_idx_i < cfg.splits_kv) + if split_is_active: + # LSE layout: [logical_kv][configured split][row]. + workspace_row = _separate_workspace_row_offset( + logical_kv_idx, + split_idx, + reduce_row_idx, + g_q_output_rows, + cfg, + ) + stats_offset = workspace_row * Int64( + SEPARATE_REDUCTION_LSE_VALUES_PER_ROW * FP32_BYTES + ) + stats_src = cutlass.inttoptr( + g_partial_stats.toint() + stats_offset, + mem_space=1, + dtype=Float32, + ) + partial_lse = stats_src.load() + + partial_o_offset = workspace_row * Int64( + cfg.headdim * PARTIAL_O_ELEMENT_BYTES + ) + Int64(reduce_col_idx) * Int64(PARTIAL_O_ELEMENT_BYTES) + partial_o_src = cutlass.inttoptr( + g_partial_o.toint() + partial_o_offset, + mem_space=1, + dtype=Int32, + ) + loaded_partial_regs = partial_o_src.load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + + new_lse, old_weight, partial_weight = merge_log2_lse( + global_lse, + partial_lse, + ) + partial_vals = unpack_normalized_vec8( + loaded_partial_regs, cfg.use_bf16_separate_partial_o + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + output_vals[elem_idx] = ( + output_vals[elem_idx] * old_weight + + partial_vals[elem_idx] * partial_weight + ) + global_lse = new_lse + + _store_parallel_reduction_output( + output_vals, + global_lse, + o_iter, + g_attention_sinks, + output_row_idx, + reduce_col_idx, + h_k_idx, + attention_sink_h_r, + grid_h_k, + attention_sink_head_idx, + cfg, + ) + + +@cute.kernel +def decode_gen_separate_reduction_kernel( + o_iter: cute.Pointer, + g_seqlens_kv: cute.Pointer, + g_cu_seqlens_q: cute.Pointer, + g_partial_o: cute.Pointer, + g_partial_stats: cute.Pointer, + g_attention_sinks: cute.Pointer, + g_q_output_rows: Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + static_full_split_prefix: cutlass.Constexpr[bool] = False, +) -> None: + """Run the 512-thread exact-split reference schedule.""" + + _reduce_exact_splits_body( + o_iter, + g_seqlens_kv, + g_cu_seqlens_q, + g_partial_o, + g_partial_stats, + g_attention_sinks, + g_q_output_rows, + cfg, + static_full_split_prefix, + ) + + +@cute.jit +def _attention_sink_log2_lse( + cfg: cutlass.Constexpr[FmhaDecodeConfig], + attention_sinks_ptr: cute.Pointer, + logical_h_k_idx: Int32, + h_r: Int32, + num_heads_kv: Int32, + local_head_idx: Int32, +) -> Float32: + """Return the sink's LSE-domain denominator state, or neutral ``-inf``.""" + + if cutlass.const_expr(not cfg.use_attention_sinks): + return Float32(-Float32.inf) + head_idx = cute.math.min( + logical_h_k_idx * h_r + local_head_idx, + h_r * num_heads_kv - Int32(1), + ) + sink_ptr = cutlass.inttoptr( + attention_sinks_ptr.toint() + cutlass.Int64(head_idx * Int32(FP32_BYTES)), + mem_space=1, + dtype=Float32, + ) + sink_lse = sink_ptr.load() * Float32(1.4426950408889634) + if cutlass.const_expr(cfg.use_fp8_qkv): + sink_lse += Float32(math.log2(448.0)) + return sink_lse + + +@cute.jit +def _store_parallel_reduction_output( + output_vals: cutlass.Array, + global_lse: Float32, + o_iter: cute.Pointer, + g_attention_sinks: cute.Pointer, + output_row_idx: Int32, + reduce_col_idx: Int32, + h_k_idx: Int32, + attention_sink_h_r: Int32, + grid_h_k: Int32, + attention_sink_head_idx: Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], +) -> None: + """Merge the optional sink and store one normalized output fragment. + + Producer CTAs already fold the softmax scale into log2-LSE and, for FP8 + output, fold the output quantization scale into normalized partial O. The + standalone reducer therefore only merges normalized states here. + """ + + sink_lse = _attention_sink_log2_lse( + cfg, + g_attention_sinks, + h_k_idx, + attention_sink_h_r, + grid_h_k, + attention_sink_head_idx, + ) + _, split_weight, _ = merge_log2_lse(global_lse, sink_lse) + + final_regs = cutlass.Array( + Int32, + PACKED_OUTPUT_REGS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + if cutlass.const_expr(cfg.use_fp8_output): + for packed_idx in cutlass.range_constexpr(FP8_PACKED_OUTPUT_REGS_PER_THREAD): + val_base = packed_idx * FP8_VALUES_PER_REG + final_regs[packed_idx] = _pack_float4_to_fp8_e4m3( + output_vals[val_base] * split_weight, + output_vals[val_base + 1] * split_weight, + output_vals[val_base + 2] * split_weight, + output_vals[val_base + 3] * split_weight, + ) + else: + for packed_idx in cutlass.range_constexpr(PACKED_OUTPUT_REGS_PER_THREAD): + val_base = packed_idx * FP16_VALUES_PER_REG + final_pair = fmul2( + (split_weight, split_weight), + (output_vals[val_base], output_vals[val_base + 1]), + ) + if cutlass.const_expr(cfg.use_bf16_output): + final_regs[packed_idx] = _pack_float2_to_bf16( + final_pair[0], final_pair[1] + ) + else: + final_regs[packed_idx] = _pack_float2_to_fp16( + final_pair[0], final_pair[1] + ) + + # Widen before the first row-stride product. Large packed-Q batches can + # exceed the signed-32-bit byte range even though each local row and + # column coordinate is individually Int32. + output_row_bytes = Int64(cfg.headdim * cfg.o_dtype_bytes) + dst_offset = Int64(output_row_idx) * output_row_bytes + Int64( + reduce_col_idx + ) * Int64(cfg.o_dtype_bytes) + dst_ptr = cutlass.inttoptr( + o_iter.toint() + dst_offset, + mem_space=1, + dtype=Int32, + ) + if cutlass.const_expr(cfg.use_fp8_output): + dst_ptr.store( + final_regs.data_ptr().load( + count=FP8_PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_REGISTER_BYTES, + ), + alignment=FP8_PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + else: + dst_ptr.store( + final_regs.data_ptr().load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_REGISTER_BYTES, + ), + alignment=PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + + +@cute.kernel +def decode_gen_parallel_separate_reduction_kernel( + o_iter: cute.Pointer, + g_seqlens_kv: cute.Pointer, + g_cu_seqlens_q: cute.Pointer, + g_partial_o: cute.Pointer, + g_partial_stats: cute.Pointer, + g_attention_sinks: cute.Pointer, + g_q_output_rows: Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + static_full_split_prefix: cutlass.Constexpr[bool] = False, +) -> None: + """Reduce split partials with a compact or clustered constexpr schedule. + + S2-S4 use one 512-thread CTA per 8 KiB output slice and reduce the exact + split count directly. Larger schedules use 128-thread CTAs over 2 KiB + slices. Each cluster rank owns 2, 4, or 8 split slots; G1 stores directly, + G16 uses a two-level 4x4 merge, and G2/G4/G8 finalize through rank zero. + """ + + # Pair with the producer's launch-dependents signal before reading any + # partial GMEM. The wait is CTA-convergent and stays outside both schedules. + if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + + if cutlass.const_expr(cfg.use_compact_parallel_reduction): + _reduce_exact_splits_body( + o_iter, + g_seqlens_kv, + g_cu_seqlens_q, + g_partial_o, + g_partial_stats, + g_attention_sinks, + g_q_output_rows, + cfg, + static_full_split_prefix, + ) + return + + thread_idx, _, _ = cute.arch.thread_idx() + block_idx_x, h_k_idx, b_idx = cute.arch.block_idx() + _, grid_h_k, _ = cute.arch.grid_dim() + cluster_rank = cute.arch.block_idx_in_cluster() + logical_kv_idx = Int64(b_idx) * Int64(grid_h_k) + Int64(h_k_idx) + attention_sink_h_r = _attention_sink_head_stride(cfg, g_q_output_rows) + + bytes_per_output_row = Int32(cfg.headdim * PARTIAL_O_ELEMENT_BYTES) + slice_idx = block_idx_x // Int32(cfg.parallel_reduction_cluster_size) + reduce_base_offset = slice_idx * Int32( + PARALLEL_REDUCTION_BYTES_PER_SLICE + ) + thread_idx * Int32(REDUCTION_BYTES_PER_THREAD) + reduce_row_idx = reduce_base_offset // bytes_per_output_row + reduce_col_idx = (reduce_base_offset % bytes_per_output_row) // Int32( + PARTIAL_O_ELEMENT_BYTES + ) + q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) + active_splits_kv = Int32(cfg.splits_kv) + if cutlass.const_expr(not static_full_split_prefix): + active_splits_kv = _reduction_active_splits_kv( + cfg, + g_seqlens_kv, + b_idx, + g_q_output_rows, + reduce_row_idx, + seq_len_q, + ) + valid_reduce_row = _q_logical_output_row_is_valid_for_seq( + cfg, + g_q_output_rows, + reduce_row_idx, + seq_len_q, + ) + output_row_idx = _q_physical_output_row_from_logical( + cfg, + g_q_output_rows, + grid_h_k, + b_idx, + h_k_idx, + reduce_row_idx, + q_token_offset, + ) + attention_sink_head_idx = _local_head_from_q_output_row( + cfg, + g_q_output_rows, + reduce_row_idx, + ) + + output_vals = cutlass.Array( + Float32, OUTPUT_VALUES_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + output_vals[elem_idx] = Float32(0.0) + global_lse = Float32(-Float32.inf) + + # Batch up to four independent GMEM loads before folding them. The final + # batch width is compile-time. Padded split slots stay neutral and never + # form GMEM pointers, including non-power-of-two split counts. + local_splits = cfg.parallel_reduction_splits_per_cta + partial_lse = cutlass.Array( + Float32, PARALLEL_REDUCTION_LOAD_BATCH, space=cutlass.AddressSpace.rmem + ) + partial_regs = cutlass.Array( + Int32, + PARALLEL_REDUCTION_LOAD_BATCH * PACKED_OUTPUT_REGS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for split_base_i in cutlass.range_constexpr( + 0, local_splits, PARALLEL_REDUCTION_LOAD_BATCH + ): + batch_width = min(PARALLEL_REDUCTION_LOAD_BATCH, local_splits - split_base_i) + split_base = Int32(split_base_i) + for jj in cutlass.range_constexpr(batch_width): + split_idx = cluster_rank * Int32(local_splits) + split_base + Int32(jj) + valid_split_idx = cutlass.const_expr( + cfg.parallel_reduction_padded_splits == cfg.max_splits_kv + ) or split_idx < Int32(cfg.max_splits_kv) + active_split_idx = split_idx < active_splits_kv + if cutlass.const_expr(static_full_split_prefix): + active_split_idx = cutlass.Boolean(True) + if valid_split_idx and active_split_idx and valid_reduce_row: + workspace_row = _separate_workspace_row_offset( + logical_kv_idx, + split_idx, + reduce_row_idx, + g_q_output_rows, + cfg, + ) + stats_offset = workspace_row * Int64( + SEPARATE_REDUCTION_LSE_VALUES_PER_ROW * FP32_BYTES + ) + stats_src = cutlass.inttoptr( + g_partial_stats.toint() + stats_offset, + mem_space=1, + dtype=Float32, + ) + partial_lse[jj] = stats_src.load() + + partial_o_offset = workspace_row * Int64( + cfg.headdim * PARTIAL_O_ELEMENT_BYTES + ) + Int64(reduce_col_idx) * Int64(PARTIAL_O_ELEMENT_BYTES) + partial_o_src = cutlass.inttoptr( + g_partial_o.toint() + partial_o_offset, + mem_space=1, + dtype=Int32, + ) + loaded_partial_regs = partial_o_src.load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + regs_base = jj * PACKED_OUTPUT_REGS_PER_THREAD + for reg_idx in cutlass.range_constexpr(PACKED_OUTPUT_REGS_PER_THREAD): + partial_regs[regs_base + reg_idx] = loaded_partial_regs[reg_idx] + + for jj in cutlass.range_constexpr(batch_width): + split_idx = cluster_rank * Int32(local_splits) + split_base + Int32(jj) + valid_split_idx = cutlass.const_expr( + cfg.parallel_reduction_padded_splits == cfg.max_splits_kv + ) or split_idx < Int32(cfg.max_splits_kv) + active_split_idx = split_idx < active_splits_kv + if cutlass.const_expr(static_full_split_prefix): + active_split_idx = cutlass.Boolean(True) + if valid_split_idx and active_split_idx and valid_reduce_row: + regs_base = jj * PACKED_OUTPUT_REGS_PER_THREAD + loaded_partial_regs = (partial_regs.data_ptr() + Int32(regs_base)).load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_REGISTER_BYTES, + ) + new_lse, old_weight, partial_weight = merge_log2_lse( + global_lse, + partial_lse[jj], + ) + partial_vals = unpack_normalized_vec8( + loaded_partial_regs, cfg.use_bf16_separate_partial_o + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + output_vals[elem_idx] = ( + output_vals[elem_idx] * old_weight + + partial_vals[elem_idx] * partial_weight + ) + global_lse = new_lse + + # G1 stores its register accumulator directly and compiles out SMEM + # publication, mapa, and cluster barriers. The branch is CTA-uniform and + # resolved at compile time. + if cutlass.const_expr(cfg.parallel_reduction_cluster_size == 1): + if valid_reduce_row: + _store_parallel_reduction_output( + output_vals, + global_lse, + o_iter, + g_attention_sinks, + output_row_idx, + reduce_col_idx, + h_k_idx, + attention_sink_h_r, + grid_h_k, + attention_sink_head_idx, + cfg, + ) + return + + # Each G2+ thread publishes one normalized ``(LSE, O)`` state using the + # profile's selected 16-bit partial type. Corresponding threads in every + # rank map to the same output row, so row validity is cluster-uniform. A + # rank with only padded split slots publishes the neutral ``(-inf, 0)``. + smem_lse = cutlass.Array( + Float32, + SEPARATE_REDUCTION_LSE_VALUES_PER_ROW * PARALLEL_REDUCTION_THREADS_PER_CTA, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + smem_partial_o = cutlass.Array( + Int32, + PACKED_OUTPUT_REGS_PER_THREAD * PARALLEL_REDUCTION_THREADS_PER_CTA, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + stats_smem_offset = thread_idx * Int32(SEPARATE_REDUCTION_LSE_VALUES_PER_ROW) + partial_o_smem_offset = thread_idx * Int32(PACKED_OUTPUT_REGS_PER_THREAD) + if valid_reduce_row: + (smem_lse.data_ptr() + stats_smem_offset).store(global_lse) + packed_o = cutlass.Array( + Int32, + PACKED_OUTPUT_REGS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for reg_idx in cutlass.range_constexpr(PACKED_OUTPUT_REGS_PER_THREAD): + elem_base = reg_idx * FP16_VALUES_PER_REG + if cutlass.const_expr(cfg.use_bf16_separate_partial_o): + packed_o[reg_idx] = _pack_float2_to_bf16( + output_vals[elem_base], output_vals[elem_base + 1] + ) + else: + packed_o[reg_idx] = _pack_float2_to_fp16( + output_vals[elem_base], output_vals[elem_base + 1] + ) + (smem_partial_o.data_ptr() + partial_o_smem_offset).store( + packed_o.data_ptr().load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_REGISTER_BYTES, + ), + alignment=PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + + # Publish every rank's local state before any distributed-SMEM read. + prims.barrier_cta_sync(0) + prims.barrier_cluster_arrive() + prims.barrier_cluster_wait() + + # For G16, ranks 0/4/8/12 first reduce their own four-rank group and + # overwrite that group's first slot. Keeping each reducer inside its source + # group prevents one reducer from overwriting a slot that another reducer + # may still be reading. G4/G8 skip this level and are consumed directly. + if cutlass.const_expr(cfg.parallel_reduction_cluster_size == 16): + is_stage_leader = ( + cluster_rank % Int32(PARALLEL_REDUCTION_FINAL_REDUCERS) + ) == Int32(0) + if is_stage_leader & valid_reduce_row: + stage_vals = cutlass.Array( + Float32, OUTPUT_VALUES_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + stage_vals[elem_idx] = Float32(0.0) + stage_lse = Float32(-Float32.inf) + peer_base = cluster_rank + for peer_offset_i in cutlass.range_constexpr( + PARALLEL_REDUCTION_FINAL_REDUCERS + ): + peer_rank = peer_base + Int32(peer_offset_i) + peer_lse = prims.mapa(smem_lse.data_ptr(), peer_rank) + peer_partial_o = prims.mapa(smem_partial_o.data_ptr(), peer_rank) + local_lse = (peer_lse + stats_smem_offset).load() + loaded_partial_regs = (peer_partial_o + partial_o_smem_offset).load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + new_lse, old_weight, partial_weight = merge_log2_lse( + stage_lse, + local_lse, + ) + partial_vals = unpack_normalized_vec8( + loaded_partial_regs, cfg.use_bf16_separate_partial_o + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + stage_vals[elem_idx] = ( + stage_vals[elem_idx] * old_weight + + partial_vals[elem_idx] * partial_weight + ) + stage_lse = new_lse + + (smem_lse.data_ptr() + stats_smem_offset).store(stage_lse) + stage_packed_o = cutlass.Array( + Int32, + PACKED_OUTPUT_REGS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for reg_idx in cutlass.range_constexpr(PACKED_OUTPUT_REGS_PER_THREAD): + elem_base = reg_idx * FP16_VALUES_PER_REG + if cutlass.const_expr(cfg.use_bf16_separate_partial_o): + stage_packed_o[reg_idx] = _pack_float2_to_bf16( + stage_vals[elem_base], stage_vals[elem_base + 1] + ) + else: + stage_packed_o[reg_idx] = _pack_float2_to_fp16( + stage_vals[elem_base], stage_vals[elem_base + 1] + ) + (smem_partial_o.data_ptr() + partial_o_smem_offset).store( + stage_packed_o.data_ptr().load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_REGISTER_BYTES, + ), + alignment=PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + + prims.barrier_cta_sync(0) + prims.barrier_cluster_arrive() + prims.barrier_cluster_wait() + + final_input_partials = cfg.parallel_reduction_cluster_size + if cutlass.const_expr(cfg.parallel_reduction_cluster_size == 16): + final_input_partials = PARALLEL_REDUCTION_FINAL_REDUCERS + + if (cluster_rank == Int32(0)) & valid_reduce_row: + final_vals = cutlass.Array( + Float32, OUTPUT_VALUES_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + final_vals[elem_idx] = Float32(0.0) + final_lse = Float32(-Float32.inf) + for peer_rank_i in cutlass.range_constexpr(final_input_partials): + peer_rank = Int32(peer_rank_i) + if cutlass.const_expr(cfg.parallel_reduction_cluster_size == 16): + peer_rank = Int32(peer_rank_i * PARALLEL_REDUCTION_FINAL_REDUCERS) + peer_lse = prims.mapa(smem_lse.data_ptr(), peer_rank) + peer_partial_o = prims.mapa(smem_partial_o.data_ptr(), peer_rank) + local_lse = (peer_lse + stats_smem_offset).load() + loaded_partial_regs = (peer_partial_o + partial_o_smem_offset).load( + count=PACKED_OUTPUT_REGS_PER_THREAD, + alignment=PACKED_OUTPUT_REGS_PER_THREAD * PACKED_REGISTER_BYTES, + ) + new_lse, old_weight, partial_weight = merge_log2_lse( + final_lse, + local_lse, + ) + partial_vals = unpack_normalized_vec8( + loaded_partial_regs, cfg.use_bf16_separate_partial_o + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + final_vals[elem_idx] = ( + final_vals[elem_idx] * old_weight + + partial_vals[elem_idx] * partial_weight + ) + final_lse = new_lse + + _store_parallel_reduction_output( + final_vals, + final_lse, + o_iter, + g_attention_sinks, + output_row_idx, + reduce_col_idx, + h_k_idx, + attention_sink_h_r, + grid_h_k, + attention_sink_head_idx, + cfg, + ) + + # Keep every peer CTA alive until rank zero has finished its DSMEM reads. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + + +@cute.jit +def fmha_decode_separate_reduction_launch( + problem_shape: tuple[Int32, Int32, Int32, Int32, Int32], + o_iter: cute.Pointer, + seqlens_kv_iter: cute.Pointer, + cu_seqlens_q_iter: cute.Pointer, + partial_o_iter: cute.Pointer, + partial_stats_iter: cute.Pointer, + attention_sinks_iter: cute.Pointer, + scale_s: Float32, + output_scale: Float32, + stream: cuda_drv.CUstream, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + static_full_split_prefix: cutlass.Constexpr[bool] = False, +) -> None: + """Launch the standalone reducer after the main split-KV decode kernel. + + The serial reference grid is ``(slice, kv_head, batch)`` with 8 KiB slices. + The production reducer keeps that geometry for compact S2-S4 and otherwise + maps cluster ranks into grid.x for each 2 KiB slice. ``scale_s`` and + ``output_scale`` remain parameters for launch-ABI compatibility; producer + CTAs already applied them to log2-LSE and normalized partial O, respectively. + ``seqlens_kv_iter`` is internal reducer metadata used only to + bound the runtime split prefix; the decode launch ABI remains unchanged. + """ + b, h_q, h_k, _, _ = problem_shape + h_r = h_q // h_k + q_output_rows = h_r + if cutlass.const_expr(cfg.max_seq_len_q > 1): + q_output_rows = h_r * Int32(cfg.max_seq_len_q) + bytes_per_slice = REDUCTION_BYTES_PER_SLICE + grid_q_output_rows = q_output_rows + if cutlass.const_expr(cfg.max_seq_len_q > 1 and cfg.heads_q_per_kv != 0): + # Grouped multi-token Q profiles lay rows out by the configured + # heads-per-KV group instead of the launch-time h_q / h_k ratio. + grid_q_output_rows = cfg.heads_q_per_kv * cfg.max_seq_len_q + # The reference schedule uses 512-thread CTAs over contiguous 8 KiB slices. + # Multiple CTAs cover wide dimensions or grouped SQ rows without changing + # the producer workspace layout. + num_reduction_slices = max( + ( + grid_q_output_rows * cfg.headdim * PARTIAL_O_ELEMENT_BYTES + + bytes_per_slice + - 1 + ) + // bytes_per_slice, + 1, + ) + if cutlass.const_expr(cfg.use_parallel_separate_reduction): + cluster_size = cfg.parallel_reduction_cluster_size + parallel_bytes_per_slice = cfg.parallel_reduction_bytes_per_slice + num_parallel_slices = max( + ( + grid_q_output_rows * cfg.headdim * PARTIAL_O_ELEMENT_BYTES + + parallel_bytes_per_slice + - 1 + ) + // parallel_bytes_per_slice, + 1, + ) + decode_gen_parallel_separate_reduction_kernel( + o_iter, + seqlens_kv_iter, + cu_seqlens_q_iter, + partial_o_iter, + partial_stats_iter, + attention_sinks_iter, + q_output_rows, + cfg, + static_full_split_prefix, + ).launch( + grid=(num_parallel_slices * cluster_size, h_k, b), + block=[cfg.parallel_reduction_threads_per_cta, 1, 1], + cluster=[cluster_size, 1, 1], + stream=stream, + use_pdl=cfg.use_parallel_separate_reduction_pdl, + ) + return + decode_gen_separate_reduction_kernel( + o_iter, + seqlens_kv_iter, + cu_seqlens_q_iter, + partial_o_iter, + partial_stats_iter, + attention_sinks_iter, + q_output_rows, + cfg, + static_full_split_prefix, + ).launch( + grid=(num_reduction_slices, h_k, b), + # Reduction parallelism is entirely in block.x; y/z are singleton + # dimensions because grid.y/grid.z already carry head and batch. + block=[REDUCTION_THREADS_PER_CTA, 1, 1], + cluster=[1, 1, 1], + stream=stream, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mask.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mask.py new file mode 100644 index 000000000000..b7096a3ca22c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mask.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared masking predicates for PrimTS attention kernels.""" + + +def kv_tile_needs_right_mask(tile_offset_k, tile_size_kv, visible_k_end): + """Return whether a KV tile crosses a row-visible right bound.""" + return tile_offset_k + tile_size_kv > visible_k_end + + +def kv_tile_is_fully_visible( + tile_offset_k, + tile_size_kv, + visible_k_begin, + visible_k_end, +): + """Return whether a KV tile is inside every row's visible interval. + + All intervals are half-open. The arguments may be Python integers in + host-side tests or CuTe DSL integer values while tracing a kernel. + """ + return (tile_offset_k >= visible_k_begin) & ( + tile_offset_k + tile_size_kv <= visible_k_end + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/__init__.py new file mode 100644 index 000000000000..4f1489736ac6 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task-scheduled MLA decode kernels.""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/__init__.py new file mode 100644 index 000000000000..a4011a2e150c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared helpers for task-scheduled MLA decode kernels.""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py new file mode 100644 index 000000000000..aa68b9efa0f1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared instruction and layout constants for MLA decode TS examples.""" + +# CUDA warp geometry used by lane ownership, butterfly reductions, and TMEM row +# addressing. Keep masks and shifts named so instruction-level code does not +# hide the warp contract behind raw bit constants. +WARP_LANES = 32 +WARP_LANE_MASK = WARP_LANES - 1 +WARP_LANE_SHIFT = 5 +HALF_WARP_LANES = 16 +HALF_WARP_MASK = HALF_WARP_LANES - 1 +QUAD_LANES = 4 +QUAD_LANE_MASK = QUAD_LANES - 1 +QUAD_LANE_SHIFT = 2 +OCTET_LANES = 8 +OCTET_LANE_MASK = OCTET_LANES - 1 +WARP_REDUCTION_BFLY_DISTANCES = (16, 8, 4, 2, 1) + +# One warpgroup is four warps. Several softmax/P/O protocols synchronize one +# full warpgroup through a named CTA barrier. +WARPGROUP_WARPS = 4 +WARPGROUP_THREADS = WARPGROUP_WARPS * WARP_LANES + +# Kernel-level TMEM lifecycle synchronization. Barrier ID 15 is reserved for +# alloc/dealloc phases in the 1CTA path; the 2CTA path uses config-owned IDs. +TMEM_LIFECYCLE_BARRIER_ID = 15 +TMEM_DEALLOC_MBAR_THREADS = WARP_LANES + +# STSM and vector-copy helpers share the tcgen05 128B-row, 16B-column swizzled +# SMEM layout. The shift constants are the log2 form of the byte units used in +# address calculations. +SMEM_ROW_BYTES = 128 +SMEM_VECTOR_BYTES = 16 +SMEM_WORD_BYTES = 4 +SMEM_ROW_BYTE_SHIFT = 7 +SMEM_VECTOR_BYTE_SHIFT = 4 +SMEM_WORD_BYTE_SHIFT = 2 +STSM_MATRIX_LANES = OCTET_LANES +STSM_MATRIX_LANE_SHIFT = 3 +STSM_X4_REG_COUNT = 4 +STSM_MATRICES_PER_WARP = QUAD_LANES +STSM_MATRICES_PER_WARP_SHIFT = 2 +STSM_WARPS_PER_SLICE = 2 +STSM_WARPS_PER_SLICE_SHIFT = 1 +STSM_ROW_BLOCK_ROWS = 16 +O_STAGE_COPY_SEGMENT_BYTES = 2048 +SWIZZLE_ROW_MASK = OCTET_LANE_MASK + +# tcgen05 TMEM load/store instruction geometry used by MLA softmax, correction, +# and epilogue paths. +TCGEN05_32B_SHAPE = "32x32b" +TCGEN05_16X256B_SHAPE = "16x256b" +TCGEN05_32B_REGS_PER_LOAD = 32 +TCGEN05_16X256B_REGS_PER_LOAD = 4 +TCGEN05_16X32BX2_BF16_P_STRIDE = 32 +TCGEN05_16X32BX2_FP8_P_STRIDE = 16 +TCGEN05_SECOND_PANEL_ADDR_OFFSET = 16 << 16 + +# Softmax scratch stores the max state first and the sum state in a second +# fixed-size panel. The offset is in Uint32 scratch words. +SOFTMAX_SCRATCH_SUM_WORD_OFFSET = 384 +SCORE_ROWS_PER_Q_PAIR = 2 +SCORE_TOKENS_PER_QK_GROUP = 8 + +# CTA-local split-reduction barrier used after the scale-writing warp publishes +# per-split LSE rescale factors in SMEM. +SPLIT_REDUCTION_SCALE_BARRIER_ID = 4 + +# SmemPResource uses adjacent CTA barriers for the two FP8 P producer instances +# after byte-transposed STSM stores have published their SMEM payloads. +SMEM_P_FP8_STORE_BARRIER_BASE_ID = 4 + +# Page-offset entries are Int32 page IDs staged through cp.async. +PAGE_OFFSET_BYTES = 4 +CP_ASYNC_CACHE_CA = "ca" + +# Dense MLA kernels specialize paged-KV addressing for this explicit ABI set. +# Each listed size exactly partitions a 128-token KV tile; other page sizes are +# outside the current kernel ABI. +SUPPORTED_MLA_PAGE_SIZES = (16, 32, 64, 128) + +# Both MLA launch families and their standalone reducers share one fixed +# workspace/scheduler split capacity. +MAX_MLA_SPLITS_KV = 128 + +# Throughput 2CTA epilogue maps 128 local threads onto two 64-row groups and +# 128-column output halves for vectorized GMEM publication. +EPILOGUE_THREAD_TILE_THREADS = 128 +EPILOGUE_THREAD_TILE_MASK = EPILOGUE_THREAD_TILE_THREADS - 1 +EPILOGUE_ROW_THREADS = 64 +EPILOGUE_ROW_MASK = EPILOGUE_ROW_THREADS - 1 +EPILOGUE_COLUMN_GROUP_SHIFT = 6 +BF16_OUTPUT_VECTOR_ELEMENTS = 8 +FP8_OUTPUT_VECTOR_ELEMENTS = 16 +PACKED_FP8_OUTPUT_REGS = 4 + +# tcgen05 SMEM descriptors are advanced in 16B units. The normal Q/K next-K +# increment is two units; the wrapped K descriptor moves backward across the +# descriptor ring to the next logical 128-wide head-dim block. +TCGEN05_DESC_NEXT_K_BLOCK_UNITS = 2 +TCGEN05_DESC_WRAPPED_K_BLOCK_UNITS = 1018 + +# Tensor-map descriptors and workspace allocations use 1024B alignment because +# the backing STensor/TMA views are 1 KiB aligned in these kernels. +TMA_DESCRIPTOR_ALIGNMENT_BYTES = 1024 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py new file mode 100644 index 000000000000..fe45148142c3 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py @@ -0,0 +1,321 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Register, TMEM, and SMEM layout helpers for MLA decode.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + +from cutlass.experimental.task_scheduling.memory import SmemAllocation +from cutlass.experimental.task_scheduling.resources import StageInfo + +from ..throughput_latency_1cta.config import MlaConfig +from .constants import ( + O_STAGE_COPY_SEGMENT_BYTES, + SMEM_ROW_BYTES, + SMEM_ROW_BYTE_SHIFT, + SMEM_VECTOR_BYTES, + SMEM_VECTOR_BYTE_SHIFT, + STSM_MATRICES_PER_WARP, + STSM_MATRICES_PER_WARP_SHIFT, + STSM_MATRIX_LANES, + STSM_MATRIX_LANE_SHIFT, + STSM_ROW_BLOCK_ROWS, + STSM_WARPS_PER_SLICE, + STSM_WARPS_PER_SLICE_SHIFT, + SWIZZLE_ROW_MASK, + TCGEN05_SECOND_PANEL_ADDR_OFFSET, +) +from .math import ceil_div, mma_k_step_for_qkv + + +# Per-task cache tuple slots. Tasks share these indices when passing common +# CTA-local coordinates and runtime sequence length into resource work methods. +_TASK_CACHE_TMEM_BASE_OFFSET = 0 +_TASK_CACHE_WARP_GRP_THREAD_IDX = 1 +_TASK_CACHE_WARP_IDX = 2 +_TASK_CACHE_LANE_IDX = 3 +_TASK_CACHE_SEQ_LEN_KV = 4 + + +def num_softmax_scale_groups(cfg: MlaConfig) -> int: + """Return the number of independent softmax scale groups per thread.""" + if cfg.kernel_variant == "keeps_mma_ab": + return 1 + return max(cfg.tile_size_q // 4, 4) + + +def num_s_regs_per_thread(cfg: MlaConfig) -> int: + """Return the number of score registers carried by each softmax thread.""" + if cfg.kernel_variant == "keeps_mma_ab": + return cfg.tile_size_kv // 2 + return num_softmax_scale_groups(cfg) * 4 + + +def softmax_scratch_words(cfg: MlaConfig) -> int: + """Return the shared scratch words needed for softmax reductions.""" + if cfg.kernel_variant == "keeps_mma_ab": + return 2 * 384 + return 4 * num_softmax_scale_groups(cfg) + + +def num_packed_p_regs(cfg: MlaConfig) -> int: + """Return the number of packed P registers per softmax thread.""" + q_repeats = max(cfg.tile_size_q // 8, 1) + return (2 if cfg.is_fp8_qkv() else 4) * q_repeats + + +def num_o_reg_pairs(cfg: MlaConfig) -> int: + """Return the number of packed O register pairs handled per thread.""" + return 2 * max(max(cfg.tile_size_q, 16) // 8, 1) + + +def num_o_repeats(cfg: MlaConfig) -> int: + """Return the O repeat count for TMEM loads and GMEM stores.""" + return max(max(cfg.tile_size_q, 16) // 8, 1) + + +def num_q_repeats(cfg: MlaConfig) -> int: + """Return the Q/head repeat count for a warpgroup row tile.""" + return max(cfg.tile_size_q // 8, 1) + + +def num_o_stsm_row_blocks(cfg: MlaConfig) -> int: + """Return the number of 16-row STSM blocks needed for one O tile.""" + return max(cfg.tile_size_q // 16, 1) + + +def num_o_tmem_loads_per_stage(cfg: MlaConfig) -> int: + """Return the number of TMEM O load groups per output stage.""" + bf16_equivalent_bytes = cfg.tile_size_q * cfg.head_dim_per_stage_v * 2 + bf16_equivalent_segments = max(1, ceil_div(bf16_equivalent_bytes, 2048)) + return max(bf16_equivalent_segments // num_o_stsm_row_blocks(cfg), 1) + + +def num_fp8_output_regs(cfg: MlaConfig) -> int: + """Return the number of packed FP8 output registers per stage.""" + return max((cfg.tile_size_q * cfg.head_dim_per_stage_v) // 512, 1) + + +def q_p_desc_k_block_wrap_bytes(cfg: MlaConfig) -> int: + """Return Q/P descriptor K-wrap distance in bytes.""" + k_block_bytes = mma_k_step_for_qkv(cfg) * cfg.qkv_dtype_bytes + k_blocks_per_smem_row = 128 // k_block_bytes + return cfg.tile_size_q * 128 - (k_blocks_per_smem_row - 1) * k_block_bytes + + +def q_p_desc_k_block_wrap_units(cfg: MlaConfig) -> int: + """Return Q/P descriptor K-wrap distance in 16-byte units.""" + return q_p_desc_k_block_wrap_bytes(cfg) // 16 + + +def tma_inner_dim_elems(cfg: MlaConfig) -> int: + """Return the maximum contiguous TMA inner dimension in Q/K/V elements.""" + return 128 // cfg.qkv_dtype_bytes + + +def tma_page_token_elems(cfg: MlaConfig) -> int: + """Return the number of page-offset tokens represented by one TMA page.""" + return cfg.num_tokens_per_page + + +@cute.jit +def p_stsm_smem_offset_bytes( + local_warp_idx: Int32, + lane_idx: Int32, + stsm_group_idx: int = 0, + stsm_row_block_idx: int = 0, + tile_size_q: int = 16, +): + """Return the SMEM byte offset for one packed P stmatrix store lane.""" + # STSM maps four 8-lane matrices per warp. Offsets below express the + # tcgen05 128B-row / 16B-column swizzled SMEM layout used for packed P. + slice_idx = local_warp_idx // Int32(STSM_WARPS_PER_SLICE) + warp_idx_in_slice = local_warp_idx % Int32(STSM_WARPS_PER_SLICE) + mtx_idx = lane_idx // Int32(STSM_MATRIX_LANES) + thr_row_idx = lane_idx % Int32(STSM_MATRIX_LANES) + if cutlass.const_expr(tile_size_q == 8): + mtx_row_idx = Int32(0) + mtx_col_idx = warp_idx_in_slice * Int32(STSM_MATRICES_PER_WARP) + mtx_idx + else: + mtx_row_idx = mtx_idx // Int32(STSM_WARPS_PER_SLICE) + mtx_col_idx = ( + warp_idx_in_slice * Int32(STSM_MATRICES_PER_WARP) + + (mtx_idx % Int32(STSM_WARPS_PER_SLICE)) + + Int32(stsm_group_idx * STSM_WARPS_PER_SLICE) + ) + return ( + slice_idx * Int32(tile_size_q * SMEM_ROW_BYTES) + + Int32(stsm_row_block_idx * STSM_ROW_BLOCK_ROWS * SMEM_ROW_BYTES) + + mtx_row_idx * Int32(STSM_MATRIX_LANES * SMEM_ROW_BYTES) + + thr_row_idx * Int32(SMEM_ROW_BYTES) + + ((mtx_col_idx ^ thr_row_idx) * Int32(SMEM_VECTOR_BYTES)) + ) + + +@cute.jit +def o_stage_stsm_and_copy_offsets( + cfg: MlaConfig, + warp_grp_thread_idx: Int32, + local_warp_idx: Int32, + lane_idx: Int32, + stsm_group_idx: int = 0, + copy_segment_idx: int = 0, +): + """Return SMEM store, load, row, and column offsets for one O stage.""" + # The copy path treats each thread as a 16B vector lane. A 2048B segment is + # one 16-row x 128B SMEM tile, matching the STSM/TMEM load granularity. + base_offset = (warp_grp_thread_idx << Int32(SMEM_VECTOR_BYTE_SHIFT)) + Int32( + copy_segment_idx * O_STAGE_COPY_SEGMENT_BYTES + ) + smem_row_idx = base_offset >> Int32(SMEM_ROW_BYTE_SHIFT) + thr_row_idx = lane_idx & Int32(SWIZZLE_ROW_MASK) + mtx_idx = lane_idx >> Int32(STSM_MATRIX_LANE_SHIFT) + # Swizzle the 16B vector column by the low three row bits to match the 128B + # shared-memory layout expected by vectorized GMEM stores. + load_smem_offset = base_offset ^ ( + (smem_row_idx & Int32(SWIZZLE_ROW_MASK)) << Int32(SMEM_VECTOR_BYTE_SHIFT) + ) + if cutlass.const_expr( + cfg.head_dim_per_stage_v * cfg.partial_o_dtype_bytes > SMEM_ROW_BYTES + ): + slice_idx = local_warp_idx >> Int32(STSM_WARPS_PER_SLICE_SHIFT) + warp_idx_in_slice = local_warp_idx & Int32(STSM_WARPS_PER_SLICE - 1) + if cutlass.const_expr(cfg.tile_size_q == 8): + mtx_row_idx = Int32(0) + mtx_col_idx = ( + warp_idx_in_slice << Int32(STSM_MATRICES_PER_WARP_SHIFT) + ) + mtx_idx + else: + mtx_row_idx = mtx_idx >> Int32(STSM_WARPS_PER_SLICE_SHIFT) + mtx_col_idx = ( + (warp_idx_in_slice << Int32(STSM_MATRICES_PER_WARP_SHIFT)) + + (mtx_idx & Int32(STSM_WARPS_PER_SLICE - 1)) + + Int32(stsm_group_idx * STSM_WARPS_PER_SLICE) + ) + smem_offset_bytes = ( + Int32(copy_segment_idx * O_STAGE_COPY_SEGMENT_BYTES) + + slice_idx * Int32(cfg.tile_size_q * SMEM_ROW_BYTES) + + (mtx_row_idx * Int32(STSM_MATRIX_LANES) + thr_row_idx) + * Int32(SMEM_ROW_BYTES) + + ((mtx_col_idx ^ thr_row_idx) * Int32(SMEM_VECTOR_BYTES)) + ) + dst_row_idx = smem_row_idx % Int32(cfg.tile_size_q) + dst_col_offset = (smem_row_idx // Int32(cfg.tile_size_q)) * Int32( + SMEM_ROW_BYTES + ) + (base_offset & Int32(SMEM_ROW_BYTES - 1)) + else: + mtx_row_idx = mtx_idx >> Int32(STSM_WARPS_PER_SLICE_SHIFT) + mtx_col_idx = mtx_idx & Int32(STSM_WARPS_PER_SLICE - 1) + seg_col_idx = ( + (local_warp_idx << Int32(STSM_WARPS_PER_SLICE_SHIFT)) + mtx_col_idx + ) ^ thr_row_idx + if cutlass.const_expr(stsm_group_idx != 0): + seg_col_idx = seg_col_idx + Int32(stsm_group_idx * STSM_WARPS_PER_SLICE) + smem_offset_bytes = ( + mtx_row_idx * Int32(STSM_MATRIX_LANES) + thr_row_idx + ) * Int32(SMEM_ROW_BYTES) + seg_col_idx * Int32(SMEM_VECTOR_BYTES) + dst_row_idx = smem_row_idx + dst_col_offset = base_offset & Int32(SMEM_ROW_BYTES - 1) + return smem_offset_bytes, load_smem_offset, dst_row_idx, dst_col_offset + + +@cute.jit +def local_q_head_idx_for_scale( + cfg: MlaConfig, + col_group_idx: Int32, + scale_idx, +): + """Return the local Q/head row controlled by a softmax scale group.""" + # Scale groups are interleaved as pairs across columns, then advanced in + # groups of eight rows to match the packed softmax register layout. + local_head_idx = col_group_idx * Int32(STSM_WARPS_PER_SLICE) + ( + Int32(scale_idx) & Int32(STSM_WARPS_PER_SLICE - 1) + ) + local_head_idx = local_head_idx + ( + Int32(scale_idx) >> Int32(STSM_WARPS_PER_SLICE_SHIFT) + ) * Int32(STSM_MATRIX_LANES) + return local_head_idx + + +def q_stage_elements(cfg: MlaConfig, stage_idx: int) -> int: + """Return the number of Q elements stored for one QK head-dim stage.""" + return cfg.tile_size_q * cfg.qk_head_stage_width(stage_idx) + + +def kv_stage_elements(cfg: MlaConfig, stage_idx: int) -> int: + """Return the number of K elements stored for one QK head-dim stage.""" + return cfg.tile_size_kv * cfg.qk_head_stage_width(stage_idx) + + +def v_stage_elements(cfg: MlaConfig, stage_idx: int) -> int: + """Return the number of V elements stored for one V head-dim stage.""" + return cfg.tile_size_kv * cfg.v_head_stage_width(stage_idx) + + +def q_stage_smem_element_offset(cfg: MlaConfig, stage_idx: int) -> int: + """Return the Q SMEM element offset for one QK head-dim stage.""" + return stage_idx * cfg.tile_size_q * cfg.head_dim_per_stage_kv + + +def kv_stage_smem_element_offset(cfg: MlaConfig, stage_idx: int) -> int: + """Return the K/V SMEM element offset for one QK head-dim stage.""" + return stage_idx * cfg.tile_size_kv * cfg.head_dim_per_stage_kv + + +@cute.jit +def head_dim_cta_offset_v(cfg: MlaConfig, cta_idx_head_dim_v): + """Return the V head-dim element offset for a split head-dim CTA.""" + if cutlass.const_expr(cta_idx_head_dim_v is None): + return Int32(0) + return Int32(cta_idx_head_dim_v) * Int32(cfg.head_dim_per_cta_v) + + +def o_stage_tmem_col_offset( + cfg: MlaConfig, + o_stage_idx, + v_stage_idx: int, +): + """Return the TMEM column offset for one O pipeline and V head-dim stage.""" + if cfg.kernel_variant == "keeps_mma_ab" and cfg.head_dim_per_cta_v > 256: + return ( + o_stage_idx * Int32(2 * cfg.tmem_o_buffer_cols) + + Int32((v_stage_idx % 2) * cfg.tmem_o_buffer_cols) + + Int32((v_stage_idx // 2) * TCGEN05_SECOND_PANEL_ADDR_OFFSET) + ) + return o_stage_idx * Int32(cfg.tmem_o_buffer_cols * cfg.v_head_dim_stages) + Int32( + v_stage_idx * cfg.tmem_o_buffer_cols + ) + + +def smem_array(context, alloc: SmemAllocation, dtype, count: int): + """Create a shared-memory array view for an allocation, or None if absent.""" + if context is None or context.smem_base is None or alloc is None: + return None + return cutlass.Array( + context.smem_base.data_ptr() + alloc.offset, + dtype=dtype, + shape=(count,), + addrspace=3, + ) + + +@cute.jit +def decode_gen_task_cache(stage_info: StageInfo): + """Return the cached per-task thread values, or a zero cache fallback.""" + if cutlass.const_expr(stage_info.task_cache is None): + zero = Int32(0) + return (zero, zero, zero, zero, zero, zero, zero, zero) + return stage_info.task_cache diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/mask.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/mask.py new file mode 100644 index 000000000000..d795dfd2b59d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/mask.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Mask policy shared by the MLA decode kernels and reference path.""" + +from enum import Enum + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + + +class MaskType(str, Enum): + """Supported speculative-decode attention masks.""" + + CAUSAL = "causal" + DENSE = "dense" + + +def normalize_mask_type(mask_type: MaskType | str) -> str: + """Return the canonical constexpr-safe string for ``mask_type``.""" + + try: + return MaskType(mask_type).value + except (TypeError, ValueError) as error: + supported = ", ".join(mask.value for mask in MaskType) + raise ValueError( + f"mask_type must be one of ({supported}), got {mask_type!r}" + ) from error + + +@cute.jit +def mask_visible_k_length( + mask_type: cutlass.Constexpr[str], + seq_len_kv, + logical_q_idx, + logical_seq_len_q, +): + """Return the dense or bottom-right-causal K length for one Q row.""" + + seq_len_kv = Int32(seq_len_kv) + if cutlass.const_expr(mask_type == MaskType.DENSE.value): + return cute.math.max(seq_len_kv, Int32(0)) + safe_seq_len_q = cute.math.max(Int32(logical_seq_len_q), Int32(1)) + safe_q_idx = cute.math.max( + Int32(0), + cute.math.min(Int32(logical_q_idx), safe_seq_len_q - Int32(1)), + ) + return cute.math.max( + seq_len_kv - (safe_seq_len_q - Int32(1) - safe_q_idx), + Int32(0), + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py new file mode 100644 index 000000000000..e07a06325518 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py @@ -0,0 +1,216 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Math, dtype, and atomic helper functions for MLA decode TS examples.""" + +from functools import partial + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Uint32 +from cutlass.experimental import primitives as prims + +# Softmax converts exp2 inputs with log2(e) and initializes masked scores to +# negative Float32 max. +LOG2_E = 1.4426950408889634074 +NEG_FLT_MAX = -3.4028235e38 + +# E4M3FN finite maximum used when clamping FP8 output conversion. +FP8_E4M3_MAX = 448.0 + +fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd="rn") +add_packed_f32x2 = partial(cute.arch.add_packed_f32x2, rnd="rn") +mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn") + +fadd2 = partial(cute.arch.add_packed_f32x2, ftz=False, rnd="rn") +fmul2 = partial(cute.arch.mul_packed_f32x2, ftz=False, rnd="rn") +ffma2 = partial(cute.arch.fma_packed_f32x2, ftz=False, rnd="rn") + + +def ceil_div(a, b): + """Return integer ceil(a / b) for positive integer-like values.""" + return (a + b - 1) // b + + +def qkv_dtype(cfg): + """Return the shared Q/K/V element dtype for MLA decode examples.""" + if cfg.is_fp8_qkv(): + return cutlass.Float8E4M3FN + return cutlass.BFloat16 + + +def output_dtype(cfg): + """Return the output element dtype for MLA decode examples.""" + if getattr(cfg, "use_fp8_output", 0) == 1: + return cutlass.Float8E4M3FN + return cutlass.BFloat16 + + +def partial_output_dtype(cfg): + """Return the split-KV partial-O dtype used before the final reduction.""" + if getattr(cfg, "num_ctas_per_seq_kv", 1) > 1: + return cutlass.BFloat16 + return output_dtype(cfg) + + +def qkv_smem_swizzle(cfg): + """Return the SMEM swizzle used by staged Q/K/V tensors.""" + if cfg.is_fp8_qkv() and cfg.head_dim_per_stage_kv == 64: + return prims.Tcgen05SmemSwizzle.SWIZZLE_64B + return prims.Tcgen05SmemSwizzle.SWIZZLE_128B + + +def qkv_smem_swizzle_for_head_dim(cfg, head_dim: int): + """Return the staged Q/K/V SMEM swizzle for one head-dimension width.""" + if cfg.is_fp8_qkv() and head_dim == 64: + return prims.Tcgen05SmemSwizzle.SWIZZLE_64B + return prims.Tcgen05SmemSwizzle.SWIZZLE_128B + + +def mma_kind_for_qkv(cfg): + """Return the tcgen05 MMA kind for QK and PV operations.""" + if cfg.is_fp8_qkv(): + return prims.Tcgen05MMAKind.F8F6F4 + return prims.Tcgen05MMAKind.F16 + + +def mma_k_step_for_qkv(cfg) -> int: + """Return the K step consumed by one tcgen05 MMA instruction.""" + return 32 if cfg.is_fp8_qkv() else 16 + + +def qkv_major_k_stride_bytes_for(cfg, head_dim: int) -> int: + """Return the shared-memory major-K stride for one Q/K/V head dimension.""" + num_smem_cols = 128 // cfg.qkv_dtype_bytes + rows_per_smem_row = max(1, num_smem_cols // head_dim) + if rows_per_smem_row == 1: + rows_per_swizzle_block = 8 + elif rows_per_smem_row == 2: + rows_per_swizzle_block = 4 + elif rows_per_smem_row == 4: + rows_per_swizzle_block = 2 + else: + rows_per_swizzle_block = 1 + return 128 * rows_per_swizzle_block + + +def qk_desc_layout_for_head_dim(cfg, head_dim: int): + """Return the UMMA descriptor layout for a QK operand head dimension.""" + if cfg.is_fp8_qkv() and head_dim == 64: + return 4 + return 2 + + +def qk_desc_layout(cfg): + """Return the UMMA descriptor layout for the latent QK operand.""" + return qk_desc_layout_for_head_dim(cfg, cfg.mma_qk_tiler_k) + + +def qk_desc_leading_byte_offset_for_head_dim(cfg, tile_rows: int, head_dim: int) -> int: + """Return the descriptor leading byte offset for one QK head dimension.""" + if cfg.is_fp8_qkv(): + return tile_rows * head_dim * cfg.qkv_dtype_bytes + return 16 + + +def qk_desc_leading_byte_offset(cfg) -> int: + """Return the descriptor leading byte offset for latent QK.""" + return qk_desc_leading_byte_offset_for_head_dim( + cfg, cfg.mma_qk_tiler[0] // cfg.num_mma_ctas, cfg.mma_qk_tiler_k + ) + + +def qk_desc_stride_byte_offset_for_head_dim(cfg, head_dim: int) -> int: + """Return the descriptor stride byte offset for one QK head dimension.""" + if cfg.is_fp8_qkv(): + return qkv_major_k_stride_bytes_for(cfg, head_dim) + return 1024 + + +def qk_desc_stride_byte_offset(cfg) -> int: + """Return the descriptor stride byte offset for latent QK.""" + return qk_desc_stride_byte_offset_for_head_dim(cfg, cfg.mma_qk_tiler_k) + + +def p_desc_layout(cfg): + """Return the UMMA descriptor layout for P in SMEM.""" + del cfg + return 4 + + +def p_desc_leading_byte_offset(cfg) -> int: + """Return the descriptor leading byte offset for P in SMEM.""" + del cfg + return 16 + + +def p_desc_stride_byte_offset(cfg) -> int: + """Return the descriptor stride byte offset for P in SMEM.""" + del cfg + return 512 + + +def neg_max_f32(): + """Return the sentinel negative max value used by online softmax.""" + return Float32(NEG_FLT_MAX) + + +@cute.jit +def pack_float2_to_bf16(v0: Float32, v1: Float32): + """Pack two Float32 values into one Int32 containing two BF16 lanes.""" + return ( + cutlass.Vector.from_elements((v0, v1), Float32) + .to(cutlass.BFloat16) + .bitcast(Int32)[0] + ) + + +@cute.jit +def float_to_u32_for_atomic_max(val: Float32): + """Encode a Float32 value so unsigned atomic max preserves float ordering.""" + bits = prims.mov_b32(val, target_type=Int32) + mask = (bits >> Int32(31)) | Int32(0x80000000) + encoded = bits ^ mask + return prims.mov_b32(encoded, target_type=Uint32) + + +@cute.jit +def u32_to_float_for_atomic_max(val: Uint32): + """Decode an unsigned atomic-max ordered value back into Float32.""" + encoded = prims.mov_b32(val, target_type=Int32) + mask = (~(encoded >> Int32(31))) | Int32(0x80000000) + bits = encoded ^ mask + return prims.mov_b32(bits, target_type=Float32) + + +@cute.jit +def smem_atomic_max_u32(ptr, val: Uint32): + """Apply a CTA-scoped unsigned max atomic to shared memory.""" + prims.atomicrmw( + prims.AtomicOp.MAX, + ptr, + val, + syncscope=prims.MemScope.CTA, + space=prims.SharedSpace.shared_cta, + ) + + +@cute.jit +def init_softmax_scratch_u32(scratch, warp_grp_thread_idx, num_entries: int): + """Initialize encoded softmax scratch entries to the negative max sentinel.""" + encoded = float_to_u32_for_atomic_max(neg_max_f32()) + for scratch_idx in cutlass.range( + warp_grp_thread_idx, Int32(num_entries), Int32(128), unroll=1 + ): + scratch[scratch_idx] = encoded diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py new file mode 100644 index 000000000000..86f1ca7cee6b --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py @@ -0,0 +1,410 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reusable low-level operations for MLA decode TS kernels.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 +from cutlass.experimental import primitives as cprims + +from .constants import ( + OCTET_LANES, + SMEM_ROW_BYTES, + SMEM_VECTOR_BYTES, + SMEM_WORD_BYTE_SHIFT, + STSM_MATRIX_LANE_SHIFT, + STSM_X4_REG_COUNT, + SOFTMAX_SCRATCH_SUM_WORD_OFFSET, + TCGEN05_16X32BX2_BF16_P_STRIDE, + TCGEN05_16X32BX2_FP8_P_STRIDE, + TCGEN05_SECOND_PANEL_ADDR_OFFSET, + WARP_LANES, + WARP_LANE_MASK, + WARP_LANE_SHIFT, + WARP_REDUCTION_BFLY_DISTANCES, +) + + +primitives_inline_ptx = cprims.inline_ptx +"""Public primitive inline-PTX entry point used by MLA helper ops.""" + +inline_ptx = cute.arch.inline_ptx +"""CuTe inline PTX entry point used by MLA helper ops.""" + + +@cute.jit +def freeze_smem_descriptor(desc): + """Copy an SMEM descriptor through a register to prevent rematerialization.""" + return inline_ptx( + "mov.b64 {$w0}, {$r0};", + write_only_types=[Int64], + read_only_args=[desc], + ) + + +@cute.jit +def warp_idx_from_warpgroup_thread(warp_grp_thread_idx): + """Return the warp index inside a warpgroup thread-id range.""" + return warp_grp_thread_idx >> Int32(WARP_LANE_SHIFT) + + +@cute.jit +def lane_idx_from_thread(thread_idx): + """Return the lane index inside one CUDA warp.""" + return thread_idx & Int32(WARP_LANE_MASK) + + +@cute.jit +def tcgen05_second_panel_addr(base_addr): + """Return the second TMEM panel address for split S/P/O panels.""" + # tcgen05 TMEM addresses encode the row in the high 16 bits. Split-panel + # layouts keep the column fixed and advance only that row field. + return base_addr + Int32(TCGEN05_SECOND_PANEL_ADDR_OFFSET) + + +@cute.jit +def tcgen05_panel_addr(base_addr, panel_idx): + """Return a TMEM address offset by ``panel_idx`` split panels.""" + # Panel index arithmetic uses the same high-16-bit row advance as the + # second-panel helper so call sites do not open-code TMEM address packing. + return base_addr + Int32(panel_idx * TCGEN05_SECOND_PANEL_ADDR_OFFSET) + + +@cute.jit +def softmax_sum_state_ptr(state_ptr): + """Return the softmax-sum scratch panel paired with a max scratch pointer.""" + # The online-softmax scratch allocation stores max and sum panels in one + # Uint32 array. The sum pointer is offset in scratch words, not bytes. + return state_ptr + Int32(SOFTMAX_SCRATCH_SUM_WORD_OFFSET) + + +@cutlass.dsl_user_op +def vector_from_scalars(values, dtype, *, loc=None, ip=None): + """Pack scalar register values into a DSL vector.""" + return cutlass.Vector.from_elements( + tuple(dtype(value) for value in values), + dtype, + loc=loc, + ip=ip, + ) + + +@cutlass.dsl_user_op +def fmax_f32(a, b, *, loc=None, ip=None): + """Return the maximum of two values as Float32.""" + return Float32( + cute.math.max(Float32(a), Float32(b), ftz=True, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + +@cutlass.dsl_user_op +def warp_reduce_max_f32(val, *, loc=None, ip=None): + """Reduce a Float32 value to the warp maximum with butterfly shuffles.""" + val = Float32(val) + for dist in WARP_REDUCTION_BFLY_DISTANCES: + val = fmax_f32( + val, + Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=val, + offset=dist, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + loc=loc, + ip=ip, + ) + ), + loc=loc, + ip=ip, + ) + return val + + +@cutlass.dsl_user_op +def warp_reduce_sum_f32(val, *, loc=None, ip=None): + """Reduce a Float32 value to the warp sum with butterfly shuffles.""" + val = Float32(val) + for dist in WARP_REDUCTION_BFLY_DISTANCES: + val = Float32( + val + + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=val, + offset=dist, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + loc=loc, + ip=ip, + ) + ) + return val + + +@cutlass.dsl_user_op +def tcgen05_ld_16x32bx2_f32( + tmem_addr, + *, + num: cutlass.Constexpr[int], + offset, + loc=None, + ip=None, +): + """Load Float32 registers from a split 16x32bx2 TMEM tile.""" + result = cprims.tcgen05_ld( + cprims.Tcgen05LdStShape.SHAPE_16X32BX2, + tmem_addr, + num=num, + offset=Int64(offset), + loc=loc, + ip=ip, + ) + if num == 1: + return result[0] + return result + + +@cutlass.dsl_user_op +def tcgen05_st_16x32bx2_f32( + tmem_addr, + value, + *, + offset, + loc=None, + ip=None, +): + """Store Float32 registers to a split 16x32bx2 TMEM tile.""" + cprims.tcgen05_st( + cprims.Tcgen05LdStShape.SHAPE_16X32BX2, + tmem_addr, + value, + offset=Int64(offset), + loc=loc, + ip=ip, + ) + + +@cute.jit +def pack_float4_to_fp8_e4m3(v0: Float32, v1: Float32, v2: Float32, v3: Float32): + """Pack four Float32 values into one E4M3 x4 register.""" + # Spell the canonical pair conversions through CuTe's public inline-PTX + # operation so the generated f8x2 instruction receives only its supported + # operands. + return inline_ptx( + "{\n" + " .reg .b16 lo;\n" + " .reg .b16 hi;\n" + " cvt.rn.satfinite.e4m3x2.f32 lo, {$r1}, {$r0};\n" + " cvt.rn.satfinite.e4m3x2.f32 hi, {$r3}, {$r2};\n" + " mov.b32 {$w0}, {lo, hi};\n" + "}", + write_only_types=[Int32], + read_only_args=[v0, v1, v2, v3], + ) + + +@cute.jit +def fp8_log2_quant_scale(): + """Return log2(448) for the E4M3 P scaling convention.""" + return Float32(8.8073549) + + +@cute.jit +def fp8_quant_scale_rcp(): + """Return the reciprocal of the shared E4M3 probability scale.""" + return Float32(1.0 / 448.0) + + +@cute.jit +def fp8_stsm_smem_dst( + smem_base_i32, + warp_grp_thread_idx, + num_trans_rows, + num_trans_cols, + stsm_idx: cutlass.Constexpr[int], +): + """Return byte-transposed STSM destination for FP8 P/O staging.""" + + num_rows = Int32(num_trans_rows) + num_bytes_per_row = Int32(num_trans_cols) + num_rows_per_smem_row = Int32(SMEM_ROW_BYTES) // num_bytes_per_row + num_segs_per_warp_per_row = num_bytes_per_row // Int32( + SMEM_VECTOR_BYTES * STSM_X4_REG_COUNT + ) + # One STSM group covers one row for the current 8-row MMA fragment. + num_stsm_per_row = 1 + num_mtx_per_col = num_rows // Int32(OCTET_LANES) + warp_idx = warp_idx_from_warpgroup_thread(warp_grp_thread_idx) + lane_idx = lane_idx_from_thread(warp_grp_thread_idx) + thr_row_idx = lane_idx & Int32(OCTET_LANES - 1) + mtx_idx = lane_idx >> Int32(STSM_MATRIX_LANE_SHIFT) + mtx_row_idx = mtx_idx % num_mtx_per_col + mtx_col_idx = mtx_idx // num_mtx_per_col + + # STSM writes one 8-row matrix per lane group. The destination is computed + # in bytes so the same helper works for the x1/x2/x4 instruction variants; + # the final pointer conversion switches back to Uint32 word addressing. + stsm_row_idx = Int32(stsm_idx % num_stsm_per_row) + stsm_col_idx = Int32(stsm_idx // num_stsm_per_row) + xor_mask = thr_row_idx // num_rows_per_smem_row + seg_col_idx = ( + warp_idx * num_segs_per_warp_per_row + mtx_col_idx + stsm_col_idx + ) ^ xor_mask + smem_offset = ( + mtx_row_idx * Int32(OCTET_LANES) + + thr_row_idx + + stsm_row_idx * Int32(WARP_LANES) + ) * num_bytes_per_row + seg_col_idx * Int32(SMEM_VECTOR_BYTES) + return smem_base_i32.data_ptr(smem_offset >> Int32(SMEM_WORD_BYTE_SHIFT)) + + +@cute.jit +def store_transposed_smem8b_x1( + smem_base_i32, + reg0: Int32, + warp_grp_thread_idx, + num_trans_rows, + num_trans_cols, + stsm_idx: cutlass.Constexpr[int] = 0, +): + """Store one FP8 STSM register into transposed SMEM layout.""" + smem_dst = fp8_stsm_smem_dst( + smem_base_i32, warp_grp_thread_idx, num_trans_rows, num_trans_cols, stsm_idx + ) + primitives_inline_ptx( + "stmatrix.sync.aligned.m16n8.x1.trans.shared.b8 [{$r0}], {{$r1}};", + read_only_args=[smem_dst, reg0], + ) + + +@cute.jit +def store_transposed_smem8b_x2( + smem_base_i32, + reg0: Int32, + reg1: Int32, + warp_grp_thread_idx, + num_trans_rows, + num_trans_cols, + stsm_idx: cutlass.Constexpr[int] = 0, +): + """Store two FP8 STSM registers into transposed SMEM layout.""" + smem_dst = fp8_stsm_smem_dst( + smem_base_i32, warp_grp_thread_idx, num_trans_rows, num_trans_cols, stsm_idx + ) + primitives_inline_ptx( + "stmatrix.sync.aligned.m16n8.x2.trans.shared.b8 [{$r0}], {{$r1}, {$r2}};", + read_only_args=[smem_dst, reg0, reg1], + ) + + +@cute.jit +def store_transposed_smem8b_x4( + smem_base_i32, + reg0: Int32, + reg1: Int32, + reg2: Int32, + reg3: Int32, + warp_grp_thread_idx, + num_trans_rows, + num_trans_cols, + stsm_idx: cutlass.Constexpr[int] = 0, +): + """Store four FP8 STSM registers into transposed SMEM layout.""" + smem_dst = fp8_stsm_smem_dst( + smem_base_i32, warp_grp_thread_idx, num_trans_rows, num_trans_cols, stsm_idx + ) + primitives_inline_ptx( + "stmatrix.sync.aligned.m16n8.x4.trans.shared.b8 [{$r0}], {{$r1}, {$r2}, {$r3}, {$r4}};", + read_only_args=[smem_dst, reg0, reg1, reg2, reg3], + ) + + +@cute.jit +def tcgen05_store_p_16x32bx2_x16(tmem_addr, regs_p, start_idx: cutlass.Constexpr[int]): + """Store 16 packed BF16 P registers to a split TMEM-P tile.""" + # The stride immediate is the BF16 lane-to-column spacing for the split + # 16x32bx2 P layout; keep it named because the FP8 path uses a different + # immediate with the same instruction shape. + inline_ptx( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32 " + f"[{{$r0}}], {TCGEN05_16X32BX2_BF16_P_STRIDE}, " + "{ {$r1}, {$r2}, {$r3}, {$r4}, {$r5}, {$r6}, {$r7}, {$r8}, " + "{$r9}, {$r10}, {$r11}, {$r12}, {$r13}, {$r14}, {$r15}, {$r16} };", + read_only_args=[ + tmem_addr, + regs_p[start_idx + 0], + regs_p[start_idx + 1], + regs_p[start_idx + 2], + regs_p[start_idx + 3], + regs_p[start_idx + 4], + regs_p[start_idx + 5], + regs_p[start_idx + 6], + regs_p[start_idx + 7], + regs_p[start_idx + 8], + regs_p[start_idx + 9], + regs_p[start_idx + 10], + regs_p[start_idx + 11], + regs_p[start_idx + 12], + regs_p[start_idx + 13], + regs_p[start_idx + 14], + regs_p[start_idx + 15], + ], + ) + + +@cute.jit +def tcgen05_store_p_fp8_16x32bx2_x16(tmem_addr, regs_p): + """Store 16 packed E4M3 P registers to a split TMEM-P tile.""" + # FP8 P uses twice as many elements per byte vector as BF16, so the TMEM + # column stride immediate is smaller even though the register count matches. + inline_ptx( + "tcgen05.st.sync.aligned.16x32bx2.x16.b32 " + f"[{{$r0}}], {TCGEN05_16X32BX2_FP8_P_STRIDE}, " + "{ {$r1}, {$r2}, {$r3}, {$r4}, {$r5}, {$r6}, {$r7}, {$r8}, " + "{$r9}, {$r10}, {$r11}, {$r12}, {$r13}, {$r14}, {$r15}, {$r16} };", + read_only_args=[ + tmem_addr, + regs_p[0], + regs_p[1], + regs_p[2], + regs_p[3], + regs_p[4], + regs_p[5], + regs_p[6], + regs_p[7], + regs_p[8], + regs_p[9], + regs_p[10], + regs_p[11], + regs_p[12], + regs_p[13], + regs_p[14], + regs_p[15], + ], + ) + + +@cute.jit +def float_to_u32_bits(val): + """Return the raw Uint32 bit pattern for a Float32 value.""" + return cprims.mov_b32(val, target_type=Uint32) + + +@cute.jit +def u32_bits_to_float(val: Uint32): + """Return the Float32 value represented by raw Uint32 bits.""" + return cprims.mov_b32(val, target_type=Float32) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/query.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/query.py new file mode 100644 index 000000000000..b1269f976be5 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/query.py @@ -0,0 +1,238 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Logical/physical query-row geometry helpers for MLA decode.""" + +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 + + +@dataclass(frozen=True) +class FlatQueryTileLayout: + """Host-side layout for consecutive ``(query, head)`` rows. + + Query tokens and heads form one affine row space ordered as + ``flat_row = query_idx * logical_num_heads_q + head_idx``. Physical MMA + tiles consume consecutive rows from that space, so only the final tile can + be partial. ``tail_rows`` is in ``[1, tile_size_q]`` and equals + ``tile_size_q`` when the final tile is full. + + Unlike the monolithic M128 helper, this generalized PrimTS layout permits + ``logical_num_heads_q > tile_size_q``. That case is required by the M8, + M16, M32, and M64 1CTA profiles: a physical tile can cover part of one + token and the following tile continues at the next logical head row. + """ + + logical_num_heads_q: int + logical_seq_len_q: int + tile_size_q: int + total_rows: int + num_tiles: int + tail_rows: int + + @classmethod + def for_tile( + cls, + logical_num_heads_q: int, + logical_seq_len_q: int, + tile_size_q: int, + ) -> "FlatQueryTileLayout": + if logical_num_heads_q <= 0: + raise ValueError("logical_num_heads_q must be positive") + if logical_seq_len_q <= 0: + raise ValueError("logical_seq_len_q must be positive") + if tile_size_q <= 0: + raise ValueError("tile_size_q must be positive") + + total_rows = logical_num_heads_q * logical_seq_len_q + num_tiles = (total_rows + tile_size_q - 1) // tile_size_q + tail_rows = total_rows - (num_tiles - 1) * tile_size_q + return cls( + logical_num_heads_q=logical_num_heads_q, + logical_seq_len_q=logical_seq_len_q, + tile_size_q=tile_size_q, + total_rows=total_rows, + num_tiles=num_tiles, + tail_rows=tail_rows, + ) + + +@cute.jit +def query_batch_bounds( + cu_seqlens_q, + batch_idx, + logical_seq_len_q: cutlass.Constexpr[int], +): + """Return the compact-storage offset and logical Q length for a batch. + + ``cu_seqlens_q`` follows the standard cumulative-offset convention. A + ``None`` value selects the fixed-length specialization, where every batch + has ``logical_seq_len_q`` rows and the batch dimension remains explicit in + the public Q/O tensors. + """ + + if cutlass.const_expr(cu_seqlens_q is None): + return Int32(0), Int32(logical_seq_len_q) + q_start = Int32(cu_seqlens_q[batch_idx]) + q_end = Int32(cu_seqlens_q[Int32(batch_idx) + Int32(1)]) + return q_start, q_end - q_start + + +@cute.jit +def runtime_flat_query_tile_valid_rows( + query_tile_idx, + tile_size_q: cutlass.Constexpr[int], + logical_num_heads_q: cutlass.Constexpr[int], + logical_seq_len_q: cutlass.Constexpr[int], + cu_seqlens_q=None, + batch_idx=None, +): + """Return the active row prefix of one physical flat-Q tile.""" + + _, query_len = query_batch_bounds( + cu_seqlens_q, + batch_idx, + logical_seq_len_q, + ) + remaining_rows = query_len * Int32(logical_num_heads_q) - Int32( + query_tile_idx + ) * Int32(tile_size_q) + return cute.math.max( + Int32(0), + cute.math.min(Int32(tile_size_q), remaining_rows), + ) + + +@cute.jit +def runtime_flat_query_tile_has_rows( + query_tile_idx, + tile_size_q: cutlass.Constexpr[int], + logical_num_heads_q: cutlass.Constexpr[int], + logical_seq_len_q: cutlass.Constexpr[int], + cu_seqlens_q=None, + batch_idx=None, +): + """Return whether one rectangular scheduler tile owns a logical Q row.""" + + return runtime_flat_query_tile_valid_rows( + query_tile_idx, + tile_size_q, + logical_num_heads_q, + logical_seq_len_q, + cu_seqlens_q, + batch_idx, + ) > Int32(0) + + +@cute.jit +def flat_query_row_state( + row_in_tile, + query_tile_idx, + tile_size_q: cutlass.Constexpr[int], + logical_num_heads_q: cutlass.Constexpr[int], + logical_seq_len_q: cutlass.Constexpr[int], + cu_seqlens_q=None, + batch_idx=None, +): + """Map one physical flat-tile row to logical and public coordinates. + + Returns ``(storage_flat_row, logical_head, safe_local_q, storage_q, + is_valid)``. Invalid final-tile rows receive safe control-flow coordinates + but remain predicated from every GMEM transaction by ``is_valid``. + """ + + local_flat_query_row = Int32(query_tile_idx) * Int32(tile_size_q) + Int32( + row_in_tile + ) + logical_q_idx = local_flat_query_row // Int32(logical_num_heads_q) + logical_head_idx = local_flat_query_row - logical_q_idx * Int32(logical_num_heads_q) + q_start, q_len = query_batch_bounds( + cu_seqlens_q, + batch_idx, + logical_seq_len_q, + ) + safe_q_len = cute.math.max(q_len, Int32(1)) + safe_logical_q_idx = cute.math.min(logical_q_idx, safe_q_len - Int32(1)) + storage_q_idx = q_start + safe_logical_q_idx + storage_flat_query_row = ( + storage_q_idx * Int32(logical_num_heads_q) + logical_head_idx + ) + is_valid = local_flat_query_row < q_len * Int32(logical_num_heads_q) + return ( + storage_flat_query_row, + logical_head_idx, + safe_logical_q_idx, + storage_q_idx, + is_valid, + ) + + +@cute.jit +def public_query_flat_row(cfg, storage_flat_query_row, batch_idx, cu_seqlens_q): + """Return the flat physical row used by public O and LSE tensors. + + Fixed-length tensors retain an explicit batch dimension, while compact + variable-length tensors already include the cumulative batch offset in + ``storage_flat_query_row``. + """ + + if cutlass.const_expr(cu_seqlens_q is None): + return Int32(batch_idx) * Int32( + cfg.logical_seq_len_q * cfg.logical_num_heads_q + ) + Int32(storage_flat_query_row) + return Int32(storage_flat_query_row) + + +@cute.jit +def split_o_element_offset( + cfg, + batch_idx, + q_idx, + head_idx, + split_idx, + dim_idx, +): + """Return a batch-dynamic partial-O offset with 64-bit-safe products. + + Batch size is intentionally absent from the compile signature, so the + complete workspace extent cannot prove that every offset fits in Int32. + The within-batch layout is still compile-time constant, however. Keep that + common bounded part in 32-bit arithmetic and widen only the batch-stride + product. Use fully widened arithmetic when a profile's per-batch layout + alone exceeds Int32. + """ + + elements_per_batch = ( + cfg.seq_len_q * cfg.num_heads_q * cfg.num_ctas_per_seq_kv * cfg.head_dim_v + ) + if cutlass.const_expr(elements_per_batch <= (1 << 31) - 1): + within_batch_offset = ( + (Int32(q_idx) * Int32(cfg.num_heads_q) + Int32(head_idx)) + * Int32(cfg.num_ctas_per_seq_kv) + + Int32(split_idx) + ) * Int32(cfg.head_dim_v) + Int32(dim_idx) + return Int64(batch_idx) * Int64(elements_per_batch) + Int64(within_batch_offset) + + return ( + ( + (Int64(batch_idx) * Int64(cfg.seq_len_q) + Int64(q_idx)) + * Int64(cfg.num_heads_q) + + Int64(head_idx) + ) + * Int64(cfg.num_ctas_per_seq_kv) + + Int64(split_idx) + ) * Int64(cfg.head_dim_v) + Int64(dim_idx) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py new file mode 100644 index 000000000000..133c3e15fe48 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py @@ -0,0 +1,336 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Captured-schedule helper functions for MLA decode TS examples.""" + +from contextlib import contextmanager + +from cutlass.experimental.task_scheduling.schedule_builder import work_tile_loop + +from .stage import MlaStage + + +def captured_loop_bounds(task_kwargs, default_start: int, default_step: int = 1): + """Consume Task loop kwargs and return captured schedule loop bounds.""" + loop_start = task_kwargs.pop("domain_start", default_start) + loop_step = task_kwargs.pop("step", default_step) + loop_end = task_kwargs.pop("domain", None) + if loop_end is None: + loop_end = 1 + return loop_start, loop_end, loop_step + + +def work_queue_tail(work_queue, *, advance_label: str = "get_and_advance_work_tile"): + """Advance the persistent work tile at the end of a task body.""" + if work_queue is None: + return + work_queue.wait() + getattr(work_queue, advance_label)() + work_queue.release() + + +def runtime_work_tile_skip_if(work_queue): + """Return a concrete work queue's public runtime-skip predicate.""" + if work_queue is None or not getattr(work_queue, "enable_runtime_skip", False): + return None + return getattr(type(work_queue), "skip_work_tile_if", None) + + +@contextmanager +def work_tile_schedule_loop( + work_queue, + *, + skip_if=None, + non_skippable_prelude=None, +): + """Wrap a captured task body in the standard TS persistent work-tile loop. + + A decode work queue may expose ``skip_work_tile_if`` for runtime-padded + tiles. Only task data work is skippable; queue wait/advance/release remains + outside the guard so every persistent task advances in lockstep. Pure + register initializers may run in ``non_skippable_prelude`` so their values + dominate the stock executor's separately guarded HEAD, LOOP, and TAIL + regions. The prelude must not issue memory or pipeline operations. + """ + if work_queue is not None: + with work_tile_loop(work_queue, skip_if=skip_if) as work_tile: + prelude_state = ( + non_skippable_prelude() if non_skippable_prelude is not None else None + ) + if skip_if is None: + yield work_tile, prelude_state + else: + with work_tile.skippable(): + yield work_tile, prelude_state + work_queue_tail(work_queue) + else: + yield None, None + + +def page_offsets_consume(smem_page_offsets, cached_page_ids=None): + """Consume staged page offsets for the next K/V TMA transfer.""" + if smem_page_offsets is None: + return None + smem_page_offsets.wait() + return smem_page_offsets.read_offsets(cached_page_ids=cached_page_ids) + + +def page_offsets_release(smem_page_offsets): + """Release the page-offset stage after the K/V load has consumed it.""" + if smem_page_offsets is None: + return + smem_page_offsets.release() + + +def page_offsets_produce(smem_page_offsets, label, *, section: MlaStage): + """Produce one page-offset stage using the named K/V slot callback.""" + smem_page_offsets.acquire() + getattr(smem_page_offsets, label)(section=section) + smem_page_offsets.commit() + + +def schedule_token_throttle_head(schedule_token_throttle): + """Publish a load-task schedule token before issuing staged loads.""" + if schedule_token_throttle is None: + return + schedule_token_throttle.acquire() + schedule_token_throttle.publish_schedule_token() + schedule_token_throttle.commit() + + +def schedule_token_throttle_tail(schedule_token_throttle): + """Consume a schedule token before the scheduler fetches more work.""" + if schedule_token_throttle is None: + return + schedule_token_throttle.wait() + schedule_token_throttle.consume_schedule_token() + schedule_token_throttle.release() + + +def staged_kv_tma_load( + smem_kv, + iterations: int, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + *, + is_v: bool, + use_next_v_pages: bool = False, +): + """Produce staged K or delayed-V TMA work for one logical k-tile.""" + for subtile_idx in range(iterations): + smem_kv.acquire() + smem_kv.tma_load( + cached_k_pages=cached_k_pages, + cached_v_pages=cached_v_pages, + cached_next_v_pages=cached_next_v_pages, + is_v=is_v, + subtile_idx=subtile_idx, + use_next_v_pages=use_next_v_pages, + ) + smem_kv.commit() + + +def staged_kv_load( + smem_kv, + *, + head_dim_stages, + producer_label, + section: MlaStage, + smem_page_offsets=None, + cached_page_ids=None, +): + """Issue all K/V producer stages for one logical K or V tile.""" + cached_page_ids = page_offsets_consume(smem_page_offsets, cached_page_ids) + for stage_idx in range(head_dim_stages): + smem_kv.acquire() + getattr(smem_kv, producer_label)( + stage_idx=stage_idx, + section=section, + cached_page_ids=cached_page_ids, + ) + smem_kv.commit() + page_offsets_release(smem_page_offsets) + return cached_page_ids + + +def staged_qk_mma( + smem_kv, + tmem_s, + iterations=None, + *, + head_dim_stages=None, + consumer_label=None, + include_acquire=True, +): + """Consume staged K descriptors and produce one QK MMA score tile.""" + if consumer_label is None: + for k_subtile_idx in range(iterations): + smem_kv.wait() + desc_k_base = smem_kv.k_desc(k_subtile_idx=k_subtile_idx) + if k_subtile_idx == 0: + tmem_s.acquire() + tmem_s.qk_mma(desc_k_base=desc_k_base, k_subtile_idx=k_subtile_idx) + smem_kv.release() + tmem_s.commit() + return + + if include_acquire: + tmem_s.acquire() + for k_subtile_idx in range(head_dim_stages): + smem_kv.wait() + kv_desc = getattr(smem_kv, consumer_label)(k_subtile_idx=k_subtile_idx) + tmem_s.qk_mma(kv_desc=kv_desc, k_subtile_idx=k_subtile_idx) + smem_kv.release() + tmem_s.commit() + + +def staged_pv_mma( + smem_kv, + smem_p, + tmem_o, + iterations=None, + *, + head_dim_stages=None, + consumer_label=None, + producer_label=None, + is_tail: bool = False, +): + """Consume staged V descriptors and produce one PV MMA output tile.""" + smem_p.wait() + if producer_label is None: + desc_p_base = smem_p.p_desc() + tmem_o.acquire() + for v_subtile_idx in range(iterations): + smem_kv.wait() + desc_v_base = smem_kv.v_desc(v_subtile_idx=v_subtile_idx) + tmem_o.pv_mma( + desc_p_base=desc_p_base, + desc_v_base=desc_v_base, + v_subtile_idx=v_subtile_idx, + is_tail=is_tail, + ) + smem_kv.release() + tmem_o.commit() + smem_p.release() + return + + smem_p.p_desc() + tmem_o.acquire() + is_tail = is_tail or ( + producer_label is not None and producer_label.find("tail") >= 0 + ) + for v_subtile_idx in range(head_dim_stages): + smem_kv.wait() + if producer_label.endswith("_0"): + v_desc_0 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) + getattr(tmem_o, producer_label)( + v_desc_0=v_desc_0, + v_subtile_idx=v_subtile_idx, + is_tail=is_tail, + ) + else: + v_desc_1 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) + getattr(tmem_o, producer_label)( + v_desc_1=v_desc_1, + v_subtile_idx=v_subtile_idx, + is_tail=is_tail, + ) + smem_kv.release() + tmem_o.commit() + smem_p.release() + + +def staged_qk_mma_k_tile(smem_k, tmem_s, iterations: int): + """Consume one whole K stage and produce one QK score tile.""" + smem_k.wait() + tmem_s.acquire() + for k_subtile_idx in range(iterations): + desc_k_base = smem_k.k_desc(k_subtile_idx=k_subtile_idx) + tmem_s.qk_mma(desc_k_base=desc_k_base, k_subtile_idx=k_subtile_idx) + tmem_s.commit() + smem_k.release() + + +def staged_pv_mma_v_tile(smem_v, smem_p, tmem_o, iterations: int): + """Consume one whole V stage plus one P stage and produce one O tile.""" + smem_p.wait() + desc_p_base = smem_p.p_desc() + smem_v.wait() + tmem_o.acquire() + for v_subtile_idx in range(iterations): + desc_v_base = smem_v.v_desc(v_subtile_idx=v_subtile_idx) + tmem_o.pv_mma( + desc_p_base=desc_p_base, + desc_v_base=desc_v_base, + v_subtile_idx=v_subtile_idx, + ) + tmem_o.commit() + smem_v.release() + smem_p.release() + + +def staged_pv_mma_v_tile_per_n( + smem_v, + smem_p, + tmem_o, + *, + iterations_pv_k: int, + iterations_pv_n: int, +): + """Consume P/V and publish one O pipeline token per PV N-slice.""" + smem_p.wait() + desc_p_base = smem_p.p_desc() + smem_v.wait() + for pv_n_idx in range(iterations_pv_n): + tmem_o.acquire() + for pv_k_idx in range(iterations_pv_k): + desc_v_base = smem_v.v_desc_n_major(pv_n_idx=pv_n_idx, pv_k_idx=pv_k_idx) + tmem_o.pv_mma_n_major( + desc_p_base=desc_p_base, + desc_v_base=desc_v_base, + pv_n_idx=pv_n_idx, + pv_k_idx=pv_k_idx, + ) + tmem_o.commit() + smem_v.release() + smem_p.release() + + +def staged_pv_mma_tmem_p( + smem_kv, + tmem_p, + tmem_o, + *, + head_dim_stages, + consumer_label, + producer_label="pv_mma_loop_tmem_p", +): + """Consume TMEM P and staged V descriptors to produce one PV output tile.""" + tmem_p.wait() + p_stage_idx = tmem_p.p_stage() + tmem_o.acquire() + is_tail = producer_label == "pv_mma_tail_tmem_p" + for v_subtile_idx in range(head_dim_stages): + smem_kv.wait() + v_desc_0 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) + getattr(tmem_o, producer_label)( + p_stage_idx=p_stage_idx, + v_desc_0=v_desc_0, + v_subtile_idx=v_subtile_idx, + is_tail=is_tail, + ) + smem_kv.release() + tmem_o.commit() + tmem_p.release() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/stage.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/stage.py new file mode 100644 index 000000000000..eb2a3d80f879 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/stage.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Kernel-local schedule-section tag for MLA decode work bodies. + +MLA decode work methods sometimes need to know which schedule section +(head/loop/tail) a call belongs to. Rather than depend on the task-scheduling +framework's ``ScheduleStageType``, the MLA schedules pass this small kernel-local +enum explicitly as a compile-time constant, so bodies branch on it with +``cutlass.const_expr(stage == MlaStage.Head)``. +""" + +import enum + + +class MlaStage(enum.Enum): + """Schedule section of a single MLA decode work call.""" + + Head = 0 + Loop = 1 + Tail = 2 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py new file mode 100644 index 000000000000..a5e7e810d770 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py @@ -0,0 +1,447 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tile and stage index helper functions for MLA decode.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + +from cutlass.experimental.task_scheduling.resources import StageInfo + +from .layout import _TASK_CACHE_SEQ_LEN_KV +from .mask import MaskType, mask_visible_k_length +from .query import ( + flat_query_row_state, + query_batch_bounds, + runtime_flat_query_tile_has_rows, +) +from .stage import MlaStage +from ..throughput_latency_1cta.config import MlaConfig + + +@cute.jit +def local_kv_tile_idx( + cfg: MlaConfig, + stage_info: StageInfo, + inst_id: int, + is_v: int, + *, + section: cutlass.Constexpr[MlaStage], +): + """Return the local K/V tile id for one staged pipeline instance.""" + if cutlass.const_expr(section == MlaStage.Head): + return Int32(inst_id) + if cutlass.const_expr(section == MlaStage.Loop): + base = stage_info.loop_offset * Int32(cfg.num_insts_kv) + if cutlass.const_expr(is_v): + return base + Int32(inst_id) + return base + Int32(cfg.num_insts_kv + inst_id) + if cutlass.const_expr(cfg.total_kv_tiles <= cfg.num_insts_kv): + return Int32(inst_id) + return stage_info.loop_end * Int32(cfg.num_insts_kv) + Int32(inst_id) + + +@cute.jit +def softmax_kv_tile_idx(cfg: MlaConfig, stage_info: StageInfo, inst_id: int): + """Return the K tile consumed by one softmax instance in the loop body.""" + return stage_info.loop_offset * Int32(cfg.num_insts_kv) + Int32(inst_id) + + +@cute.jit +def runtime_total_kv_tiles(cfg: MlaConfig, seq_len_kv): + """Return ceil(seq_len_kv / tile_size_kv) at runtime.""" + return (seq_len_kv + Int32(cfg.tile_size_kv - 1)) // Int32(cfg.tile_size_kv) + + +def active_split_kv_count( + seq_len_kv: int, + tile_size_kv: int, + num_insts_kv: int, + configured_splits_kv: int, +) -> int: + """Return the host form of the runtime active split-prefix rule.""" + + if seq_len_kv < 0: + raise ValueError("seq_len_kv must be non-negative") + if tile_size_kv <= 0: + raise ValueError("tile_size_kv must be positive") + if num_insts_kv <= 0: + raise ValueError("num_insts_kv must be positive") + if configured_splits_kv <= 0: + raise ValueError("configured_splits_kv must be positive") + total_kv_tiles = (seq_len_kv + tile_size_kv - 1) // tile_size_kv + groups_per_cta = (total_kv_tiles + configured_splits_kv * num_insts_kv - 1) // ( + configured_splits_kv * num_insts_kv + ) + local_kv_tiles = max(groups_per_cta * num_insts_kv, num_insts_kv) + return (total_kv_tiles + local_kv_tiles - 1) // local_kv_tiles + + +def runtime_split_pruning_is_profitable(configured_splits_kv: int) -> bool: + """Return whether split pruning can retire enough 1CTA work to pay for itself. + + A configured S2/S3 launch can retire at most one or two mainloop CTAs per + logical tile, while still requiring every cluster rank to publish a neutral + partial and perform its static row-owner reduction. Hardware validation showed + that this does not amortize the runtime activity branch. Starting at S4, + contracted requests retire enough K/V task graphs to recover the control + cost. This is a topology rule, not a problem-shape list or user knob. + """ + + if configured_splits_kv <= 0: + raise ValueError("configured_splits_kv must be positive") + return configured_splits_kv >= 4 + + +@cute.jit +def _runtime_configured_local_kv_tiles(cfg: MlaConfig, seq_len_kv): + """Return the instruction-aligned local span for configured split capacity.""" + + total_kv_tiles = runtime_total_kv_tiles(cfg, seq_len_kv) + num_insts_kv = Int32(cfg.num_insts_kv) + tiles_per_group = Int32(cfg.num_ctas_per_seq_kv) * num_insts_kv + num_groups = (total_kv_tiles + tiles_per_group - Int32(1)) // tiles_per_group + return cute.math.max(num_groups * num_insts_kv, num_insts_kv) + + +@cute.jit +def runtime_num_ctas_kv(cfg: MlaConfig, seq_len_kv): + """Return the active prefix of configured split-KV CTAs. + + One CTA's minimum useful K unit is ``tile_size_kv * num_insts_kv``. + Launch and workspace shapes retain the configured maximum split count; + runtime-short sequences activate only the prefix that owns real work. + """ + if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): + return Int32(1) + total_kv_tiles = runtime_total_kv_tiles(cfg, seq_len_kv) + local_kv_tiles = _runtime_configured_local_kv_tiles(cfg, seq_len_kv) + return (total_kv_tiles + local_kv_tiles - Int32(1)) // local_kv_tiles + + +@cute.jit +def runtime_local_kv_tiles(cfg: MlaConfig, seq_len_kv): + """Return the padded local KV tile count assigned to each KV CTA group.""" + if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): + return runtime_total_kv_tiles(cfg, seq_len_kv) + return _runtime_configured_local_kv_tiles(cfg, seq_len_kv) + + +@cute.jit +def runtime_base_seq_len_kv(cfg: MlaConfig, cache_seqs, batch_idx): + """Return the raw KV sequence length for a batch row.""" + if cutlass.const_expr(cache_seqs is None): + return Int32(cfg.seq_len_kv) + return Int32(cache_seqs[batch_idx]) + + +@cute.jit +def runtime_seq_len_kv_for_logical_q( + cfg: MlaConfig, + cache_seqs, + batch_idx, + logical_q_idx, + cu_seqlens_q=None, +): + """Return the mask-visible KV length for one logical Q row. + + Dense decode exposes the full runtime cache. Bottom-right causal decode + removes the speculative K positions following ``logical_q_idx``. + """ + _, logical_seq_len_q = query_batch_bounds( + cu_seqlens_q, + batch_idx, + cfg.logical_seq_len_q, + ) + return mask_visible_k_length( + cfg.mask_type, + runtime_base_seq_len_kv(cfg, cache_seqs, batch_idx), + logical_q_idx, + logical_seq_len_q, + ) + + +@cute.jit +def runtime_seq_len_kv_for_q( + cfg: MlaConfig, + cache_seqs, + batch_idx, + cta_idx_q, + cu_seqlens_q=None, +): + """Return the KV domain shared by flat query rows in a CTA. + + The last physical row has the largest causal domain. Row-causal softmax + applies the remaining row-specific mask; all other modes can use + this CTA-visible length with the ordinary dense tail predicate. + """ + if cutlass.const_expr(cfg.mask_type == MaskType.DENSE.value): + return runtime_base_seq_len_kv(cfg, cache_seqs, batch_idx) + _, _, logical_q_idx, _, _ = flat_query_row_state( + Int32(cfg.tile_size_q - 1), + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q, + batch_idx, + ) + return runtime_seq_len_kv_for_logical_q( + cfg, + cache_seqs, + batch_idx, + logical_q_idx, + cu_seqlens_q, + ) + + +@cute.jit +def runtime_query_tile_is_active( + cfg: MlaConfig, + cu_seqlens_q, + batch_idx, + cta_idx_q, +): + """Return whether a configured Q tile owns any runtime query rows.""" + + query_is_active = cutlass.Boolean(True) + if cutlass.const_expr(cu_seqlens_q is not None): + query_is_active = runtime_flat_query_tile_has_rows( + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q, + batch_idx, + ) + return query_is_active + + +@cute.jit +def runtime_split_tile_is_active( + cfg: MlaConfig, + cache_seqs, + cu_seqlens_q, + batch_idx, + cta_idx_q, + cta_idx_kv, +): + """Return whether a configured split rank owns runtime KV work.""" + + seq_len_kv = runtime_seq_len_kv_for_q( + cfg, + cache_seqs, + batch_idx, + cta_idx_q, + cu_seqlens_q, + ) + return Int32(cta_idx_kv) < runtime_num_ctas_kv(cfg, seq_len_kv) + + +@cute.jit +def runtime_work_tile_activity( + cfg: MlaConfig, + cache_seqs, + cu_seqlens_q, + batch_idx, + cta_idx_q, + cta_idx_kv, +): + """Return the independent runtime Q and split-KV activity predicates.""" + + query_is_active = runtime_query_tile_is_active( + cfg, + cu_seqlens_q, + batch_idx, + cta_idx_q, + ) + split_is_active = runtime_split_tile_is_active( + cfg, + cache_seqs, + cu_seqlens_q, + batch_idx, + cta_idx_q, + cta_idx_kv, + ) + return query_is_active, split_is_active + + +@cute.jit +def runtime_work_tile_is_active( + cfg: MlaConfig, + cache_seqs, + cu_seqlens_q, + batch_idx, + cta_idx_q, + cta_idx_kv, +): + """Return whether a Q/split work tile owns any runtime data.""" + + query_is_active, split_is_active = runtime_work_tile_activity( + cfg, + cache_seqs, + cu_seqlens_q, + batch_idx, + cta_idx_q, + cta_idx_kv, + ) + return query_is_active and split_is_active + + +@cute.jit +def runtime_seq_len_kv_for_query_row( + cfg: MlaConfig, + cache_seqs, + batch_idx, + cta_idx_q, + row_in_tile, + cu_seqlens_q=None, +): + """Return the KV length visible to one physical flat-query row.""" + _, _, logical_q_idx, _, _ = flat_query_row_state( + row_in_tile, + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q, + batch_idx, + ) + return runtime_seq_len_kv_for_logical_q( + cfg, + cache_seqs, + batch_idx, + logical_q_idx, + cu_seqlens_q, + ) + + +@cute.jit +def runtime_seq_len_kv_from_task_cache( + cfg: MlaConfig, + task_cache, + cta_idx_q, + cu_seqlens_q=None, + batch_idx=None, +): + """Return the CTA-visible KV domain from the task-cached raw length.""" + seq_len_kv = Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]) + if cutlass.const_expr(cfg.mask_type == MaskType.DENSE.value): + return seq_len_kv + _, logical_seq_len_q = query_batch_bounds( + cu_seqlens_q, + batch_idx, + cfg.logical_seq_len_q, + ) + _, _, logical_q_idx, _, _ = flat_query_row_state( + Int32(cfg.tile_size_q - 1), + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q, + batch_idx, + ) + return mask_visible_k_length( + cfg.mask_type, + seq_len_kv, + logical_q_idx, + logical_seq_len_q, + ) + + +@cute.jit +def global_kv_tile_idx( + cfg: MlaConfig, + local_tile_idx, + seq_len_kv, + cta_idx_kv, +): + """Map a local KV tile id to the global KV tile id for split-KV mode.""" + if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): + return local_tile_idx + return Int32(cta_idx_kv) * runtime_local_kv_tiles(cfg, seq_len_kv) + local_tile_idx + + +@cute.jit +def attr_or_work_tile_idx(attr, stage_info: StageInfo, coord_idx: int): + """Return an explicit attribute value or the matching work-tile coordinate.""" + if cutlass.const_expr(attr is None): + return Int32(stage_info.work_tile.tile_idx[coord_idx]) + return Int32(attr) + + +@cute.jit +def batch_idx_for_stage(attr, stage_info: StageInfo): + """Return the batch index for resources that do not pack heads into z.""" + return attr_or_work_tile_idx(attr, stage_info, 2) + + +@cute.jit +def batch_idx_for_stage_cfg(attr, cfg: MlaConfig, stage_info: StageInfo): + """Return the batch index from a batch-major combined batch/head coordinate.""" + if cutlass.const_expr(attr is None): + _cta_idx_q, _cta_idx_head_dim, batch_head_idx = stage_info.work_tile.tile_idx + del _cta_idx_q, _cta_idx_head_dim + batch_head_idx = Int32(batch_head_idx) + return batch_head_idx // Int32(cfg.num_ctas_for_all_heads) + return Int32(attr) + + +@cute.jit +def head_idx_for_stage(attr, cfg: MlaConfig, stage_info: StageInfo): + """Return the base Q/head row from a batch-major combined batch/head tile.""" + if cutlass.const_expr(attr is None): + _cta_idx_q, _cta_idx_head_dim, batch_head_idx = stage_info.work_tile.tile_idx + del _cta_idx_q, _cta_idx_head_dim + batch_head_idx = Int32(batch_head_idx) + head_tile_idx = batch_head_idx % Int32(cfg.num_ctas_for_all_heads) + return head_tile_idx * Int32(cfg.tile_size_q) + return Int32(attr) + + +@cute.jit +def cta_idx_q_for_stage(attr, stage_info: StageInfo): + """Return the Q CTA index for the current stage.""" + return attr_or_work_tile_idx(attr, stage_info, 0) + + +@cute.jit +def cta_idx_head_dim_v_for_stage(attr, stage_info: StageInfo): + """Return the V head-dim CTA index for the current stage.""" + return attr_or_work_tile_idx(attr, stage_info, 1) + + +@cute.jit +def cta_idx_kv_for_stage(attr, stage_info: StageInfo): + """Return the KV CTA index, defaulting to zero for non-split KV.""" + if cutlass.const_expr(attr is None): + return Int32(0) + return Int32(attr) + + +@cute.jit +def staged_kv_head_dim_call_idx( + cfg: MlaConfig, + stage_info: StageInfo, + inst_id: int, + is_v: int, + *, + stage_idx: cutlass.Constexpr[int], + section: cutlass.Constexpr[MlaStage], +): + """Return the head-dim stage index within a K or V staged load group.""" + del cfg, stage_info, inst_id, is_v, section + return stage_idx diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile_scheduler.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile_scheduler.py new file mode 100644 index 000000000000..b9965c1620b2 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile_scheduler.py @@ -0,0 +1,317 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MLA static tile scheduler helpers used by TS decode kernels. + +The high-throughput MLA TS kernel schedules work in +``(cluster_idx, seq_q_idx, batch_idx, split_kv_idx)`` coordinates. CUTLASS's +generic static scheduler does not expose that decomposition, so the MLA example +keeps this small scheduler locally instead of importing the bare-metal MLA +runner. +""" + +import cutlass +import cutlass.cute as cute + + +def _is_positive_power_of_two(value: int) -> bool: + """Return whether a host-known scheduler extent is a power of two.""" + + return ( + isinstance(value, int) + and not isinstance(value, bool) + and value > 0 + and (value & (value - 1)) == 0 + ) + + +@cute.jit +def divmod_constexpr_power_of_two_or_fdd( + dividend, + constexpr_divisor: cutlass.Constexpr[int], + fallback_divisor: cute.FastDivmodDivisor, +): + """Use shift/mask for a power-of-two constexpr, otherwise the existing FDD.""" + + if cutlass.const_expr(_is_positive_power_of_two(constexpr_divisor)): + shift = constexpr_divisor.bit_length() - 1 + return ( + dividend >> cutlass.Int32(shift), + dividend & cutlass.Int32(constexpr_divisor - 1), + ) + return divmod(dividend, fallback_divisor) + + +class MLAStaticTileSchedulerParams: + """Static scheduler parameters for MLA split-KV work tiles.""" + + def __init__( + self, + is_persistent: bool, + problem_shape_b: cute.Int32, + problem_shape_s: cute.Int32, + cluster_shape_mnk: cute.Shape, + split_kv: cutlass.Int32, + *, + problem_shape_b_fdd: cute.FastDivmodDivisor = None, + problem_shape_s_fdd: cute.FastDivmodDivisor = None, + split_kv_fdd: cute.FastDivmodDivisor = None, + loc=None, + ip=None, + ): + """Initialize static scheduler dimensions and fast-divmod divisors.""" + self.is_persistent = is_persistent + self.problem_shape_b = problem_shape_b + self.problem_shape_s = problem_shape_s + self.problem_shape_b_fdd = problem_shape_b_fdd + self.problem_shape_s_fdd = problem_shape_s_fdd + self.cluster_shape_mnk = cluster_shape_mnk + self.split_kv = split_kv + self.split_kv_fdd = split_kv_fdd + if cutlass.const_expr(problem_shape_b_fdd is None): + self.problem_shape_b_fdd = cute.fast_divmod_create_divisor( + problem_shape_b, loc=loc, ip=ip + ) + if cutlass.const_expr(problem_shape_s_fdd is None): + self.problem_shape_s_fdd = cute.fast_divmod_create_divisor( + problem_shape_s, loc=loc, ip=ip + ) + if cutlass.const_expr(split_kv_fdd is None): + self.split_kv_fdd = cute.fast_divmod_create_divisor( + split_kv, loc=loc, ip=ip + ) + self.loc = loc + self.ip = ip + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.problem_shape_b) + values += cutlass.extract_mlir_values(self.problem_shape_s) + values += cutlass.extract_mlir_values(self.split_kv) + values += cutlass.extract_mlir_values(self.problem_shape_b_fdd) + values += cutlass.extract_mlir_values(self.problem_shape_s_fdd) + values += cutlass.extract_mlir_values(self.split_kv_fdd) + return values + + def __new_from_mlir_values__(self, values): + problem_shape_b = cutlass.new_from_mlir_values( + self.problem_shape_b, (values[0],) + ) + problem_shape_s = cutlass.new_from_mlir_values( + self.problem_shape_s, (values[1],) + ) + split_kv = cutlass.new_from_mlir_values(self.split_kv, (values[2],)) + problem_shape_b_fdd = cutlass.new_from_mlir_values( + self.problem_shape_b_fdd, (values[3],) + ) + problem_shape_s_fdd = cutlass.new_from_mlir_values( + self.problem_shape_s_fdd, (values[4],) + ) + split_kv_fdd = cutlass.new_from_mlir_values(self.split_kv_fdd, (values[5],)) + return MLAStaticTileSchedulerParams( + self.is_persistent, + problem_shape_b, + problem_shape_s, + self.cluster_shape_mnk, + split_kv, + problem_shape_b_fdd=problem_shape_b_fdd, + problem_shape_s_fdd=problem_shape_s_fdd, + split_kv_fdd=split_kv_fdd, + loc=self.loc, + ip=self.ip, + ) + + +def create_mla_static_tile_scheduler_params( + is_persistent: bool, + problem_shape_b: cute.Int32, + problem_shape_s: cute.Int32, + cluster_shape_mnk: cute.Shape, + split_kv: cutlass.Int32, +) -> MLAStaticTileSchedulerParams: + """Create MLA static scheduler parameters from host/runtime dimensions.""" + return MLAStaticTileSchedulerParams( + is_persistent, problem_shape_b, problem_shape_s, cluster_shape_mnk, split_kv + ) + + +class WorkTileInfo: + """One MLA scheduler work tile and validity bit.""" + + def __init__(self, blk_coord: cute.Coord, is_valid: bool): + self.blk_coord = blk_coord + self.is_valid = cutlass.Boolean(is_valid) + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.blk_coord) + values += cutlass.extract_mlir_values(self.is_valid) + return values + + def __new_from_mlir_values__(self, values): + new_tile_idx = cutlass.new_from_mlir_values(self.blk_coord, values[:-1]) + new_is_valid_tile = cutlass.new_from_mlir_values(self.is_valid, [values[-1]]) + return WorkTileInfo(new_tile_idx, new_is_valid_tile) + + @property + def is_valid_tile(self) -> cutlass.Boolean: + """Return whether the current tile should execute.""" + return self.is_valid + + @property + def tile_idx(self) -> cute.Coord: + """Return ``(cluster_idx, seq_q_idx, batch_idx, split_kv_idx)``.""" + return self.blk_coord + + +class MLAStaticTileScheduler: + """Static and persistent MLA tile scheduler. + + Persistent mode grid-strides through the logical MLA tile space by SM count. + Non-persistent mode maps each CTA directly to one logical work tile. + """ + + def __init__( + self, + params: MLAStaticTileSchedulerParams, + current_work_linear_idx: cutlass.Int32, + blk_coord: cute.Coord, + grid_shape: cute.Shape, + *, + is_valid: bool = True, + loc=None, + ip=None, + ): + """Initialize scheduler state for the current CTA.""" + self.params = params + self.blk_coord = blk_coord + self.grid_shape = grid_shape + self.current_work_linear_idx = current_work_linear_idx + if params.is_persistent: + self.persistent_blk_layout = cute.make_layout( + ( + params.cluster_shape_mnk[0], + params.problem_shape_s, + params.problem_shape_b, + params.split_kv, + ), + loc=loc, + ip=ip, + ) + self.num_blocks = cute.size(self.persistent_blk_layout, loc=loc, ip=ip) + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + else: + self.is_valid = cutlass.Boolean(is_valid) + self.loc = loc + self.ip = ip + + @staticmethod + def get_grid_shape( + params: MLAStaticTileSchedulerParams, + max_active_clusters: int, + *, + loc=None, + ip=None, + ) -> cute.Shape: + """Return launch grid shape for static or persistent scheduling.""" + grid_shape = ( + params.cluster_shape_mnk[0], + params.problem_shape_b * params.problem_shape_s, + params.split_kv, + ) + if params.is_persistent: + return ( + cutlass.min( + max_active_clusters * cute.size(params.cluster_shape_mnk), + cute.size(grid_shape, loc=loc, ip=ip), + ), + 1, + 1, + ) + return grid_shape + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + """Decode the current CTA or persistent linear index into one work tile.""" + is_valid = ( + self.current_work_linear_idx < self.num_blocks + if self.params.is_persistent + else self.is_valid + ) + + if self.params.is_persistent: + current_work_cluster_batch, cluster_idx = ( + self.current_work_linear_idx // self.params.cluster_shape_mnk[0], + self.current_work_linear_idx % self.params.cluster_shape_mnk[0], + ) + current_work_s_batch, s_idx = divmod( + current_work_cluster_batch, self.params.problem_shape_s_fdd + ) + current_work_b_batch, b_idx = divmod( + current_work_s_batch, self.params.problem_shape_b_fdd + ) + _, split_kv_idx = divmod(current_work_b_batch, self.params.split_kv_fdd) + blk_coord = (cluster_idx, s_idx, b_idx, split_kv_idx) + else: + s_idx, b_idx = divmod(self.blk_coord[1], self.params.problem_shape_b_fdd) + blk_coord = (self.blk_coord[0], s_idx, b_idx, self.blk_coord[2]) + + return WorkTileInfo(blk_coord, is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + """Return the initial work tile for this CTA.""" + return self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + """Advance to the next persistent tile or invalidate static work.""" + del loc, ip + if self.params.is_persistent: + self.current_work_linear_idx += advance_count * self.num_persistent_sm + else: + self.is_valid = cutlass.Boolean(False) + + def __extract_mlir_values__(self): + values = cutlass.extract_mlir_values(self.params) + values.extend(cutlass.extract_mlir_values(self.current_work_linear_idx)) + values.extend(cutlass.extract_mlir_values(self.blk_coord)) + values.extend(cutlass.extract_mlir_values(self.grid_shape)) + if not self.params.is_persistent: + values.extend(cutlass.extract_mlir_values(self.is_valid)) + return values + + def __new_from_mlir_values__(self, values): + expected_values = 13 if self.params.is_persistent else 14 + assert len(values) == expected_values + new_params = cutlass.new_from_mlir_values(self.params, values[0:6]) + new_current_work_linear_idx = cutlass.new_from_mlir_values( + self.current_work_linear_idx, [values[6]] + ) + new_blk_coord = cutlass.new_from_mlir_values(self.blk_coord, values[7:10]) + new_grid_shape = cutlass.new_from_mlir_values(self.grid_shape, values[10:13]) + new_is_valid = True + if not self.params.is_persistent: + new_is_valid = cutlass.new_from_mlir_values(self.is_valid, [values[13]]) + return MLAStaticTileScheduler( + new_params, + new_current_work_linear_idx, + new_blk_coord, + new_grid_shape, + is_valid=new_is_valid, + ) + + +def create_mla_static_tile_scheduler( + params: MLAStaticTileSchedulerParams, + blk_coord: cute.Coord, + grid_shape: cute.Shape, +) -> MLAStaticTileScheduler: + """Create the MLA static tile scheduler for the current CTA.""" + return MLAStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py new file mode 100644 index 000000000000..c0e29277c516 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py @@ -0,0 +1,339 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Explicit selector for MLA TS kernel families. + +The runner asks this module to choose between the throughput 2CTA M128 schedule +and the throughput-latency 1CTA schedule. Selection is explicit: +``requested_policy`` is honored first, and there is no implicit fallback between +families. The returned ``MlaKernelDecision`` tells the caller whether the +requested policy is implemented for the shape; callers must reject or report +``implementation_ready=False`` before constructing a kernel. + +The 2CTA predicate is a pure shape/feature eligibility check. The 1CTA path is +additionally gated by profile enumeration, and the returned +``config``/``profile_name`` are populated only when a matching 1CTA profile +exists. Invalid policy names raise ``ValueError`` in +``normalize_mla_kernel_policy``; invalid explicit profile names are propagated +from ``make_throughput_latency_mla_config``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, cast + +from .helpers.constants import SUPPORTED_MLA_PAGE_SIZES +from .throughput_latency_1cta.config import ( + MlaConfig, + SUPPORTED_TILE_SIZE_Q, + enumerate_throughput_latency_mla_profiles, + make_throughput_latency_mla_config, + tile_size_q_from_profile_name, +) + + +MlaKernelPolicy = Literal["throughput_2cta", "throughput_latency_1cta"] +"""User-visible TS MLA policy names accepted by the runner.""" + +MlaKernelName = MlaKernelPolicy +"""Concrete TS MLA kernel family selected for a launch.""" + + +@dataclass(frozen=True) +class MlaKernelDecision: + """Python-side result of explicit MLA TS kernel selection. + + ``requested_policy`` and ``selected_kernel`` are equal because selection does + not fall back. ``throughput_latency_candidate`` means the requested shape has at + least one throughput-latency 1CTA profile. ``implementation_ready`` is the + final launchability bit for the requested policy. ``reason`` is a + user-facing explanation for ready or rejected decisions. ``config`` is + ``None`` for throughput 2CTA decisions and for unsupported 1CTA requests. + ``available_profiles`` lists the profile names visible for explicit 1CTA + selection. + """ + + requested_policy: MlaKernelPolicy + selected_kernel: MlaKernelName + throughput_latency_candidate: bool + implementation_ready: bool + reason: str + config: MlaConfig | None + profile_name: str | None = None + available_profiles: tuple[str, ...] = () + + +def normalize_mla_kernel_policy(policy: str) -> MlaKernelPolicy: + """Return a typed MLA TS policy or raise ``ValueError`` for bad input.""" + + if policy not in ("throughput_2cta", "throughput_latency_1cta"): + raise ValueError( + "TS MLA kernel policy must be 'throughput_2cta', " + "or 'throughput_latency_1cta'" + ) + return cast(MlaKernelPolicy, policy) + + +def select_default_mla_kernel_policy( + num_heads: int, + seq_len_q: int, + *, + one_cta_split_kv: int | None = None, + two_cta_split_kv: int | None = None, +) -> MlaKernelPolicy: + """Choose a native family, preferring direct output over split reduction. + + The M64 1CTA schedule owns one native tile of logical query rows. Within + that tile, 2CTA is selected only when it can write the final output directly + while 1CTA would require a split-K reduction. This compares task-graph + structure; it does not contain shape, dtype, or measured crossover tables. + """ + + if (one_cta_split_kv is None) != (two_cta_split_kv is None): + raise ValueError("automatic MLA split topology must be provided as a pair") + if one_cta_split_kv is not None: + if one_cta_split_kv <= 0 or two_cta_split_kv is None or two_cta_split_kv <= 0: + raise ValueError("automatic MLA split counts must be positive") + + if num_heads * seq_len_q <= max(SUPPORTED_TILE_SIZE_Q): + if one_cta_split_kv is not None: + if one_cta_split_kv > 1 and two_cta_split_kv == 1: + return "throughput_2cta" + return "throughput_latency_1cta" + return "throughput_2cta" + + +def resolve_mla_kernel_policy( + policy: str | None, + num_heads: int, + seq_len_q: int, + *, + one_cta_split_kv: int | None = None, + two_cta_split_kv: int | None = None, +) -> tuple[MlaKernelPolicy, str]: + """Resolve an explicit or automatic TS MLA kernel policy.""" + + if policy in (None, "", "auto"): + return ( + select_default_mla_kernel_policy( + num_heads, + seq_len_q, + one_cta_split_kv=one_cta_split_kv, + two_cta_split_kv=two_cta_split_kv, + ), + "auto", + ) + return normalize_mla_kernel_policy(policy), "explicit" + + +def is_throughput_2cta_mla_supported_shape( + *, + batch_size: int, + num_heads: int, + seq_len_q: int, + seq_len_k: int, + latent_dim: int, + rope_dim: int, + page_size: int, + dtype: str, + out_dtype: str = "bf16", +) -> bool: + """Return whether the 2CTA M128 TS MLA path is eligible. + + The predicate is intentionally side-effect free: it only checks the dense + MLA shape constraints that can be known before kernel construction. + """ + + del batch_size + return ( + dtype in ("bf16", "e4m3") + and out_dtype in ("bf16", "e4m3") + and latent_dim == 512 + and rope_dim == 64 + and page_size in SUPPORTED_MLA_PAGE_SIZES + and 1 <= num_heads <= 128 + and seq_len_q >= 1 + and seq_len_k >= 1 + ) + + +def _make_config_for_profile( + *, + batch_size: int, + num_heads: int, + seq_len_q: int, + seq_len_k: int, + latent_dim: int, + rope_dim: int, + page_size: int, + qkv_dtype: str, + o_dtype: str, + profile_name: str | None, + tile_size_q: int | None, + max_active_clusters: int, + explicit_split_kv: int | None, + explicit_persistent: bool | None, +) -> MlaConfig: + """Build the 1CTA config for one candidate policy profile.""" + + return make_throughput_latency_mla_config( + batch_size=batch_size, + num_heads_q=num_heads, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_k, + latent_dim=latent_dim, + rope_dim=rope_dim, + num_tokens_per_page=page_size, + qkv_dtype=qkv_dtype, + o_dtype=o_dtype, + profile=profile_name, + tile_size_q=tile_size_q, + max_active_clusters=max_active_clusters, + explicit_split_kv=explicit_split_kv, + explicit_persistent=explicit_persistent, + ) + + +def select_mla_ts_kernel( + *, + requested_policy: MlaKernelPolicy, + batch_size: int, + num_heads: int, + seq_len_q: int, + seq_len_k: int, + latent_dim: int, + rope_dim: int, + page_size: int, + dtype: str = "bf16", + out_dtype: str = "bf16", + throughput_latency_profile: str | None = None, + throughput_latency_tile_size_q: int | None = None, + max_active_clusters: int, + throughput_latency_split_kv: int | None = None, + throughput_latency_persistent: bool | None = None, +) -> MlaKernelDecision: + """Resolve one of the two explicit TS MLA kernel families. + + ``requested_policy`` is never rewritten to another family. For + throughput 2CTA requests, ``implementation_ready`` mirrors the static 2CTA + eligibility predicate. For throughput-latency 1CTA requests, it mirrors + profile availability and includes the selected profile config. A bad policy + string raises ``ValueError`` through ``normalize_mla_kernel_policy``; a bad + explicit profile raises from the config factory. + """ + + requested_policy = normalize_mla_kernel_policy(requested_policy) + throughput_2cta_candidate = is_throughput_2cta_mla_supported_shape( + batch_size=batch_size, + num_heads=num_heads, + seq_len_q=seq_len_q, + seq_len_k=seq_len_k, + latent_dim=latent_dim, + rope_dim=rope_dim, + page_size=page_size, + dtype=dtype, + out_dtype=out_dtype, + ) + profile_tile_size_q = throughput_latency_tile_size_q + if profile_tile_size_q is None: + profile_tile_size_q = tile_size_q_from_profile_name(throughput_latency_profile) + + profiles = enumerate_throughput_latency_mla_profiles( + batch_size=batch_size, + num_heads_q=num_heads, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_k, + latent_dim=latent_dim, + rope_dim=rope_dim, + num_tokens_per_page=page_size, + tile_size_q=profile_tile_size_q, + max_active_clusters=max_active_clusters, + qkv_dtype=dtype, + explicit_split_kv=throughput_latency_split_kv, + explicit_persistent=throughput_latency_persistent, + ) + profile_names = tuple(profile.name for profile in profiles) + supported_dtype_pair = dtype in ("bf16", "e4m3") and out_dtype in ( + "bf16", + "e4m3", + ) + throughput_latency_candidate = supported_dtype_pair and bool(profiles) + default_profile = None + if throughput_latency_candidate: + default_profile = throughput_latency_profile + if default_profile is None and profile_names: + default_profile = profile_names[0] + + if requested_policy == "throughput_2cta": + reason = ( + "forced throughput 2CTA M128 TS MLA path" + if throughput_2cta_candidate + else ( + "throughput 2CTA M128 TS MLA path requires BF16 or E4M3 input " + "with BF16 or E4M3 output" + if dtype not in ("bf16", "e4m3") or out_dtype not in ("bf16", "e4m3") + else "shape/features outside the throughput 2CTA M128 TS MLA path" + ) + ) + return MlaKernelDecision( + requested_policy=requested_policy, + selected_kernel="throughput_2cta", + throughput_latency_candidate=throughput_latency_candidate, + implementation_ready=throughput_2cta_candidate, + reason=reason, + config=None, + profile_name=None, + available_profiles=profile_names, + ) + + cfg = ( + _make_config_for_profile( + batch_size=batch_size, + num_heads=num_heads, + seq_len_q=seq_len_q, + seq_len_k=seq_len_k, + latent_dim=latent_dim, + rope_dim=rope_dim, + page_size=page_size, + qkv_dtype=dtype, + o_dtype=out_dtype, + profile_name=default_profile, + tile_size_q=throughput_latency_tile_size_q, + max_active_clusters=max_active_clusters, + explicit_split_kv=throughput_latency_split_kv, + explicit_persistent=throughput_latency_persistent, + ) + if throughput_latency_candidate + else None + ) + reason = ( + "forced throughput-latency 1CTA MLA TS path" + if throughput_latency_candidate + else ( + "throughput-latency 1CTA MLA TS path requires BF16 or E4M3 input with BF16 or E4M3 output" + if not supported_dtype_pair + else "shape/features outside the throughput-latency 1CTA MLA TS path" + ) + ) + return MlaKernelDecision( + requested_policy=requested_policy, + selected_kernel="throughput_latency_1cta", + throughput_latency_candidate=throughput_latency_candidate, + implementation_ready=throughput_latency_candidate, + reason=reason, + config=cfg, + profile_name=default_profile, + available_profiles=profile_names, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/parallel_reduction_topology.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/parallel_reduction_topology.py new file mode 100644 index 000000000000..7ee4779fac6e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/parallel_reduction_topology.py @@ -0,0 +1,304 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Host validation and topology for MLA's parallel standalone reducer.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .helpers.constants import MAX_MLA_SPLITS_KV + +_TARGET_SPLITS_PER_RANK = 8 +_PARALLEL_REDUCER_MIN_SPLITS_PER_RANK = 32 +_PARALLEL_REDUCER_MAX_CLUSTER_WAVES = 4 +_SUPPORTED_CLUSTER_SIZES = (1, 2, 4, 8, 16) +_Q128_OUTPUT_ELEMENTS_PER_ROW = 512 +_OUTPUT_ELEMENTS_PER_WORK_UNIT = 1024 +# Some workspace layout products remain 32-bit. Keep the parallel reducer below +# that boundary until every producer and tensor-layout stride is 64-bit safe. +_MAX_PARTIAL_O_WORKSPACE_ELEMENTS = 2**31 - 1 + + +def should_use_q128_g1_parallel_reducer( + *, + batch_size: int, + physical_rows_per_batch: int, + producer_ctas: int, + reference_rows_per_cta: int, + physical_sm_count: int, +) -> bool: + """Prefer one-row G1 only for an underfilled reference launch. + + The Q128 reference reducer coarsens several rows into one CTA. That is + efficient once it fills the machine, but a sub-wave launch leaves most SMs + idle. The one-row G1 reducer is useful only when the producer supplies at + least half a physical-SM wave of split work and its expanded row grid + remains within the same four-wave pressure bound used by the clustered + Q128 topology. These occupancy bounds avoid a per-shape dispatch table. + """ + + _validate_positive_int(batch_size, "batch_size") + _validate_positive_int(physical_rows_per_batch, "physical_rows_per_batch") + _validate_positive_int(producer_ctas, "producer_ctas") + _validate_positive_int(reference_rows_per_cta, "reference_rows_per_cta") + _validate_positive_int(physical_sm_count, "physical_sm_count") + + reference_ctas = ( + batch_size + * (physical_rows_per_batch + reference_rows_per_cta - 1) + // reference_rows_per_cta + ) + parallel_ctas = batch_size * physical_rows_per_batch + return ( + reference_ctas < physical_sm_count + and producer_ctas >= (physical_sm_count + 1) // 2 + and parallel_ctas <= physical_sm_count * _PARALLEL_REDUCER_MAX_CLUSTER_WAVES + ) + + +def _next_power_of_two(value: int) -> int: + return 1 << (value - 1).bit_length() + + +def _validate_splits_kv(splits_kv: int) -> None: + if isinstance(splits_kv, bool) or not isinstance(splits_kv, int): + raise TypeError("splits_kv must be an integer") + if not 1 <= splits_kv <= MAX_MLA_SPLITS_KV: + raise ValueError(f"splits_kv must be in [1, {MAX_MLA_SPLITS_KV}]") + + +def _validate_positive_int(value: int, name: str) -> None: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if value <= 0: + raise ValueError(f"{name} must be positive") + + +def _validate_cluster_size(max_cluster_size: int) -> None: + if max_cluster_size not in _SUPPORTED_CLUSTER_SIZES: + raise ValueError( + f"max_cluster_size must be one of {set(_SUPPORTED_CLUSTER_SIZES)}" + ) + + +def choose_q64_parallel_reducer_cluster_size( + splits_kv: int, + *, + base_clusters: int, + sm_count: int, + max_cluster_size: int = 16, +) -> int: + """Choose Q64 G from actual split work and a four-wave cluster bound. + + Keep at least one warp (32) of actual splits on every rank to amortize + DSMEM publication and merging. A G>1 candidate is accepted only when the + clustered launch fits in four waves; G1 remains the unrestricted fallback. + With S<=128, the per-rank work bound limits production Q64 to G<=4. + """ + + _validate_splits_kv(splits_kv) + _validate_positive_int(base_clusters, "base_clusters") + _validate_positive_int(sm_count, "sm_count") + _validate_cluster_size(max_cluster_size) + + split_limit = max(1, splits_kv // _PARALLEL_REDUCER_MIN_SPLITS_PER_RANK) + split_cluster_size = 1 << (split_limit.bit_length() - 1) + split_cluster_size = min(split_cluster_size, max_cluster_size) + + for cluster_size in reversed(_SUPPORTED_CLUSTER_SIZES): + if cluster_size > split_cluster_size: + continue + clusters_per_wave = sm_count // cluster_size + if clusters_per_wave == 0: + continue + waves = (base_clusters + clusters_per_wave - 1) // clusters_per_wave + if waves <= _PARALLEL_REDUCER_MAX_CLUSTER_WAVES: + return cluster_size + return 1 + + +def validate_parallel_reduction_workspace( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + splits_kv: int, + head_dim: int, +) -> int: + """Validate and return the normalized partial-O workspace element count.""" + + _validate_positive_int(batch_size, "batch_size") + _validate_positive_int(num_heads_q, "num_heads_q") + _validate_positive_int(seq_len_q, "seq_len_q") + _validate_splits_kv(splits_kv) + _validate_positive_int(head_dim, "head_dim") + if splits_kv == 1: + raise ValueError("parallel reduction requires splits_kv in [2, 128]") + + workspace_elements = batch_size * num_heads_q * seq_len_q * splits_kv * head_dim + if workspace_elements > _MAX_PARTIAL_O_WORKSPACE_ELEMENTS: + raise ValueError( + "parallel reduction requires fewer than 2^31 partial-O workspace " + "elements until every layout stride is qualified as 64-bit safe" + ) + return workspace_elements + + +@dataclass(frozen=True) +class ParallelReductionTopology: + """Padded split slots distributed uniformly across one CTA cluster.""" + + actual_splits: int + padded_splits: int + cluster_size: int + slots_per_rank: int + interleaved: bool = False + + @property + def padding_slots(self) -> int: + return self.cluster_size * self.slots_per_rank - self.actual_splits + + def split_for_slot(self, rank: int, slot: int) -> int | None: + """Map a rank-local slot to a split, or ``None`` for padding.""" + + if not 0 <= rank < self.cluster_size: + raise ValueError(f"rank must be in [0, {self.cluster_size})") + if not 0 <= slot < self.slots_per_rank: + raise ValueError(f"slot must be in [0, {self.slots_per_rank})") + + split = ( + slot * self.cluster_size + rank + if self.interleaved + else rank * self.slots_per_rank + slot + ) + return split if split < self.actual_splits else None + + +def make_balanced_parallel_reduction_topology( + splits_kv: int, + *, + cluster_size: int, +) -> ParallelReductionTopology | None: + """Distribute actual splits cyclically with fewer than G padding slots.""" + + _validate_splits_kv(splits_kv) + _validate_cluster_size(cluster_size) + if splits_kv == 1: + return None + slots_per_rank = (splits_kv + cluster_size - 1) // cluster_size + capacity_splits = cluster_size * slots_per_rank + return ParallelReductionTopology( + actual_splits=splits_kv, + padded_splits=capacity_splits, + cluster_size=cluster_size, + slots_per_rank=slots_per_rank, + interleaved=True, + ) + + +def make_parallel_reduction_topology( + splits_kv: int, + *, + max_cluster_size: int = 16, +) -> ParallelReductionTopology | None: + """Build the S2..S128 topology, returning ``None`` for the S1 bypass. + + S2..S16 use exact-capacity G1 so small reductions avoid both padding and + cluster exchange. S17+ retain the power-of-two split capacity and + split-derived power-of-two cluster topology. Slots beyond ``actual_splits`` + must skip loads and computation. + """ + + _validate_splits_kv(splits_kv) + _validate_cluster_size(max_cluster_size) + if splits_kv == 1: + return None + if splits_kv <= 16: + return ParallelReductionTopology( + actual_splits=splits_kv, + padded_splits=splits_kv, + cluster_size=1, + slots_per_rank=splits_kv, + ) + + padded_splits = _next_power_of_two(splits_kv) + cluster_size = min( + max_cluster_size, + _next_power_of_two( + (splits_kv + _TARGET_SPLITS_PER_RANK - 1) // _TARGET_SPLITS_PER_RANK + ), + ) + return ParallelReductionTopology( + actual_splits=splits_kv, + padded_splits=padded_splits, + cluster_size=cluster_size, + slots_per_rank=padded_splits // cluster_size, + ) + + +def make_q128_wave_limited_parallel_reduction_topology( + splits_kv: int, + *, + logical_rows: int, + physical_sm_count: int, + max_cluster_size: int = 16, +) -> ParallelReductionTopology | None: + """Build Q128/D512 topology targeting four output-work waves. + + One base work unit is 1,024 output elements. Start from the split-derived + topology, then reduce G until those work units fit within a four-wave + pressure target, or until G1 is reached. This is a work proxy rather than a + literal count of launched CTA waves. If G changes, recompute the minimum + contiguous split capacity supported by the retained ranks. + """ + + _validate_splits_kv(splits_kv) + _validate_positive_int(logical_rows, "logical_rows") + _validate_positive_int(physical_sm_count, "physical_sm_count") + _validate_cluster_size(max_cluster_size) + + topology = make_parallel_reduction_topology( + splits_kv, + max_cluster_size=max_cluster_size, + ) + if topology is None: + return None + base_work_units = ( + logical_rows * _Q128_OUTPUT_ELEMENTS_PER_ROW + + _OUTPUT_ELEMENTS_PER_WORK_UNIT + - 1 + ) // _OUTPUT_ELEMENTS_PER_WORK_UNIT + cluster_size = topology.cluster_size + while cluster_size > 1: + work_units_per_wave = physical_sm_count // cluster_size + if work_units_per_wave == 0: + cluster_size //= 2 + continue + pressure_waves = ( + base_work_units + work_units_per_wave - 1 + ) // work_units_per_wave + if pressure_waves <= _PARALLEL_REDUCER_MAX_CLUSTER_WAVES: + break + cluster_size //= 2 + + if cluster_size == topology.cluster_size: + return topology + slots_per_rank = (splits_kv + cluster_size - 1) // cluster_size + return ParallelReductionTopology( + actual_splits=splits_kv, + padded_splits=cluster_size * slots_per_rank, + cluster_size=cluster_size, + slots_per_rank=slots_per_rank, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/__init__.py new file mode 100644 index 000000000000..e9d40fca344c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Two-CTA throughput policy for task-scheduled MLA decode.""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py new file mode 100644 index 000000000000..af89becc56eb --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py @@ -0,0 +1,602 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration for the throughput 2CTA MLA decode TS kernel. + +The throughput 2CTA policy uses a 2CTA M128 schedule. BF16 and FP8 inputs share +the same explicit configuration structure and output-reduction contract. +""" + +from dataclasses import dataclass +from typing import Tuple + +from ..helpers.constants import MAX_MLA_SPLITS_KV, SUPPORTED_MLA_PAGE_SIZES +from ..helpers.mask import MaskType, normalize_mask_type + + +# Softmax converts natural-scale scores to exp2 with log2(e). +LOG2_E = 1.4426950408889634074 + +# The separate reducer follows the public MLA output contract: partial O is +# stored as BF16 while LSE and the final accumulation remain FP32. One +# 512-thread CTA owns eight D512 rows, with each thread moving one 16-byte +# BF16 vector. Keeping these values derived makes the row packing explicit +# and prevents the launch geometry from drifting away from the workspace +# representation. +PARTIAL_O_BITS = 16 +REDUCTION_THREADS_PER_CTA = 512 +REDUCTION_VECTOR_BYTES = 16 +REDUCTION_VALUES_PER_THREAD = REDUCTION_VECTOR_BYTES * 8 // PARTIAL_O_BITS +REDUCTION_THREADS_PER_ROW = 512 // REDUCTION_VALUES_PER_THREAD +REDUCTION_ROWS_PER_CTA = REDUCTION_THREADS_PER_CTA // REDUCTION_THREADS_PER_ROW + +# PV consumes V from SMEM in 32-token K blocks. Physical KV pages may be +# smaller or larger, but TMA must assemble this fixed block geometry before +# tcgen05 advances the V descriptor to the next K block. +V_SMEM_K_BLOCK_TOKENS = 32 + +# Each transpose-TMA issue stages one 64-element slice of the latent dimension. +V_TMA_LATENT_ELEMENTS = 64 + + +def ceil_div(a: int, b: int) -> int: + """Return ``ceil(a / b)`` for positive integer divisors.""" + + if b <= 0: + raise ValueError(f"divisor must be positive, got {b}") + return (a + b - 1) // b + + +def compute_split_kv( + *, + batch_size: int, + num_q_tiles: int, + seq_len_kv: int, + mma_qk_tiler_mn: Tuple[int, int] = (128, 128), + max_active_blocks: int, +) -> int: + """Choose the throughput 2CTA split-KV count for a concrete launch. + + The heuristic tries to expose enough K-split work to fill the available CTA + slots without creating extra partial K waves. The result is capped to keep + the reduction grid bounded. + """ + + if batch_size <= 0: + raise ValueError(f"batch_size must be positive, got {batch_size}") + if num_q_tiles <= 0: + raise ValueError(f"num_q_tiles must be positive, got {num_q_tiles}") + if seq_len_kv <= 0: + raise ValueError(f"seq_len_kv must be positive, got {seq_len_kv}") + if max_active_blocks <= 0: + raise ValueError(f"max_active_blocks must be positive, got {max_active_blocks}") + if mma_qk_tiler_mn[1] <= 0: + raise ValueError( + f"mma_qk_tiler_mn[1] must be positive, got {mma_qk_tiler_mn[1]}" + ) + + max_splits = ceil_div(seq_len_kv, mma_qk_tiler_mn[1]) + blocks_per_batch = max(1, max_active_blocks // batch_size // (num_q_tiles * 2)) + split_heur = min(max_splits, blocks_per_batch) + k_waves = ceil_div(max_splits, split_heur) + split_wave_aware = ceil_div(max_splits, k_waves) + return min(split_wave_aware, MAX_MLA_SPLITS_KV) + + +def compute_workspace_size( + *, + tile_size_q: int, + num_q_tiles: int, + latent_dim: int, + batch_size: int, + split_kv: int, + partial_o_dtype, + lse_dtype, +) -> int: + """Return the physical flat-tile split-KV workspace size in bytes.""" + + if split_kv == 1: + return 0 + if tile_size_q <= 0: + raise ValueError(f"tile_size_q must be positive, got {tile_size_q}") + if num_q_tiles <= 0: + raise ValueError(f"num_q_tiles must be positive, got {num_q_tiles}") + partial_rows = batch_size * tile_size_q * num_q_tiles * split_kv + return partial_rows * ( + latent_dim * partial_o_dtype.width // 8 + lse_dtype.width // 8 + ) + + +@dataclass +class MlaDecodeConfig: + """MLA decode kernel configuration. + + All values are plain Python ints/tuples so they can be used as + ``Constexpr`` in DSL code. + """ + + # Architecture. Dense MLA uses 512 latent channels plus 64 RoPE channels; + # the throughput path uses one 2-CTA cluster per M128 tile. + latent_dim: int = 512 + rope_dim: int = 64 + num_mma_ctas: int = 2 + cluster_shape_mnk: Tuple[int, int, int] = (2, 1, 1) + + # MMA tile shapes. QK runs over 128x128 K tiles, while PV writes the 512 V + # head in two 256-column passes. + mma_qk_tiler_mn: Tuple[int, int] = (128, 128) + mma_pv_tiler_mn: Tuple[int, int] = (128, 256) + mma_qk_tiler_k: int = 64 # = rope_dim + mma_qk_tiler: Tuple[int, int, int] = (128, 128, 64) + mma_qk_rope_tiler: Tuple[int, int, int] = (128, 128, 64) + mma_pv_tiler: Tuple[int, int, int] = (128, 256, 32) + + # Iteration counts derived from the fixed dense MLA dimensions and MMA + # tile shapes above. + iterations_qk_latent: int = 8 # latent_dim / mma_qk_tiler_k = 512/64 + iterations_qk_rope: int = 1 # rope_dim / mma_qk_tiler_k = 64/64 + iterations_qk: int = 9 # latent + rope + iterations_pv_k: int = 4 # mma_qk_tiler[1] / mma_pv_tiler[2] = 128/32 + iterations_pv_n: int = 2 # latent_dim / mma_pv_tiler[1] = 512/256 + # BF16 keeps the hardware-legal K64/PV-K32 transactions, but publishes two + # consecutive transactions under one TMA/UMMA pipeline stage. This + # matches a 128-element head-dimension stage without requiring a TMA box + # wider than the 128-byte swizzle permits. FP8 already uses native K128 + # QK stages and separate whole-tile K/V resources. + kv_subtiles_per_stage: int = 2 + iterations_qk_latent_stages: int = 4 + iterations_qk_stages: int = 5 + iterations_pv_stages: int = 4 + + # Pipeline stage counts for the captured schedule resources. The combined + # K/V stage count keeps enough delayed-V stages live for K-before-V overlap. + load_q_stage: int = 1 + load_k_stage: int = 3 + load_v_stage: int = 2 + load_kv_stage: int = 7 + mma_s_stage: int = 2 + p_mma_stage: int = 2 + p_cor_stage: int = 2 + mma_o_stage: int = 1 + + # Base BF16 warp assignments for the 12-warp CTA. Softmax and correction + # each own a contiguous four-warp group; the remaining warps issue MMA and + # TMA (including register-held page IDs) or provide scheduler/alignment + # roles. The FP8 factory extends the CTA to 16 warps for its second softmax + # group and split QK/PV schedule. + compute_warp_ids: Tuple[int, ...] = (0, 1, 2, 3) + correction_warp_ids: Tuple[int, ...] = (4, 5, 6, 7) + mma_warp_id: int = 8 + load_tma_warp_id: int = 9 + pv_mma_warp_id: int = 11 + empty_warp_ids: Tuple[int, ...] = (11,) + second_compute_warp_ids: Tuple[int, ...] = () + num_softmax_groups: int = 1 + + num_compute_warps: int = 4 + threads_per_warp: int = 32 + threads_per_cta: int = 384 # 12 warps * 32 + warps_in_n: int = 2 + + # Register budgets passed to setmaxnreg for the high-register softmax and + # correction groups; all other warps use the lower shared budget. + softmax_reg_num: int = 192 + correction_reg_num: int = 208 + other_reg_num: int = 96 + + # Named barrier IDs and thread counts. IDs are local to this kernel's + # manual synchronization protocol and are kept away from the TMEM barrier. + softmax_sync_bar_id: int = 2 + softmax_sync_threads: int = 128 # 4 warps * 32 + epilogue_sync_bar_id: int = 3 + epilogue_sync_threads: int = 128 # 4 warps * 32 + softmax_order_bar_0_id: int = 5 + softmax_order_bar_1_id: int = 6 + + # TMEM sync barrier (for alloc/dealloc) + tmem_sync_bar_id: int = 1 + tmem_sync_bar_threads: int = 0 # computed in make_config + + # TMEM layout offsets. The full 512-column TMEM budget is reserved so S, + # O, and correction-factor columns can use fixed offsets. + num_tmem_cols: int = 512 + tmem_o_offset: int = 0 # computed + correction_factor_offset: int = 0 # computed + + # SMEM element counts (per-stage or total) + smem_q_latent_elems: int = 0 + smem_q_rope_elems: int = 0 + smem_kc_elems: int = 0 + smem_vc_elems: int = 0 + smem_k_stage_elems: int = 0 + smem_v_stage_elems: int = 0 + smem_p_elems: int = 0 + softmax_exchange_elems: int = 128 + + # Page geometry. Physical page-table geometry is independent of the V + # transpose-TMA microtile used to assemble the tcgen05 SMEM operand. + page_size: int = 32 + kc_page_tile_size: int = 32 + v_tma_token_count: int = 32 + + # Data types + qkv_dtype: str = "bf16" + o_dtype: str = "bf16" + qkv_dtype_bytes: int = 2 + o_dtype_bytes: int = 2 + use_bf16_output: int = 1 + use_fp8_output: int = 0 + + # TMA byte counts + tma_copy_q_bytes: int = 0 + tma_copy_kc_bytes: int = 0 + tma_copy_vc_bytes: int = 0 + tma_copy_k_tile_bytes: int = 0 + tma_copy_v_tile_bytes: int = 0 + tma_kc_subtile_bytes: int = 0 + tma_vc_subtile_bytes: int = 0 + + # Scheduling. The runner normally supplies max_active_clusters from + # HardwareInfo; 56 is the SM100-class fallback used by local construction + # paths that do not query hardware. + use_fp8_split_mma_schedule: bool = False + use_fp8_dual_softmax_schedule: bool = False + max_active_clusters: int = 56 + is_persistent: bool = True + is_var_seq: bool = False + # Use block_split_kvs[batch] as a per-batch cap before runtime K contracts + # the useful split prefix. Grid and workspace geometry retain the maximum. + is_var_split_kv: bool = False + # Causal is bottom-right aligned for speculative decode. Dense still masks + # the ordinary per-batch KV tail at ``cache_seqs[batch]``. + mask_type: str = MaskType.CAUSAL.value + + @property + def tokens_per_k_tile(self) -> int: + """Return the logical KV-token count covered by one QK tile.""" + + return self.mma_qk_tiler[1] + + @property + def tokens_per_k_cta(self) -> int: + """Return the KV-token count owned by one CTA in the 2CTA cluster.""" + + return self.tokens_per_k_tile // self.num_mma_ctas + + @property + def pages_per_k_tile(self) -> int: + """Return the physical page count spanned by one logical K tile.""" + + return self.tokens_per_k_tile // self.page_size + + @property + def pages_per_k_cta(self) -> int: + """Return page IDs consumed by one CTA, including shared-page K tiles.""" + + return max(1, ceil_div(self.pages_per_k_tile, self.num_mma_ctas)) + + @property + def tokens_per_v_tile(self) -> int: + """Return the logical KV-token count covered by all PV K iterations.""" + + return self.mma_pv_tiler[2] * self.iterations_pv_k + + @property + def pages_per_v_tile(self) -> int: + """Return the physical page count spanned by one logical V tile.""" + + return self.tokens_per_v_tile // self.page_size + + @property + def pages_per_v_subtile(self) -> int: + """Return physical page IDs consumed by one staged BF16 PV iteration.""" + + return max(1, ceil_div(self.pages_per_v_tile, self.iterations_pv_k)) + + @property + def v_subtiles_per_page(self) -> int: + """Return staged BF16 PV iterations that share one physical page ID.""" + + return max(1, ceil_div(self.iterations_pv_k, self.pages_per_v_tile)) + + @property + def v_tma_copies_per_subtile(self) -> int: + """Return transpose-TMA copies needed to assemble one PV K subtile.""" + + return self.mma_pv_tiler[2] // self.v_tma_token_count + + def is_fp8_qkv(self) -> bool: + """Return whether Q/K/V tensors use E4M3 data.""" + + return self.qkv_dtype == "e4m3" + + +def make_mla_decode_config( + mma_qk_tiler_mn: Tuple[int, int] = (128, 128), + mma_pv_tiler_mn: Tuple[int, int] = (128, 256), + rope_dim: int = 64, + page_size: int = 32, + qkv_dtype: str = "bf16", + o_dtype: str = "bf16", + max_active_clusters: int = 56, + is_persistent: bool = True, + is_var_seq: bool = False, + is_var_split_kv: bool = False, + mask_type: MaskType | str = MaskType.CAUSAL, +) -> MlaDecodeConfig: + """Create and populate a MlaDecodeConfig from problem parameters.""" + cfg = MlaDecodeConfig() + cfg.mma_qk_tiler_mn = mma_qk_tiler_mn + cfg.mma_pv_tiler_mn = mma_pv_tiler_mn + cfg.rope_dim = rope_dim + cfg.page_size = page_size + cfg.qkv_dtype = qkv_dtype + cfg.o_dtype = o_dtype + cfg.max_active_clusters = max_active_clusters + cfg.is_persistent = is_persistent + cfg.is_var_seq = is_var_seq + cfg.is_var_split_kv = is_var_split_kv + cfg.mask_type = normalize_mask_type(mask_type) + + def _require_positive(name: str, value: int) -> None: + """Validate that a named configuration value is positive.""" + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + + def _require_divisible( + dividend_name: str, dividend: int, divisor_name: str, divisor: int + ) -> None: + """Validate that one named configuration value divides another.""" + _require_positive(dividend_name, dividend) + _require_positive(divisor_name, divisor) + if dividend % divisor != 0: + raise ValueError( + f"{dividend_name}={dividend} must be divisible by " + f"{divisor_name}={divisor}" + ) + + for idx, value in enumerate(mma_qk_tiler_mn): + _require_positive(f"mma_qk_tiler_mn[{idx}]", value) + for idx, value in enumerate(mma_pv_tiler_mn): + _require_positive(f"mma_pv_tiler_mn[{idx}]", value) + if rope_dim < 0: + raise ValueError(f"rope_dim must be non-negative, got {rope_dim}") + _require_positive("page_size", page_size) + if page_size not in SUPPORTED_MLA_PAGE_SIZES: + raise ValueError( + f"page_size must be one of {SUPPORTED_MLA_PAGE_SIZES}, got {page_size}" + ) + if page_size > mma_qk_tiler_mn[1] or mma_qk_tiler_mn[1] % page_size != 0: + raise ValueError( + "page_size must exactly partition the throughput 2CTA K tile: " + f"mma_qk_tiler_mn[1]={mma_qk_tiler_mn[1]}, page_size={page_size}" + ) + + if qkv_dtype not in ("bf16", "e4m3"): + raise ValueError(f"unsupported qkv_dtype={qkv_dtype!r}") + if o_dtype not in ("bf16", "e4m3"): + raise ValueError(f"unsupported o_dtype={o_dtype!r}") + cfg.qkv_dtype_bytes = 1 if qkv_dtype == "e4m3" else 2 + cfg.o_dtype_bytes = 1 if o_dtype == "e4m3" else 2 + cfg.use_bf16_output = int(o_dtype == "bf16") + cfg.use_fp8_output = int(o_dtype == "e4m3") + cfg.use_fp8_split_mma_schedule = qkv_dtype == "e4m3" + cfg.use_fp8_dual_softmax_schedule = qkv_dtype == "e4m3" + cfg.empty_warp_ids = () if cfg.use_fp8_split_mma_schedule else (cfg.pv_mma_warp_id,) + if cfg.use_fp8_split_mma_schedule: + cfg.second_compute_warp_ids = (12, 13, 14, 15) + cfg.num_softmax_groups = 2 + cfg.threads_per_cta = cfg.threads_per_warp * 16 + cfg.softmax_reg_num = 160 + cfg.correction_reg_num = 160 + cfg.other_reg_num = 32 + cfg.mma_o_stage = 2 + cfg.epilogue_sync_bar_id = 4 + + # Derived MMA tilers. FP8 latent QK uses K=128 while the separate RoPE MMA + # keeps K=64. + if cfg.rope_dim > 0: + cfg.mma_qk_tiler_k = cfg.rope_dim * (2 if qkv_dtype == "e4m3" else 1) + else: + cfg.mma_qk_tiler_k = 128 if qkv_dtype == "e4m3" else 64 + _require_divisible( + "latent_dim", cfg.latent_dim, "mma_qk_tiler_k", cfg.mma_qk_tiler_k + ) + _require_divisible( + "mma_qk_tiler_mn[1] * mma_qk_tiler_k", + mma_qk_tiler_mn[1] * cfg.mma_qk_tiler_k, + "mma_pv_tiler_mn[1]", + mma_pv_tiler_mn[1], + ) + cfg.mma_qk_tiler = (mma_qk_tiler_mn[0], mma_qk_tiler_mn[1], cfg.mma_qk_tiler_k) + cfg.mma_qk_rope_tiler = (mma_qk_tiler_mn[0], mma_qk_tiler_mn[1], cfg.rope_dim) + pv_k = mma_qk_tiler_mn[1] * cfg.mma_qk_tiler_k // mma_pv_tiler_mn[1] + _require_divisible("mma_qk_tiler_mn[1]", mma_qk_tiler_mn[1], "pv_k", pv_k) + _require_divisible( + "latent_dim", cfg.latent_dim, "mma_pv_tiler_mn[1]", mma_pv_tiler_mn[1] + ) + cfg.mma_pv_tiler = (mma_pv_tiler_mn[0], mma_pv_tiler_mn[1], pv_k) + + # Iteration counts + cfg.iterations_qk_latent = cfg.latent_dim // cfg.mma_qk_tiler_k + cfg.iterations_qk_rope = 1 if cfg.rope_dim > 0 else 0 + cfg.iterations_qk = cfg.iterations_qk_latent + cfg.iterations_qk_rope + cfg.iterations_pv_k = cfg.mma_qk_tiler[1] // cfg.mma_pv_tiler[2] + cfg.iterations_pv_n = cfg.latent_dim // cfg.mma_pv_tiler[1] + cfg.kv_subtiles_per_stage = 1 if qkv_dtype == "e4m3" else 2 + _require_divisible( + "iterations_qk_latent", + cfg.iterations_qk_latent, + "kv_subtiles_per_stage", + cfg.kv_subtiles_per_stage, + ) + _require_divisible( + "iterations_pv_k * iterations_pv_n", + cfg.iterations_pv_k * cfg.iterations_pv_n, + "kv_subtiles_per_stage", + cfg.kv_subtiles_per_stage, + ) + cfg.iterations_qk_latent_stages = ( + cfg.iterations_qk_latent // cfg.kv_subtiles_per_stage + ) + cfg.iterations_qk_stages = cfg.iterations_qk_latent_stages + cfg.iterations_qk_rope + cfg.iterations_pv_stages = ( + cfg.iterations_pv_k * cfg.iterations_pv_n // cfg.kv_subtiles_per_stage + ) + if cfg.tokens_per_v_tile != cfg.tokens_per_k_tile: + raise ValueError( + "PV K iterations must cover the same token tile as QK: " + f"tokens_per_v_tile={cfg.tokens_per_v_tile}, " + f"tokens_per_k_tile={cfg.tokens_per_k_tile}" + ) + + # Page-offset tile sizes + num_mma_ctas = cfg.cluster_shape_mnk[0] + cfg.num_mma_ctas = num_mma_ctas + _require_divisible( + "mma_qk_tiler_mn[0]", cfg.mma_qk_tiler[0], "num_mma_ctas", num_mma_ctas + ) + _require_divisible( + "mma_qk_tiler_mn[1]", cfg.mma_qk_tiler[1], "num_mma_ctas", num_mma_ctas + ) + if ( + cfg.page_size != cfg.tokens_per_k_tile + and cfg.tokens_per_k_cta % cfg.page_size != 0 + ): + raise ValueError( + "page_size must partition each CTA's K tile unless both CTAs share " + "one full-tile page: " + f"tokens_per_k_cta={cfg.tokens_per_k_cta}, page_size={cfg.page_size}" + ) + _require_divisible( + "mma_pv_tiler_mn[0]", cfg.mma_pv_tiler[0], "num_mma_ctas", num_mma_ctas + ) + _require_divisible( + "mma_pv_tiler_mn[1]", cfg.mma_pv_tiler[1], "num_mma_ctas", num_mma_ctas + ) + cfg.kc_page_tile_size = min(cfg.page_size, cfg.tokens_per_k_cta) + cfg.v_tma_token_count = min(cfg.page_size, V_SMEM_K_BLOCK_TOKENS) + _require_divisible( + "V_SMEM_K_BLOCK_TOKENS", + V_SMEM_K_BLOCK_TOKENS, + "v_tma_token_count", + cfg.v_tma_token_count, + ) + _require_divisible( + "mma_pv_tiler[2]", + cfg.mma_pv_tiler[2], + "v_tma_token_count", + cfg.v_tma_token_count, + ) + + # SMEM sizes (elements) + cfg.smem_q_latent_elems = ( + (cfg.mma_qk_tiler[0] // num_mma_ctas) + * cfg.mma_qk_tiler[2] + * cfg.iterations_qk_latent + * cfg.load_q_stage + ) + cfg.smem_q_rope_elems = ( + (cfg.mma_qk_rope_tiler[0] // num_mma_ctas) + * cfg.mma_qk_rope_tiler[2] + * cfg.load_q_stage + ) + k_latent_subtile_elems = cfg.mma_qk_tiler[1] // num_mma_ctas * cfg.mma_qk_tiler[2] + k_rope_subtile_elems = ( + cfg.mma_qk_rope_tiler[1] // num_mma_ctas * cfg.mma_qk_rope_tiler[2] + ) + if qkv_dtype == "e4m3": + cfg.smem_k_stage_elems = ( + k_latent_subtile_elems * cfg.iterations_qk_latent + + k_rope_subtile_elems * cfg.iterations_qk_rope + ) + cfg.smem_kc_elems = cfg.smem_k_stage_elems * cfg.load_k_stage + else: + cfg.smem_k_stage_elems = k_latent_subtile_elems * cfg.kv_subtiles_per_stage + cfg.smem_kc_elems = cfg.smem_k_stage_elems * cfg.load_kv_stage + v_subtile_elems = cfg.mma_pv_tiler[1] // num_mma_ctas * cfg.mma_pv_tiler[2] + cfg.smem_v_stage_elems = v_subtile_elems * cfg.iterations_pv_k * cfg.iterations_pv_n + cfg.smem_vc_elems = ( + cfg.smem_v_stage_elems * cfg.load_v_stage if qkv_dtype == "e4m3" else 0 + ) + cfg.smem_p_elems = ( + (cfg.mma_pv_tiler[0] // num_mma_ctas) + * cfg.mma_pv_tiler[2] + * cfg.iterations_pv_k + * cfg.p_mma_stage + ) + cfg.softmax_exchange_elems = ( + cfg.num_compute_warps + * cfg.threads_per_warp + * (1 if not cfg.use_fp8_dual_softmax_schedule else 2) + ) + + # TMEM layout + cfg.tmem_o_offset = cfg.mma_s_stage * cfg.mma_qk_tiler[1] // cfg.warps_in_n + cfg.correction_factor_offset = cfg.tmem_o_offset + cfg.latent_dim // cfg.warps_in_n + + # TMA byte counts + q_latent_tile_bytes = ( + cfg.mma_qk_tiler[0] // num_mma_ctas * cfg.mma_qk_tiler[2] * cfg.qkv_dtype_bytes + ) + q_rope_tile_bytes = ( + cfg.mma_qk_rope_tiler[0] + // num_mma_ctas + * cfg.mma_qk_rope_tiler[2] + * cfg.qkv_dtype_bytes + ) + cfg.tma_copy_q_bytes = ( + q_latent_tile_bytes * num_mma_ctas * cfg.iterations_qk_latent + + q_rope_tile_bytes * num_mma_ctas * cfg.iterations_qk_rope + ) + cfg.tma_copy_kc_bytes = ( + cfg.mma_qk_tiler[1] + // num_mma_ctas + * cfg.mma_qk_tiler[2] + * cfg.qkv_dtype_bytes + * num_mma_ctas + ) + cfg.tma_copy_vc_bytes = ( + cfg.mma_pv_tiler[1] + // num_mma_ctas + * cfg.mma_pv_tiler[2] + * cfg.qkv_dtype_bytes + * num_mma_ctas + ) + if cfg.tma_copy_kc_bytes != cfg.tma_copy_vc_bytes: + raise ValueError( + "K and V TMA subtile byte counts must match: " + f"tma_copy_kc_bytes={cfg.tma_copy_kc_bytes}, " + f"tma_copy_vc_bytes={cfg.tma_copy_vc_bytes}" + ) + cfg.tma_kc_subtile_bytes = cfg.tma_copy_kc_bytes * cfg.kv_subtiles_per_stage + cfg.tma_vc_subtile_bytes = cfg.tma_copy_vc_bytes * cfg.kv_subtiles_per_stage + cfg.tma_copy_k_tile_bytes = ( + cfg.smem_k_stage_elems * cfg.qkv_dtype_bytes * num_mma_ctas + ) + cfg.tma_copy_v_tile_bytes = ( + cfg.tma_copy_vc_bytes * cfg.iterations_pv_k * cfg.iterations_pv_n + ) + + # TMEM sync barrier thread count: QK MMA + softmax + correction, plus + # the FP8 PV-MMA warp when it participates in TMEM O production. + cfg.tmem_sync_bar_threads = ( + cfg.threads_per_warp * (2 if cfg.use_fp8_split_mma_schedule else 1) + + cfg.threads_per_warp * cfg.num_compute_warps + + cfg.threads_per_warp * cfg.num_compute_warps + ) + if cfg.use_fp8_dual_softmax_schedule: + cfg.tmem_sync_bar_threads += cfg.threads_per_warp * cfg.num_compute_warps + + return cfg diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py new file mode 100644 index 000000000000..5ea23ffb4238 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py @@ -0,0 +1,1873 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Throughput 2CTA MLA decode TS kernel implementation. + +Warp-specialized MLA decode kernel using the CUTLASS task-scheduling framework. + +The graph depends on dtype and scheduler policy. BF16 uses 12 warps and a +combined TMA/MMA schedule; CLC adds a scheduler task to that graph. FP8 uses +16 warps, separate K/V and QK/PV tasks, and two softmax groups. In both paths, +softmax and correction own the first two four-warp groups. + +Entry points: + - build_mla_decode_task_manager() -- pure Python, used for validation only + - MlaDecodeTs -- class with @cute.kernel for GPU execution +""" + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.experimental.cuda as cuda +from cutlass import Int32, Int64 +from cutlass.cute.testing import assert_ as runtime_assert +from cutlass.cute.nvgpu import OperandMajorMode, tcgen05 +from cutlass.experimental import primitives as prims + +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + PipelineConfig, + TileSchedulerConfig, +) +from cutlass.experimental.task_scheduling.enums import SignalingThreads +from cutlass.experimental.task_scheduling.task_manager import TaskManager +from ...tensor_map import ( + create_tensor_map_ragged_from_tensor, + create_tensor_map_tiled_from_view, +) + +from .config import ( + LOG2_E, + REDUCTION_ROWS_PER_CTA, + REDUCTION_THREADS_PER_ROW, + V_TMA_LATENT_ELEMENTS, + MlaDecodeConfig, + make_mla_decode_config, +) +from ..helpers.constants import MAX_MLA_SPLITS_KV, TMEM_DEALLOC_MBAR_THREADS +from ..helpers.mask import MaskType, mask_visible_k_length, normalize_mask_type +from ..helpers.query import ( + FlatQueryTileLayout, + flat_query_row_state, + query_batch_bounds, + runtime_flat_query_tile_has_rows, +) +from .work_partition import ( + runtime_split_kv_cap, + runtime_split_tile_range, +) +from .resources import ( + PageOffsetWindowResource, + SmemQResource, + SmemKResource, + SmemKVResource, + SmemVResource, + SmemPResource, + TmemSResource, + TmemCorrResource, + TmemOResource, + GmemOResource, + MlaWorkQueue, + WorkThrottleBarrierResource, +) +from ..helpers.math import qkv_dtype +from ..parallel_reduction_topology import ( + ParallelReductionTopology, + make_q128_wave_limited_parallel_reduction_topology, + should_use_q128_g1_parallel_reducer, + validate_parallel_reduction_workspace, +) +from .parallel_reduction import ( + PARALLEL_REDUCTION_HEAD_DIM, + PARALLEL_REDUCTION_THREADS, + run_parallel_reduction_kernel, +) +from .reduction import run_reduction_kernel +from ..helpers.tile_scheduler import ( + MLAStaticTileSchedulerParams, + MLAStaticTileScheduler, + create_mla_static_tile_scheduler_params, + divmod_constexpr_power_of_two_or_fdd, +) +from .tasks import ( + MlaClcTask, + MlaInterleavedTask, + MlaTask, + create_load_k_task, + create_load_v_task, + create_load_tma_task, + create_mma_task, + create_mma_qk_direct_task, + create_mma_pv_direct_task, + create_softmax_task, + create_correction_task, + create_padding_task, + create_scheduler_task, +) + + +def ceil_div(a, b): + """Return the ceiling of a divided by b.""" + return (a + b - 1) // b + + +def build_mla_decode_task_manager( + cfg: MlaDecodeConfig, + # SMEM arrays (None for validate-only) + smem_q_latent_arr=None, + smem_q_rope_arr=None, + smem_kc_arr=None, + smem_vc_arr=None, + smem_p_arr=None, + # TMA descriptors (None for validate-only) + tma_desc_q_latent=None, + tma_desc_q_rope=None, + tma_desc_c_latent=None, + tma_desc_c_rope=None, + tma_desc_c_transpose=None, + # Page-offset tensor + page_offsets=None, + # Runtime coordinates + blk_coord=None, + tidx=None, + # GMEM output tensors + output=None, + acc_output=None, + lse=None, + acc_lse=None, + # Domain (k_tile_count) + domain=4, + # Persistent loop support + work_queue=None, + # For k_index_base computation (split_kv support) + cache_seqs=None, + cu_seqlens_q=None, + split_kv=None, + logical_num_heads_q=128, + logical_seq_len_q=1, + tiled_mma_qk=None, + verbose=False, + exhaustive_deadlock_race_check=True, +) -> "tuple[TaskManager, list[MemoryResource], dict[str, MemoryResource]]": + """Build the MLA decode TaskManager with all resources, tasks, and dependency graph. + + The throughput 2CTA graph uses Q/K/V TMA loads, a register-held BF16 page-ID + window, QK/PV MMA, softmax, and correction/store tasks. Validation-only + calls pass a concrete integer ``domain`` and no runtime tensors; JIT calls + pass symbolic runtime state and skip schedule validation. The dependency + graph relies on TaskManager DMA-order validation to keep SMEM producers + alive until async TMEM consumers have launched. + + Parameters + ---------- + cfg : MlaDecodeConfig + Kernel-wide configuration. + domain : int or symbolic + Number of k-tile iterations (loop domain for all tasks). + Use an integer for validation-only mode. + + Returns + ------- + TaskManager, list[MemoryResource], dict + Configured TaskManager, TMEM resources list, named resource dict. + """ + # ────────────────────────────────────────────────────────────── + # Cluster / CTA layout + # ────────────────────────────────────────────────────────────── + cluster_shape_vmnk = (cfg.num_mma_ctas, 1, 1, 1) + WARP_SIZE = 32 + Agent = pipeline.Agent + use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + use_work_throttle = use_clc_dynamic and not cfg.use_fp8_split_mma_schedule + non_interleaved_task_class = MlaClcTask if use_clc_dynamic else MlaTask + task_domain = ( + MlaClcTask.get_domain + if use_clc_dynamic and not isinstance(domain, int) + else domain + ) + + # Cooperative groups + tma_producer_group = pipeline.CooperativeGroup(Agent.Thread) # elect_one TMA + umma_hw_group = pipeline.CooperativeGroup(Agent.Thread) # UMMA hardware + + # Softmax warps: 4 warps * 32 threads * 2 CTAs for cluster-scoped consumer + compute_group_cluster = pipeline.CooperativeGroup( + Agent.Thread, cfg.num_compute_warps * WARP_SIZE * cfg.num_mma_ctas + ) + # Softmax warps: 4 warps * 32 threads (local CTA only) + compute_group_local = pipeline.CooperativeGroup( + Agent.Thread, cfg.num_compute_warps * WARP_SIZE + ) + # Correction warps: same layout + correction_group_cluster = pipeline.CooperativeGroup( + Agent.Thread, cfg.num_compute_warps * WARP_SIZE * cfg.num_mma_ctas + ) + correction_group_local = pipeline.CooperativeGroup( + Agent.Thread, cfg.num_compute_warps * WARP_SIZE + ) + # ────────────────────────────────────────────────────────────── + # Pipeline configs + # ────────────────────────────────────────────────────────────── + + # SmemQ: TmaUmma, 1 stage, LoadTma -> Mma + smem_q_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.load_q_stage, + num_bytes=cfg.tma_copy_q_bytes, + producer_group=tma_producer_group, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + num_bytes_per_warp_per_cta=(cfg.tma_copy_q_bytes // cfg.num_mma_ctas), + ) + + if cfg.use_fp8_split_mma_schedule: + # FP8 uses independent whole-tile K and V pipelines so QK and PV have + # separate consumer states. This mirrors the native FP8 schedule shape + # without changing the BF16 combined-KV path. + smem_k_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.load_k_stage, + num_bytes=cfg.tma_copy_k_tile_bytes, + producer_group=tma_producer_group, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + num_bytes_per_warp_per_cta=(cfg.tma_copy_k_tile_bytes // cfg.num_mma_ctas), + ) + smem_v_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.load_v_stage, + num_bytes=cfg.tma_copy_v_tile_bytes, + producer_group=tma_producer_group, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + num_bytes_per_warp_per_cta=(cfg.tma_copy_v_tile_bytes // cfg.num_mma_ctas), + ) + smem_kv_pipeline_cfg = None + else: + # SmemKV: TmaUmma, LoadTma -> QK/PV MMA. + smem_kv_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.load_kv_stage, + num_bytes=cfg.tma_kc_subtile_bytes, + producer_group=tma_producer_group, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + num_bytes_per_warp_per_cta=(cfg.tma_kc_subtile_bytes // cfg.num_mma_ctas), + ) + smem_k_pipeline_cfg = None + smem_v_pipeline_cfg = None + + # TmemS: UmmaAsync, 2 stages, MmaTask -> SoftmaxTask + tmem_s_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=cfg.mma_s_stage, + producer_group=umma_hw_group, + consumer_group=compute_group_cluster, + cta_layout_vmnk=cluster_shape_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + interleave_stride=(1, 1, 2, 2) if cfg.use_fp8_dual_softmax_schedule else 1, + ) + + # SmemP: AsyncUmma, 2 stages, SoftmaxTask -> MmaTask + smem_p_pipeline_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=cfg.p_mma_stage, + producer_group=compute_group_cluster, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + interleave_stride=(2, 2, 1, 1) if cfg.use_fp8_dual_softmax_schedule else 1, + ) + + # TmemCorr: AsyncAsync, 2 stages, Softmax -> Correction + tmem_corr_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=cfg.p_cor_stage, + producer_group=compute_group_local, + consumer_group=correction_group_local, + cta_layout_vmnk=cluster_shape_vmnk, + interleave_stride=(2, 2, 1, 1) if cfg.use_fp8_dual_softmax_schedule else 1, + ) + + # TmemO: UmmaAsync, 1 stage, Mma -> Correction + tmem_o_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=cfg.mma_o_stage, + producer_group=umma_hw_group, + consumer_group=correction_group_cluster, + cta_layout_vmnk=cluster_shape_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + ) + + # OInit removed: TMEM visibility is ensured by kernel-level named + # barrier sync before task_manager.run(), so no pipeline needed. + + work_throttle = None + if use_work_throttle: + work_throttle = WorkThrottleBarrierResource( + pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( + num_stages=2, + producer_group=pipeline.CooperativeGroup(Agent.Thread, WARP_SIZE), + consumer_group=pipeline.CooperativeGroup(Agent.Thread, WARP_SIZE), + cta_layout_vmnk=cluster_shape_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ), + name="work_throttle", + ) + + # ────────────────────────────────────────────────────────────── + # Create resource instances + # ────────────────────────────────────────────────────────────── + + page_offset_window = PageOffsetWindowResource( + page_offsets=page_offsets, + cfg=cfg, + pipeline_config=None, + name="page_offset_window", + ) + + smem_q = SmemQResource( + smem_q_latent=smem_q_latent_arr, + smem_q_rope=smem_q_rope_arr, + tma_desc_q_latent=tma_desc_q_latent, + tma_desc_q_rope=tma_desc_q_rope, + cu_seqlens_q=cu_seqlens_q, + logical_num_heads_q=logical_num_heads_q, + logical_seq_len_q=logical_seq_len_q, + cfg=cfg, + pipeline_config=smem_q_pipeline_cfg, + name="smem_q", + ) + + if cfg.use_fp8_split_mma_schedule: + smem_k = SmemKResource( + smem_k=smem_kc_arr, + page_offsets=page_offsets, + tma_desc_c_latent=tma_desc_c_latent, + tma_desc_c_rope=tma_desc_c_rope, + logical_seq_len_q=logical_seq_len_q, + cfg=cfg, + pipeline_config=smem_k_pipeline_cfg, + name="smem_k", + ) + smem_v = SmemVResource( + smem_v=smem_vc_arr, + page_offsets=page_offsets, + tma_desc_c_transpose=tma_desc_c_transpose, + logical_seq_len_q=logical_seq_len_q, + cfg=cfg, + pipeline_config=smem_v_pipeline_cfg, + name="smem_v", + ) + smem_kv = None + else: + smem_kv = SmemKVResource( + smem_kv=smem_kc_arr, + tma_desc_c_latent=tma_desc_c_latent, + tma_desc_c_rope=tma_desc_c_rope, + tma_desc_c_transpose=tma_desc_c_transpose, + cfg=cfg, + pipeline_config=smem_kv_pipeline_cfg, + name="smem_kv", + ) + smem_k = None + smem_v = None + + tmem_s = TmemSResource( + smem_q_latent=smem_q_latent_arr, + smem_q_rope=smem_q_rope_arr, + smem_p=smem_p_arr, + smem_exchange=None, # Set at runtime + softmax_scale_log2=None, # Set at runtime + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + split_kv=split_kv, + logical_num_heads_q=logical_num_heads_q, + logical_seq_len_q=logical_seq_len_q, + tiled_mma_qk=tiled_mma_qk, + cfg=cfg, + pipeline_config=tmem_s_pipeline_cfg, + name="tmem_s", + ) + + smem_p = SmemPResource( + smem_p=smem_p_arr, + cfg=cfg, + pipeline_config=smem_p_pipeline_cfg, + name="smem_p", + ) + + tmem_corr = TmemCorrResource( + cfg=cfg, + pipeline_config=tmem_corr_pipeline_cfg, + name="tmem_corr", + ) + + tmem_o = TmemOResource( + cfg=cfg, + tmem_corr_ref=tmem_corr, + pipeline_config=tmem_o_pipeline_cfg, + name="tmem_o", + ) + + gmem_o = GmemOResource( + cfg=cfg, + output=output, + partial_output=acc_output, + lse=lse, + partial_lse=acc_lse, + tmem_o_ref=tmem_o, + tmem_corr_ref=tmem_corr, + output_scale=None, # set at runtime + softmax_scale_log2=None, # set at runtime + smem_exchange=None, # set at runtime + split_kv=None, # set at runtime + cu_seqlens_q=cu_seqlens_q, + logical_num_heads_q=logical_num_heads_q, + logical_seq_len_q=logical_seq_len_q, + name="gmem_o", + ) + + # ────────────────────────────────────────────────────────────── + # Create tasks + # ────────────────────────────────────────────────────────────── + + # Warpgroup 2 tasks (warps 8-11) use the low-register producer budget. + # Keep this explicit for both validation and codegen so the scheduler's + # register-budget check accounts for the real MLA warpgroup layout. + wg2_reg_count = cfg.other_reg_num + + if cfg.use_fp8_split_mma_schedule: + load_k_task = create_load_k_task( + smem_q, + smem_k, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + load_v_task = create_load_v_task( + smem_v, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + mma_task = None + mma_qk_task = create_mma_qk_direct_task( + smem_q, + smem_k, + tmem_s, + iterations_qk=cfg.iterations_qk, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + mma_pv_task = create_mma_pv_direct_task( + smem_v, + smem_p, + tmem_o, + iterations_pv=cfg.iterations_pv_k * cfg.iterations_pv_n, + iterations_pv_k=cfg.iterations_pv_k, + iterations_pv_n=cfg.iterations_pv_n, + per_n_o_pipeline=cfg.use_fp8_split_mma_schedule, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + load_tma_task = None + else: + load_k_task = None + load_v_task = None + load_tma_task = create_load_tma_task( + page_offset_window, + smem_q, + smem_kv, + iterations_qk=cfg.iterations_qk_stages, + iterations_pv=cfg.iterations_pv_stages, + work_queue=work_queue, + task_class=non_interleaved_task_class, + domain=task_domain, + num_registers=wg2_reg_count, + ) + mma_task = create_mma_task( + smem_q, + smem_kv, + smem_p, + tmem_s, + tmem_o, + iterations_qk=cfg.iterations_qk_stages, + iterations_pv=cfg.iterations_pv_stages, + work_queue=work_queue, + work_throttle=work_throttle, + task_class=non_interleaved_task_class, + domain=task_domain, + num_registers=wg2_reg_count, + ) + mma_qk_task = None + mma_pv_task = None + + # Warpgroup 0 (warps 0-3): all 4 warps are in SoftmaxTask, + # so setmaxnreg.inc 192 is safe (all warps participate). + softmax_task = create_softmax_task( + tmem_s, + tmem_corr, + smem_p, + work_queue=work_queue, + task_class=( + MlaInterleavedTask + if cfg.use_fp8_dual_softmax_schedule + else non_interleaved_task_class + ), + domain=(domain if cfg.use_fp8_dual_softmax_schedule else task_domain), + domain_start=0, + step=2 if cfg.use_fp8_dual_softmax_schedule else 1, + num_registers=cfg.softmax_reg_num, + softmax_group_id=0, + ) + if cfg.use_fp8_dual_softmax_schedule: + second_softmax_task = create_softmax_task( + tmem_s, + tmem_corr, + smem_p, + work_queue=work_queue, + task_class=MlaInterleavedTask, + domain=domain, + domain_start=1, + step=2, + num_registers=cfg.softmax_reg_num, + warp_idx=cfg.second_compute_warp_ids[0], + name="SoftmaxOddTask", + softmax_group_id=1, + ) + else: + second_softmax_task = None + + # Warpgroup 1 (warps 4-7): all 4 warps are in CorrectionTask, + # so setmaxnreg.inc 208 is safe (all warps participate). + correction_task = create_correction_task( + tmem_corr, + tmem_o, + gmem_o, + iterations_pv_n=cfg.iterations_pv_n, + per_n_o_pipeline=cfg.use_fp8_split_mma_schedule, + work_queue=work_queue, + task_class=non_interleaved_task_class, + domain=task_domain, + num_registers=cfg.correction_reg_num, + ) + + if cfg.use_fp8_split_mma_schedule: + task_list = [ + load_k_task, + load_v_task, + ] + task_list.extend([mma_qk_task, softmax_task]) + if second_softmax_task is not None: + task_list.append(second_softmax_task) + task_list.extend([mma_pv_task, correction_task]) + else: + task_list = [ + load_tma_task, + ] + task_list.extend([mma_task, softmax_task, correction_task]) + + if not cfg.use_fp8_split_mma_schedule and use_clc_dynamic: + padding_task = create_padding_task( + work_queue=work_queue, + task_class=non_interleaved_task_class, + domain=task_domain, + num_registers=wg2_reg_count, + warp_idx=10, + ) + task_list.append(padding_task) + scheduler_task = create_scheduler_task( + work_queue, + work_throttle, + task_class=non_interleaved_task_class, + num_registers=wg2_reg_count, + ) + task_list.append(scheduler_task) + elif not cfg.use_fp8_split_mma_schedule: + # BF16 uses only three one-warp producer tasks in warpgroup 2. Keep an + # explicit warp-11 placeholder so register validation sees the complete + # four-warp group with the low-register producer budget. + padding_task = create_padding_task( + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + warp_idx=10, + num_warps=2, + ) + task_list.append(padding_task) + + # ────────────────────────────────────────────────────────────── + # Dependency graph + # ────────────────────────────────────────────────────────────── + if cfg.use_fp8_split_mma_schedule: + resource_dependency_graph = { + smem_q: [], # Q loads (independent) + smem_k: [], # K loads read page offsets directly + smem_v: [], # V loads read page offsets directly + tmem_s: [smem_k, smem_q], # QK MMA needs K and Q + smem_p: [tmem_s], # softmax reads S -> writes P + tmem_corr: [tmem_s], # softmax produces correction + tmem_o: [smem_p, smem_v], # PV MMA needs P and V + gmem_o: [tmem_o, tmem_corr], # epilogue needs O + correction + } + dma_consumer_release_labels = { + (smem_k, tmem_s): {"k_desc"}, + (smem_v, tmem_o): {"v_desc", "v_desc_n_major"}, + } + else: + resource_dependency_graph = { + page_offset_window: [], # register-held page-table window + smem_q: [], # Q loads (independent) + # Page offsets are cached into registers inside LoadTmaTask before K/V TMA. + smem_kv: [], + tmem_s: [smem_kv, smem_q], # QK MMA needs K and Q + smem_p: [tmem_s], # softmax reads S -> writes P + tmem_corr: [tmem_s], # softmax produces correction + tmem_o: [smem_p, smem_kv], # PV MMA needs P and V + gmem_o: [tmem_o, tmem_corr], # epilogue needs O + correction + } + dma_consumer_release_labels = { + (smem_kv, tmem_s): {"k_desc"}, + (smem_kv, tmem_o): {"v_desc"}, + } + if work_queue is not None: + if use_clc_dynamic: + for resource, dependencies in tuple(resource_dependency_graph.items()): + if ( + resource is not work_queue + and resource is not page_offset_window + and work_queue not in dependencies + ): + dependencies.append(work_queue) + resource_dependency_graph[work_queue] = ( + [work_queue, work_throttle] + if work_throttle is not None + else [work_queue] + ) + if work_throttle is not None: + # The leader MMA produces both S and the throttle token after + # it observes Q for the current work tile. + resource_dependency_graph[work_throttle] = [tmem_s] + else: + resource_dependency_graph[work_queue] = [] + + # ────────────────────────────────────────────────────────────── + # Create TaskManager + # ────────────────────────────────────────────────────────────── + skip = not isinstance(domain, int) + task_manager = TaskManager( + tasks=task_list, + resource_dependency_graph=resource_dependency_graph, + dma_consumer_release_labels=dma_consumer_release_labels, + skip_validation=skip, + verbose=verbose, + exhaustive_deadlock_race_check=exhaustive_deadlock_race_check, + ) + + tmem_resources = [tmem_s, tmem_o, tmem_corr] + named_resources = { + "tmem_s": tmem_s, + "tmem_o": tmem_o, + "tmem_corr": tmem_corr, + "gmem_o": gmem_o, + } + return task_manager, tmem_resources, named_resources + + +# GPU Kernel Class +# ===================================================================== + + +class MlaDecodeTs: + """Warp-specialised MLA decode kernel using the TS framework. + + Usage:: + + mla = MlaDecodeTs() + mla(q_latent, q_rope, c_latent, c_rope, page_offsets, + o, lse, workspace, split_kv, cache_seqs, cu_seqlens_q, + block_split_kvs, softmax_scale, output_scale, stream) + """ + + def __init__( + self, + acc_dtype=None, + lse_dtype=None, + mma_qk_tiler_mn=(128, 128), + mma_pv_tiler_mn=(128, 256), + max_active_clusters=56, + page_size=32, + is_persistent=True, + is_var_seq=False, + is_var_split_kv=False, + static_split_kv=None, + static_seq_len_k=None, + qkv_dtype="bf16", + out_dtype="bf16", + rope_dim=64, + num_heads=128, + seq_len_q=1, + batch_size=1, + mask_type: MaskType | str = MaskType.CAUSAL, + ): + """ + Parameters + ---------- + acc_dtype : cutlass dtype, optional + Accumulator dtype (default: Float32). + lse_dtype : cutlass dtype, optional + LSE dtype (default: Float32). + mma_qk_tiler_mn : tuple, optional + MMA tiler shape (M, N) for QK gemm (default: (128, 128)). + mma_pv_tiler_mn : tuple, optional + MMA tiler shape (M, N) for PV gemm (default: (128, 256)). + max_active_clusters : int, optional + Maximum number of active clusters (default: 56). + page_size : int, optional + KV cache page size in tokens (default: 32). + is_persistent : bool, optional + Use persistent kernel scheduling (default: True). + is_var_seq : bool, optional + Enable variable KV-cache sequence lengths (default: False). + is_var_split_kv : bool, optional + Use ``block_split_kvs[batch]`` as a per-batch split cap. Device + scheduling further contracts that cap from each tile's valid K. + static_split_kv : int or None, optional + Compile-time maximum split-KV capacity. The grid and workspace use + this value, while device scheduling skips the inactive split suffix + for shorter runtime K lengths. + static_seq_len_k : int or None, optional + Compile-time K length for fixed-length launches. Used only for + non-empty split-KV-1 specialisation; variable cache_seqs launches + keep the generic runtime domain path. + qkv_dtype : str, optional + Q/K/V dtype name. + out_dtype : str, optional + Output dtype name. + rope_dim : int, optional + RoPE head dimension. + num_heads : int, optional + Logical query-head count used to derive the flat-row tile count. + seq_len_q : int, optional + Logical query length used to derive the flat-row tile count. + batch_size : int, optional + Host-known batch size used to qualify the standalone reducer + topology. + mask_type : MaskType or str, optional + ``causal`` (default) for bottom-right speculative decoding or + ``dense`` for full per-batch KV visibility. + """ + import cutlass as _cutlass + + if acc_dtype is None: + acc_dtype = _cutlass.Float32 + if lse_dtype is None: + lse_dtype = _cutlass.Float32 + + self.acc_dtype = acc_dtype + self.lse_dtype = lse_dtype + self.mma_qk_tiler_mn = mma_qk_tiler_mn + self.mma_pv_tiler_mn = mma_pv_tiler_mn + self.max_active_clusters = max_active_clusters + self.page_size = page_size + self.is_persistent = is_persistent + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv + self.static_split_kv = static_split_kv + self.reduction_split_capacity = ( + static_split_kv if static_split_kv is not None else MAX_MLA_SPLITS_KV + ) + if not 1 <= self.reduction_split_capacity <= MAX_MLA_SPLITS_KV: + raise ValueError(f"static_split_kv must be in [1, {MAX_MLA_SPLITS_KV}]") + self.static_seq_len_k = static_seq_len_k + self.qkv_dtype = qkv_dtype + self.out_dtype = out_dtype + self.rope_dim = rope_dim + self.num_heads = num_heads + self.seq_len_q = seq_len_q + if isinstance(batch_size, bool) or not isinstance(batch_size, int): + raise TypeError("batch_size must be an integer") + if batch_size <= 0: + raise ValueError("batch_size must be positive") + self.batch_size = batch_size + self.mask_type = normalize_mask_type(mask_type) + self.query_tile_layout = FlatQueryTileLayout.for_tile( + num_heads, seq_len_q, mma_qk_tiler_mn[0] + ) + self.num_q_tiles = self.query_tile_layout.num_tiles + self.tail_q_rows = self.query_tile_layout.tail_rows + self.parallel_reduction_topology: ParallelReductionTopology | None = None + self.use_parallel_reduction = False + self._parallel_reduction_shape_is_eligible = ( + not is_var_split_kv + and static_split_kv is not None + and 2 <= static_split_kv <= 128 + and mma_qk_tiler_mn[0] == 128 + and acc_dtype == _cutlass.Float32 + and lse_dtype == _cutlass.Float32 + and not is_persistent + ) + self._configure_parallel_reduction_topology() + + def _effective_reduction_shape(self) -> tuple[int, int]: + """Return physical row/tile extents used by the split workspace.""" + + return self.mma_qk_tiler_mn[0], self.num_q_tiles + + def compile_signature(self) -> tuple[object, ...]: + """Return the complete batch-independent JIT identity.""" + + return ( + self.acc_dtype, + self.lse_dtype, + self.mma_qk_tiler_mn, + self.mma_pv_tiler_mn, + self.max_active_clusters, + self.page_size, + self.is_persistent, + self.is_var_seq, + self.is_var_split_kv, + self.static_split_kv, + self.static_seq_len_k, + self.qkv_dtype, + self.out_dtype, + self.rope_dim, + self.num_heads, + self.seq_len_q, + self.mask_type, + self.reduction_split_capacity, + self.query_tile_layout, + self.num_q_tiles, + self.tail_q_rows, + self._parallel_reduction_shape_is_eligible, + self.use_parallel_reduction, + self.parallel_reduction_topology, + ) + + def _configure_parallel_reduction_topology(self) -> None: + """Refresh reducer topology after any host-side launch-shape update.""" + + self.parallel_reduction_topology = None + self.use_parallel_reduction = False + physical_tile_rows, num_query_tiles = self._effective_reduction_shape() + + # Every reducer shares this normalized partial workspace. Validate its + # full configured capacity, not only high-split clustered launches. + if self.reduction_split_capacity > 1: + validate_parallel_reduction_workspace( + batch_size=self.batch_size, + num_heads_q=physical_tile_rows, + seq_len_q=num_query_tiles, + splits_kv=self.reduction_split_capacity, + head_dim=PARALLEL_REDUCTION_HEAD_DIM, + ) + + if not self._parallel_reduction_shape_is_eligible: + return + assert self.static_split_kv is not None + topology = make_q128_wave_limited_parallel_reduction_topology( + self.static_split_kv, + logical_rows=(self.num_heads * self.seq_len_q * self.batch_size), + physical_sm_count=self.max_active_clusters * 2, + max_cluster_size=8, + ) + parallel_g1_grid_has_no_padded_rows = ( + self.query_tile_layout.total_rows == physical_tile_rows * num_query_tiles + ) + # Small split counts use G1 only when row coarsening leaves a sub-wave + # grid and the producer generated enough work to amortize one CTA per + # physical row. Padded M128 tails and intermediate split counts retain + # the compact reference grid; high split counts may use clusters. + use_small_split_g1 = ( + self.static_split_kv <= 16 + and topology is not None + and topology.cluster_size == 1 + and parallel_g1_grid_has_no_padded_rows + and should_use_q128_g1_parallel_reducer( + batch_size=self.batch_size, + physical_rows_per_batch=(physical_tile_rows * num_query_tiles), + producer_ctas=( + self.batch_size * num_query_tiles * self.static_split_kv * 2 + ), + reference_rows_per_cta=REDUCTION_ROWS_PER_CTA, + physical_sm_count=self.max_active_clusters * 2, + ) + ) + use_high_split_cluster = ( + self.static_split_kv > 32 + and topology is not None + and topology.cluster_size > 1 + ) + if use_small_split_g1 or use_high_split_cluster: + self.parallel_reduction_topology = topology + self.use_parallel_reduction = True + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_offsets: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor | None, + block_split_kvs: cute.Tensor, + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: object, + ): + """Execute the MLA decode TS kernel.""" + cfg = make_mla_decode_config( + mma_qk_tiler_mn=self.mma_qk_tiler_mn, + mma_pv_tiler_mn=self.mma_pv_tiler_mn, + rope_dim=self.rope_dim, + page_size=self.page_size, + qkv_dtype=self.qkv_dtype, + o_dtype=self.out_dtype, + max_active_clusters=self.max_active_clusters, + is_persistent=self.is_persistent, + is_var_seq=self.is_var_seq, + is_var_split_kv=self.is_var_split_kv, + mask_type=self.mask_type, + ) + physical_tile_rows = self.mma_qk_tiler_mn[0] + num_query_tiles = self.num_q_tiles + # Fixed public tensors retain [H,D,SQ,B]/[H,SQ,B]; variable-Q tensors + # compact batches into [H,D,totalQ]/[H,totalQ]. The scheduler/workspace + # use physical flat-query tile coordinates, while resources map + # each valid row back to its logical fixed or ragged storage location. + if cutlass.const_expr(cu_seqlens_q is not None): + runtime_assert( + cute.size(cu_seqlens_q) == cute.size(cache_seqs) + Int32(1), + "cu_seqlens_q must contain one more offset than cache_seqs", + ) + batch_size = cute.size(cu_seqlens_q) - Int32(1) + else: + batch_size = cute.size(o.shape[3]) + runtime_assert( + batch_size == cute.size(cache_seqs), + "fixed output batch size must match cache_seqs", + ) + + runtime_assert( + q_latent.stride[2] == q_latent.shape[0] * q_latent.stride[0], + "q_latent must be compact across the head and query dimensions", + ) + runtime_assert( + q_rope.stride[2] == q_rope.shape[0] * q_rope.stride[0], + "q_rope must be compact across the head and query dimensions", + ) + runtime_assert( + o.stride[1] == 1, + "o must have a contiguous dimension axis", + ) + runtime_assert( + o.stride[0] == o.shape[1] * o.stride[1], + "o must be compact from the dimension axis into the head axis", + ) + runtime_assert( + o.stride[2] == o.shape[0] * o.stride[0], + "o must be compact from the head axis into the query axis", + ) + if cutlass.const_expr(cu_seqlens_q is None): + runtime_assert( + o.stride[3] == o.shape[2] * o.stride[2], + "o must be compact from the query axis into the batch axis", + ) + runtime_assert( + lse.stride[0] == 1, + "lse must have a contiguous head axis", + ) + runtime_assert( + lse.stride[1] == lse.shape[0] * lse.stride[0], + "lse must be compact from the head axis into the query axis", + ) + if cutlass.const_expr(cu_seqlens_q is None): + runtime_assert( + lse.stride[2] == lse.shape[1] * lse.stride[1], + "lse must be compact from the query axis into the batch axis", + ) + + def _flatten_query_rows_for_tma(t): + """Expose logical ``(SQ, H)`` as one bounded TMA row dimension.""" + if cutlass.const_expr(cu_seqlens_q is not None): + return cute.make_tensor( + t.iterator, + cute.make_layout( + (t.shape[1], t.shape[0] * t.shape[2]), + stride=(t.stride[1], t.stride[0]), + ), + ) + return cute.make_tensor( + t.iterator, + cute.make_layout( + (t.shape[1], t.shape[0] * t.shape[2], t.shape[3]), + stride=(t.stride[1], t.stride[0], t.stride[3]), + ), + ) + + # Create TMA descriptors (same as bare metal) + + # Keep the descriptor extent at the logical H*SQ row count. The final + # physical query tile still requests a full M tile; tensor-map OOB fill + # supplies zeros for its tail without changing the public Q shape. + q_latent_tma = _flatten_query_rows_for_tma(q_latent) + if cutlass.const_expr(cu_seqlens_q is not None): + tma_desc_q_latent = create_tensor_map_ragged_from_tensor( + q_latent_tma, + box_dims=(cfg.mma_qk_tiler[2], cfg.mma_qk_tiler[0] // 2), + ragged_dim=1, + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_q_latent = create_tensor_map_tiled_from_view( + q_latent_tma, + box_dims=(cfg.mma_qk_tiler[2], cfg.mma_qk_tiler[0] // 2, 1), + stride_order=(0, 1, 2), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + if cutlass.const_expr(cfg.rope_dim > 0): + q_rope_tma = _flatten_query_rows_for_tma(q_rope) + q_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + q_rope_swizzle = cuda.TensorMapSwizzle.s64b + if cutlass.const_expr(cu_seqlens_q is not None): + tma_desc_q_rope = create_tensor_map_ragged_from_tensor( + q_rope_tma, + box_dims=( + cfg.mma_qk_rope_tiler[2], + cfg.mma_qk_rope_tiler[0] // 2, + ), + ragged_dim=1, + stride_order=(0, 1), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_q_rope = create_tensor_map_tiled_from_view( + q_rope_tma, + box_dims=( + cfg.mma_qk_rope_tiler[2], + cfg.mma_qk_rope_tiler[0] // 2, + 1, + ), + stride_order=(0, 1, 2), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_q_rope = tma_desc_q_latent + + c_latent_tma = cute.make_tensor( + c_latent.iterator, + cute.select(c_latent.layout, mode=[1, 0, 2]), + ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_latent_tma, + box_dims=(cfg.mma_qk_tiler[2], cfg.kc_page_tile_size, 1), + stride_order=(0, 1, 2), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + if cutlass.const_expr(cfg.rope_dim > 0): + c_rope_tma = cute.make_tensor( + c_rope.iterator, + cute.select(c_rope.layout, mode=[1, 0, 2]), + ) + c_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + c_rope_swizzle = cuda.TensorMapSwizzle.s64b + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_rope_tma, + box_dims=(cfg.mma_qk_rope_tiler[2], cfg.kc_page_tile_size, 1), + stride_order=(0, 1, 2), + swizzle=c_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_c_rope = tma_desc_c_latent + + c_transpose_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.mma_pv_tiler[2] == 64): + c_transpose_swizzle = cuda.TensorMapSwizzle.s64b + c_latent_transpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor( + c_latent.iterator, c_latent_transpose_layout + ) + # The physical page controls only the GMEM coordinate. TMA assembles + # fixed K32 SMEM blocks for PV from one or more page-bounded copies. + tma_desc_c_transpose = create_tensor_map_tiled_from_view( + c_latent_transpose, + box_dims=(V_TMA_LATENT_ELEMENTS, cfg.v_tma_token_count, 1), + stride_order=(0, 1, 2), + swizzle=c_transpose_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + softmax_scale_log2 = softmax_scale * LOG2_E + + kernel_split_kv = ( + Int32(self.static_split_kv) + if cutlass.const_expr(self.static_split_kv is not None) + else split_kv + ) + + # Compute grid + tile_sched_params = create_mla_static_tile_scheduler_params( + self.is_persistent, + batch_size, + Int32(num_query_tiles), + cfg.cluster_shape_mnk, + kernel_split_kv, + ) + use_clc_dynamic = self.is_persistent and not cfg.is_fp8_qkv() + clc_tile_sched_params = None + if cutlass.const_expr(use_clc_dynamic): + # Keep the physical query-tile dimension in grid X and flatten + # only split/batch into grid Z. Besides avoiding a hot-path + # S/B decode for every stolen tile, this preserves the natural + # 2CTA query-cluster raster used by the nonpersistent launch. + clc_tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + ( + cfg.cluster_shape_mnk[0] * Int32(num_query_tiles), + 1, + batch_size * kernel_split_kv, + ), + cfg.cluster_shape_mnk, + ) + grid = clc_tile_sched_params.get_grid_shape() + else: + grid = MLAStaticTileScheduler.get_grid_shape( + tile_sched_params, self.max_active_clusters + ) + + # Initialize workspace for split_kv > 1 + acc_o, acc_lse = self.initialize_workspace( + Int32(physical_tile_rows), + cfg.latent_dim, # D + Int32(num_query_tiles), + batch_size, + kernel_split_kv, + workspace, + ) + # A one-wave producer can publish its dependent reducer launch while + # retiring, hiding launch latency without admitting reducer CTAs into + # an actively grid-striding persistent producer. + use_one_wave_reducer_pdl = acc_o is not None and not self.is_persistent + + self.split_kv_kernel( + tma_desc_q_latent, + tma_desc_q_rope, + tma_desc_c_latent, + tma_desc_c_rope, + tma_desc_c_transpose, + c_latent, + c_rope, + page_offsets, + o, + lse, + acc_o, + acc_lse, + kernel_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + softmax_scale_log2, + output_scale, + tile_sched_params, + clc_tile_sched_params, + ).launch( + grid=grid, + block=[cfg.threads_per_cta, 1, 1], + cluster=cfg.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + use_pdl=use_one_wave_reducer_pdl, + ) + + # Reduction kernel: combine per-split results when split_kv > 1 + if cutlass.const_expr(acc_o is not None): + if cutlass.const_expr(self.use_parallel_reduction): + topology = self.parallel_reduction_topology + self.parallel_reduction_kernel( + o, + lse, + acc_o, + acc_lse, + kernel_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + ).launch( + grid=( + physical_tile_rows * topology.cluster_size, + num_query_tiles, + batch_size, + ), + block=[PARALLEL_REDUCTION_THREADS, 1, 1], + cluster=[topology.cluster_size, 1, 1], + stream=stream, + min_blocks_per_mp=1, + use_pdl=use_one_wave_reducer_pdl, + ) + else: + logical_query_rows = self.num_heads * self.seq_len_q + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + kernel_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + ).launch( + grid=( + ceil_div(logical_query_rows, REDUCTION_ROWS_PER_CTA), + 1, + batch_size, + ), + block=[REDUCTION_ROWS_PER_CTA * REDUCTION_THREADS_PER_ROW, 1, 1], + smem=( + REDUCTION_ROWS_PER_CTA + * self.reduction_split_capacity + * self.lse_dtype.width + // 8 + ), + stream=stream, + min_blocks_per_mp=2, + use_pdl=use_one_wave_reducer_pdl, + ) + + @cute.kernel + def split_kv_kernel( + self, + tma_desc_q_latent: cutlass.GridConstant[cuda.TensorMap], + tma_desc_q_rope: cutlass.GridConstant[cuda.TensorMap], + tma_desc_c_latent: cutlass.GridConstant[cuda.TensorMap], + tma_desc_c_rope: cutlass.GridConstant[cuda.TensorMap], + tma_desc_c_transpose: cutlass.GridConstant[cuda.TensorMap], + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_offsets: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + acc_o: cute.Tensor, + acc_lse: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor | None, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + tile_sched_params: MLAStaticTileSchedulerParams, + clc_tile_sched_params: object, + ) -> None: + """MLA decode TS kernel: persistent tile-scheduled execution.""" + cfg = make_mla_decode_config( + mma_qk_tiler_mn=self.mma_qk_tiler_mn, + mma_pv_tiler_mn=self.mma_pv_tiler_mn, + rope_dim=self.rope_dim, + page_size=self.page_size, + qkv_dtype=self.qkv_dtype, + o_dtype=self.out_dtype, + max_active_clusters=self.max_active_clusters, + is_persistent=self.is_persistent, + is_var_seq=self.is_var_seq, + is_var_split_kv=self.is_var_split_kv, + mask_type=self.mask_type, + ) + num_query_tiles = self.num_q_tiles + use_clc_dynamic = self.is_persistent and not cfg.is_fp8_qkv() + tiled_mma_qk = None + if cutlass.const_expr(cfg.is_fp8_qkv()): + tiled_mma_qk = sm100_utils.make_trivial_tiled_mma( + qkv_dtype(cfg), + qkv_dtype(cfg), + OperandMajorMode.K, + OperandMajorMode.K, + self.acc_dtype, + tcgen05.CtaGroup.TWO, + cfg.mma_qk_tiler[:2], + ) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + cluster_idx, _, _ = cute.arch.block_idx() + + mma_tile_coord_v = cluster_idx % 2 + + # Prefetch TMA descriptors on MMA warp + if warp_idx == cfg.mma_warp_id: + prims.prefetch_tensormap(tma_desc_q_latent.get_ptr()) + prims.prefetch_tensormap(tma_desc_q_rope.get_ptr()) + prims.prefetch_tensormap(tma_desc_c_latent.get_ptr()) + prims.prefetch_tensormap(tma_desc_c_rope.get_ptr()) + prims.prefetch_tensormap(tma_desc_c_transpose.get_ptr()) + + # Allocate SMEM + qkv_element_dtype = qkv_dtype(cfg) + smem_q_latent_arr = cutlass.Array( + qkv_element_dtype, + cfg.smem_q_latent_elems, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + smem_q_rope_arr = cutlass.Array( + qkv_element_dtype, + max(1, cfg.smem_q_rope_elems), + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + smem_kc_arr = cutlass.Array( + qkv_element_dtype, + cfg.smem_kc_elems, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + smem_vc_arr = None + if cutlass.const_expr(cfg.use_fp8_split_mma_schedule): + smem_vc_arr = cutlass.Array( + qkv_element_dtype, + cfg.smem_vc_elems, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + smem_p_arr = cutlass.Array( + qkv_element_dtype, + cfg.smem_p_elems, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + softmax_exchange_arr = cutlass.Array( + self.acc_dtype, + cfg.softmax_exchange_elems, + space=cutlass.AddressSpace.smem, + alignment=4, + ) + epilogue_exchange_arr = cutlass.Array( + self.acc_dtype, + cfg.num_compute_warps * cfg.threads_per_warp, + space=cutlass.AddressSpace.smem, + alignment=4, + ) + tmem_holding_buf_arr = cutlass.Array( + Int32, 1, space=cutlass.AddressSpace.smem, alignment=4 + ) + tmem_dealloc_mbar_arr = cutlass.Array( + Int64, 1, space=cutlass.AddressSpace.smem, alignment=8 + ) + clc_response_ptr = None + if cutlass.const_expr(use_clc_dynamic): + # Keep both response stages in the ordinary dynamic-SMEM arena. + # ``alloc_smem`` creates a separately rounded static section, which + # needlessly exceeds this kernel's near-capacity SMEM budget. + clc_response_arr = cutlass.Array( + cutlass.Int128, + 2, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + clc_response_ptr = cute.make_ptr( + cutlass.Int128, + clc_response_arr.data_ptr(), + mem_space=cutlass.AddressSpace.smem, + ) + + # Init dealloc mbarrier + if warp_idx == cfg.mma_warp_id: + if prims.elect_sync(): + prims.mbarrier_init(tmem_dealloc_mbar_arr, TMEM_DEALLOC_MBAR_THREADS) + + # setmaxnreg.dec: ALL warps in warpgroup 2 (warps 8-11) + if warp_idx >= 8: + prims.setmaxregister(cfg.other_reg_num, prims.SetMaxRegisterAction.DECREASE) + + # Workaround: avoid CSE-induced register spills. The tile + # scheduler calls grid_dim(), producing nctaid SSA values in the shared + # prologue; CSE merges them across tasks, preventing per-task tile + # schedulers from getting independent copies and causing register spills. + # Remove once the compiler scopes nctaid reads per task. + # Tile decomposition — compute blk_coord from ctaid.x WITHOUT creating + # a full tile scheduler. + if cutlass.const_expr(use_clc_dynamic): + query_cluster_idx, _, split_batch_idx = cute.arch.block_idx() + seq_q_idx = query_cluster_idx // Int32(cfg.cluster_shape_mnk[0]) + cluster_idx = query_cluster_idx % Int32(cfg.cluster_shape_mnk[0]) + split_kv_idx, batch_idx = divmod_constexpr_power_of_two_or_fdd( + split_batch_idx, + None, + tile_sched_params.problem_shape_b_fdd, + ) + blk_coord = (cluster_idx, seq_q_idx, batch_idx, split_kv_idx) + elif cutlass.const_expr(self.is_persistent): + current_work_linear_idx = cute.arch.block_idx()[0] + current_work_cluster_batch = current_work_linear_idx // Int32( + cfg.cluster_shape_mnk[0] + ) + cluster_idx = current_work_linear_idx % Int32(cfg.cluster_shape_mnk[0]) + current_work_after_seq_q, seq_q_idx = divmod_constexpr_power_of_two_or_fdd( + current_work_cluster_batch, + num_query_tiles, + tile_sched_params.problem_shape_s_fdd, + ) + current_work_after_batch, batch_idx = divmod_constexpr_power_of_two_or_fdd( + current_work_after_seq_q, + None, + tile_sched_params.problem_shape_b_fdd, + ) + _, split_kv_idx = divmod( + current_work_after_batch, tile_sched_params.split_kv_fdd + ) + blk_coord = (cluster_idx, seq_q_idx, batch_idx, split_kv_idx) + else: + cluster_idx, seq_batch_idx, split_kv_idx = cute.arch.block_idx() + seq_q_idx, batch_idx = divmod_constexpr_power_of_two_or_fdd( + seq_batch_idx, + None, + tile_sched_params.problem_shape_b_fdd, + ) + blk_coord = (cluster_idx, seq_q_idx, batch_idx, split_kv_idx) + tile_cluster_idx, tile_seq_q_idx, tile_batch_idx, tile_split_kv_idx = blk_coord + del tile_cluster_idx + + fixed_nonempty_single_split = ( + not self.is_var_seq + and not self.is_var_split_kv + and self.static_split_kv == 1 + and cu_seqlens_q is None + ) + if cutlass.const_expr(fixed_nonempty_single_split): + max_split_kv = Int32(1) + exit_early = False + elif cutlass.const_expr(not self.is_persistent): + # Every task already derives its graph-live K/Q domain through the + # work queue. Let a zero-domain task skip its data schedule instead + # of recomputing the same metadata in an all-warp CTA prologue. + max_split_kv = ( + Int32(self.static_split_kv) + if cutlass.const_expr(self.static_split_kv is not None) + else split_kv + ) + exit_early = False + else: + # Compute the initial tile's active Q/split state for static early + # exit. Persistent CTAs stay live because a later grid-stride tile + # can be active even when their initial logical tile is padded. + # The full k_tile_count is recomputed per-task inside + # MlaTask._run_task_body_persistent to keep SSA live ranges short. + if cutlass.const_expr(self.static_seq_len_k is not None): + K = Int32(self.static_seq_len_k) + else: + K = cache_seqs[tile_batch_idx] + + _, q_len = query_batch_bounds( + cu_seqlens_q, + tile_batch_idx, + self.seq_len_q, + ) + query_tile_has_rows = runtime_flat_query_tile_has_rows( + tile_seq_q_idx, + self.mma_qk_tiler_mn[0], + self.num_heads, + self.seq_len_q, + cu_seqlens_q, + tile_batch_idx, + ) + if cutlass.const_expr( + cfg.mask_type == MaskType.CAUSAL.value and self.seq_len_q > 1 + ): + _, _, logical_q_idx, _, _ = flat_query_row_state( + Int32(self.mma_qk_tiler_mn[0] - 1), + tile_seq_q_idx, + self.mma_qk_tiler_mn[0], + self.num_heads, + self.seq_len_q, + cu_seqlens_q, + tile_batch_idx, + ) + K = mask_visible_k_length(cfg.mask_type, K, logical_q_idx, q_len) + K = K if query_tile_has_rows else Int32(0) + # split_kv is the static launch/workspace capacity. Runtime K + # contracts the optional per-batch cap to an active prefix. + max_split_kv = ( + Int32(self.static_split_kv) + if cutlass.const_expr(self.static_split_kv is not None) + else split_kv + ) + if cutlass.const_expr( + self.static_split_kv is not None and not self.is_var_split_kv + ): + split_kv_cap = max_split_kv + else: + split_kv_cap = runtime_split_kv_cap( + max_split_kv, + self.is_var_split_kv, + block_split_kvs, + tile_batch_idx, + ) + k_tile_total = (K + cfg.mma_qk_tiler[1] - 1) // cfg.mma_qk_tiler[1] + _, k_tile_count = runtime_split_tile_range( + k_tile_total, + split_kv_cap, + tile_split_kv_idx, + ) + exit_early = k_tile_count <= Int32(0) + if cutlass.const_expr(self.is_persistent): + # A physical persistent CTA may own an empty initial split but + # later grid-stride to a nonempty batch/Q tile. Let MlaTask + # skip empty logical tiles instead of terminating the CTA. + exit_early = False + + if exit_early: + # The ptxas-generated CTA_2 TMEM lifecycle barrier is initialized + # in the shared prologue. Order that initialization across both + # CTAs before tcgen05_alloc, matching the active path below. + prims.fence_mbarrier_init() + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + + # TMEM alloc + sync even for empty CTAs to keep cluster in sync. + if warp_idx == cfg.mma_warp_id: + prims.tcgen05_alloc( + tmem_holding_buf_arr, cfg.num_tmem_cols, group="cta_2" + ) + prims.tcgen05_relinquish_alloc_permit(group="cta_2") + participates_in_tmem_sync = warp_idx <= cfg.mma_warp_id + if cutlass.const_expr(cfg.use_fp8_split_mma_schedule): + participates_in_tmem_sync = participates_in_tmem_sync or ( + warp_idx == cfg.pv_mma_warp_id + ) + if cutlass.const_expr(cfg.use_fp8_dual_softmax_schedule): + participates_in_tmem_sync = participates_in_tmem_sync or ( + warp_idx >= cfg.second_compute_warp_ids[0] + and warp_idx <= cfg.second_compute_warp_ids[-1] + ) + if participates_in_tmem_sync: + prims.barrier_cta_sync( + cfg.tmem_sync_bar_id, thread_count=cfg.tmem_sync_bar_threads + ) + else: + # Build TaskManager and initialize pipelines BEFORE tcgen05_alloc. + # Pipeline mbarrier init must happen first because ptxas inserts + # TMEM lifecycle barrier code around tcgen05_alloc that uses + # static SMEM at offset 0x40, which overlaps with the region + # that pipeline mbarrier init writes to. Without this ordering, + # the lifecycle barrier is uninitialized and crashes. + + # Create one MLA-coordinate work queue for persistent scheduling. + # Static dispatch wraps MLAStaticTileScheduler and needs no + # response pipeline. CLC dispatch installs the fetch pipeline and + # response parameters below; all participating tasks still consume + # the same cached per-tile MLA coordinate and K-domain state. + work_queue_pipeline_config = None + work_queue_tile_scheduler_config = None + if cutlass.const_expr(use_clc_dynamic): + work_queue_pipeline_config = ( + PipelineConfig.create_clc_fetch_async_pipeline_cfg( + num_stages=2, + num_bytes=16, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + # CTA1 omits both the leader-only MMA warp and the + # leader-only scheduler warp. Exclude both from the + # cluster-wide empty-barrier arrival count. + cfg.threads_per_cta * cfg.num_mma_ctas + - 2 * cfg.threads_per_warp, + ), + cta_layout_vmnk=(cfg.num_mma_ctas, 1, 1, 1), + producer_signaling_threads=SignalingThreads.CtaLeader, + consumer_signaling_threads=SignalingThreads.All, + ) + ) + work_queue_tile_scheduler_config = TileSchedulerConfig.create_clc_dynamic_persistent_tile_scheduler_params( + tile_scheduler_params=clc_tile_sched_params, + response_ptr=clc_response_ptr, + ) + + mla_work_queue = MlaWorkQueue( + tile_sched_params=tile_sched_params, + cache_seqs=cache_seqs, + split_kv=max_split_kv, + block_split_kvs=block_split_kvs, + is_var_split_kv=self.is_var_split_kv, + cfg=cfg, + static_split_kv=self.static_split_kv, + static_seq_len_k=self.static_seq_len_k, + cu_seqlens_q=cu_seqlens_q, + logical_num_heads_q=self.num_heads, + logical_seq_len_q=self.seq_len_q, + static_problem_shape_b=None, + static_problem_shape_s=num_query_tiles, + use_clc_dynamic=use_clc_dynamic, + tile_scheduler_config=work_queue_tile_scheduler_config, + pipeline_config=work_queue_pipeline_config, + name="mla_work_queue", + ) + + task_manager, _tmem_resources, named_res = build_mla_decode_task_manager( + cfg=cfg, + smem_q_latent_arr=smem_q_latent_arr, + smem_q_rope_arr=smem_q_rope_arr, + smem_kc_arr=smem_kc_arr, + smem_vc_arr=smem_vc_arr, + smem_p_arr=smem_p_arr, + tma_desc_q_latent=tma_desc_q_latent.get_ptr(), + tma_desc_q_rope=tma_desc_q_rope.get_ptr(), + tma_desc_c_latent=tma_desc_c_latent.get_ptr(), + tma_desc_c_rope=tma_desc_c_rope.get_ptr(), + tma_desc_c_transpose=tma_desc_c_transpose.get_ptr(), + page_offsets=page_offsets, + blk_coord=blk_coord, + tidx=tidx, + output=o, + acc_output=acc_o, + lse=lse, + acc_lse=acc_lse, + domain=Int32(1), # dummy; MlaTask recomputes per-task to avoid spills + work_queue=mla_work_queue, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + split_kv=max_split_kv, + logical_num_heads_q=self.num_heads, + logical_seq_len_q=self.seq_len_q, + tiled_mma_qk=tiled_mma_qk, + ) + + # Initialize pipelines (creates mbarriers in SMEM) + task_manager.setup_resources_and_tasks() + + # Fence mbarrier init then cluster sync so both CTAs see + # initialized barriers before tcgen05_alloc. The alloc with + # CTA_2 group uses ptxas-generated lifecycle barriers in + # static SMEM [0,1024) that require cross-CTA coordination. + prims.fence_mbarrier_init() + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + + # TMEM allocation AFTER lifecycle init + cluster sync (warp 8 only) + if warp_idx == cfg.mma_warp_id: + prims.tcgen05_alloc( + tmem_holding_buf_arr, cfg.num_tmem_cols, group="cta_2" + ) + prims.tcgen05_relinquish_alloc_permit(group="cta_2") + + # tcgen05.alloc.cta_group::2 is itself cluster-synchronous; the + # pre-allocation cluster barrier above protects lifecycle-mbarrier + # initialization, while another full-cluster barrier here only + # serializes the first consumers of the published TMEM base. + + # TMEM sync barrier: only warps 0-8 participate + participates_in_tmem_sync = warp_idx <= cfg.mma_warp_id + if cutlass.const_expr(cfg.use_fp8_split_mma_schedule): + participates_in_tmem_sync = participates_in_tmem_sync or ( + warp_idx == cfg.pv_mma_warp_id + ) + if cutlass.const_expr(cfg.use_fp8_dual_softmax_schedule): + participates_in_tmem_sync = participates_in_tmem_sync or ( + warp_idx >= cfg.second_compute_warp_ids[0] + and warp_idx <= cfg.second_compute_warp_ids[-1] + ) + if participates_in_tmem_sync: + prims.barrier_cta_sync( + cfg.tmem_sync_bar_id, thread_count=cfg.tmem_sync_bar_threads + ) + + # Only the warps synchronized above consume TMEM resources. Do + # not let the loader/page/padding warps race the allocation's + # publication slot merely to materialize an unused pointer. + tmem_base_addr = Int32(0) + if participates_in_tmem_sync: + tmem_base_addr = tmem_holding_buf_arr.load() + + # Set TMEM base address on resources + named_res["tmem_s"].tmem_base_addr = tmem_base_addr + named_res["tmem_o"].tmem_base_addr = tmem_base_addr + named_res["tmem_corr"].tmem_base_addr = tmem_base_addr + + # Set runtime params on TmemS (softmax) + named_res["tmem_s"].softmax_scale_log2 = softmax_scale_log2 + named_res["tmem_s"].smem_exchange = softmax_exchange_arr + + named_res[ + "tmem_corr" + ].smem_exchange = epilogue_exchange_arr.data_ptr().toint(Int32) + + # Set runtime params on GmemO (epilogue) + named_res["gmem_o"].output_scale = output_scale + named_res["gmem_o"].softmax_scale_log2 = softmax_scale_log2 + named_res["gmem_o"].smem_exchange = epilogue_exchange_arr.data_ptr().toint( + Int32 + ) + named_res["gmem_o"].split_kv = max_split_kv + + task_manager.run() + + # TMEM deallocation (MMA warp) + if warp_idx == cfg.mma_warp_id: + cta_rank = cute.arch.make_warp_uniform(mma_tile_coord_v) + peer_cta_rank = cta_rank ^ 1 + peer_mbar = prims.mapa(tmem_dealloc_mbar_arr, peer_cta_rank) + prims.mbarrier_arrive(peer_mbar) + while not prims.mbarrier_try_wait_parity(tmem_dealloc_mbar_arr, 0): + pass + tmem_arr_for_dealloc = prims.make_tmem_ptr( + tmem_holding_buf_arr.load(), self.acc_dtype + ) + prims.tcgen05_dealloc( + tmem_arr_for_dealloc, + cfg.num_tmem_cols, + group="cta_2", + ) + # Each producer CTA publishes completion only after its paired + # TMEM lifetime and all scheduled output work have retired. + if cutlass.const_expr(acc_o is not None and not self.is_persistent): + if prims.elect_sync(): + prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) + + @cute.jit + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + workspace: cute.Tensor, + ): + """Construct acc_o and acc_lse tensors from the workspace buffer.""" + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + # Workspace strides are aligned to 256 bits, expressed in Float16 + # elements because split-KV partial O is BF16 even for FP8 output. + align = 256 // cutlass.Float16.width + acc_o_layout = cute.make_layout( + (H, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, + cute.assume(split_kv * H * D, align), + cute.assume(H * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=cutlass.BFloat16) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( + (H, split_kv, S, B), + stride=(split_kv, 1, H * split_kv, H * split_kv * S), + ) + acc_lse_iter = cute.recast_ptr( + workspace.iterator + + Int64(cute.cosize(acc_o_layout)) * Int64(cutlass.BFloat16.width // 8), + dtype=self.lse_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @cute.kernel + def reduction_kernel( + self, + output: cute.Tensor, + lse: cute.Tensor, + acc_output: cute.Tensor, + acc_lse: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor | None, + block_split_kvs: cute.Tensor, + ): + """Dispatch the throughput 2CTA split-KV reduction body.""" + cfg = make_mla_decode_config( + mma_qk_tiler_mn=self.mma_qk_tiler_mn, + mma_pv_tiler_mn=self.mma_pv_tiler_mn, + rope_dim=self.rope_dim, + page_size=self.page_size, + qkv_dtype=self.qkv_dtype, + o_dtype=self.out_dtype, + mask_type=self.mask_type, + ) + if cutlass.const_expr(not self.is_persistent): + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + run_reduction_kernel( + self, + output, + lse, + acc_output, + acc_lse, + split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + cfg, + self.reduction_split_capacity, + REDUCTION_ROWS_PER_CTA, + ) + + @cute.kernel + def parallel_reduction_kernel( + self, + output: cute.Tensor, + lse: cute.Tensor, + acc_output: cute.Tensor, + acc_lse: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor | None, + block_split_kvs: cute.Tensor, + ): + """Dispatch the high-split fixed-D512 cluster reducer.""" + + cfg = make_mla_decode_config( + mma_qk_tiler_mn=self.mma_qk_tiler_mn, + mma_pv_tiler_mn=self.mma_pv_tiler_mn, + rope_dim=self.rope_dim, + page_size=self.page_size, + qkv_dtype=self.qkv_dtype, + o_dtype=self.out_dtype, + mask_type=self.mask_type, + ) + if cutlass.const_expr(not self.is_persistent): + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + topology = self.parallel_reduction_topology + run_parallel_reduction_kernel( + self.num_heads, + self.seq_len_q, + self.is_var_split_kv, + output, + lse, + acc_output, + acc_lse, + split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + cfg, + topology.actual_splits, + topology.cluster_size, + topology.slots_per_rank, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py new file mode 100644 index 000000000000..2512ef3dcfa1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py @@ -0,0 +1,490 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parallel standalone split-KV reduction for throughput 2CTA MLA.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 +from cutlass.experimental import primitives as prims + +from ...separate_reduction import finalize_log2_sum_exp, normalized_lse_weight +from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ..helpers.mask import MaskType, mask_visible_k_length +from ..helpers.math import ceil_div +from ..helpers.ops import fmax_f32, warp_reduce_max_f32, warp_reduce_sum_f32 +from ..helpers.query import flat_query_row_state, query_batch_bounds +from .work_partition import ( + runtime_row_prefix_active_split_count, + runtime_split_kv_cap, +) + +# One standalone-reducer thread loads an aligned BF16 vec4 partial-O fragment +# and accumulates its four values in FP32 registers. A 128-thread CTA therefore +# covers exactly one D=512 MLA row. These constants are deliberately local to +# the 2CTA policy; the reference reducer keeps its four-warp launch unchanged. +PARALLEL_REDUCTION_THREADS = 128 +PARALLEL_REDUCTION_ELEMENTS_PER_THREAD = 4 +PARALLEL_REDUCTION_HEAD_DIM = ( + PARALLEL_REDUCTION_THREADS * PARALLEL_REDUCTION_ELEMENTS_PER_THREAD +) + + +@cute.jit +def _parallel_reduction_row_state( + tile_size_q: cutlass.Constexpr[int], + num_heads: cutlass.Constexpr[int], + seq_len_q: cutlass.Constexpr[int], + is_var_split_kv: cutlass.Constexpr[bool], + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_kv, + row_in_tile, + query_tile_idx, + batch_idx, + cfg, +): + """Return logical row coordinates and its runtime active split count. + + Keep this arithmetic identical to the serial reducer below. A causal row + can see fewer K tiles than the last row in its producer tile, and a + variable-length batch can expose fewer real partitions than the compiled + static split capacity. The parallel reducer must ignore both kinds of + empty workspace rows. + """ + + ( + _, + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + ) = flat_query_row_state( + row_in_tile, + query_tile_idx, + tile_size_q, + num_heads, + seq_len_q, + cu_seqlens_q, + batch_idx, + ) + # Producer and reducer share the same per-batch cap and configured-span group + # partition. ``split_kv`` remains the grid/workspace capacity. + split_kv_cap = runtime_split_kv_cap( + split_kv, + is_var_split_kv, + block_split_kvs, + batch_idx, + ) + tile_k = cache_seqs[batch_idx] + row_k = tile_k + if cutlass.const_expr(cfg.mask_type == MaskType.CAUSAL.value and seq_len_q > 1): + _, logical_seq_len_q = query_batch_bounds( + cu_seqlens_q, + batch_idx, + seq_len_q, + ) + _, _, tile_last_logical_q_idx, _, _ = flat_query_row_state( + Int32(tile_size_q - 1), + query_tile_idx, + tile_size_q, + num_heads, + seq_len_q, + cu_seqlens_q, + batch_idx, + ) + tile_k = mask_visible_k_length( + cfg.mask_type, + tile_k, + tile_last_logical_q_idx, + logical_seq_len_q, + ) + row_k = mask_visible_k_length( + cfg.mask_type, + cache_seqs[batch_idx], + logical_q_idx, + logical_seq_len_q, + ) + tile_k_tile_total = (tile_k + cfg.mma_qk_tiler[1] - 1) // cfg.mma_qk_tiler[1] + row_k_tile_total = (row_k + cfg.mma_qk_tiler[1] - 1) // cfg.mma_qk_tiler[1] + active_split_kv = runtime_row_prefix_active_split_count( + row_k_tile_total, + tile_k_tile_total, + split_kv_cap, + ) + return ( + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + active_split_kv, + ) + + +@cute.jit +def _store_parallel_reduction_result( + output, + lse, + output_vals, + global_lse, + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + element_idx, + tidx, + batch_idx, + cu_seqlens_q, +): + """Publish one normalized FP32 output fragment and final LSE.""" + + if tidx == Int32(0) and query_is_valid: + if cutlass.const_expr(cu_seqlens_q is not None): + lse[logical_head_idx, storage_q_idx] = global_lse + else: + lse[logical_head_idx, logical_q_idx, batch_idx] = global_lse + + out_element_dtype = output.element_type + output_regs = cutlass.Array( + out_element_dtype, + PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + output_regs[j] = out_element_dtype(output_vals[j]) + + if query_is_valid: + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + dim_idx = element_idx + Int32(j) + if cutlass.const_expr(cu_seqlens_q is not None): + output[logical_head_idx, dim_idx, storage_q_idx] = output_regs[j] + else: + output[logical_head_idx, dim_idx, logical_q_idx, batch_idx] = ( + output_regs[j] + ) + + +@cute.jit +def run_parallel_reduction_kernel( + num_heads: cutlass.Constexpr[int], + seq_len_q: cutlass.Constexpr[int], + is_var_split_kv: cutlass.Constexpr[bool], + output, + lse, + acc_output, + acc_lse, + split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + cfg, + actual_splits: cutlass.Constexpr[int], + cluster_size: cutlass.Constexpr[int], + slots_per_rank: cutlass.Constexpr[int], +): + """Reduce one D=512 row cooperatively across a padded CTA cluster. + + Rank ``r`` owns ``slots_per_rank`` contiguous split slots. Slots beyond + ``actual_splits`` and row-specific inactive splits perform no GMEM access + or arithmetic. G1 writes its result directly; for G2/G4/G8 every rank + publishes a neutral-or-valid ``(FP32 LSE, BF16 O[512])`` state to DSMEM and + rank zero performs the final merge. Arithmetic remains FP32. + + For a one-wave producer, the caller waits for its post-TMEM dependent-launch + signal before entering this body. + """ + + block_idx_x, query_tile_idx, batch_idx = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_idx = tidx % Int32(cfg.threads_per_warp) + cluster_rank = cute.arch.block_idx_in_cluster() + row_in_tile = block_idx_x // Int32(cluster_size) + ( + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + active_split_kv, + ) = _parallel_reduction_row_state( + cfg.mma_qk_tiler[0], + num_heads, + seq_len_q, + is_var_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_kv, + row_in_tile, + query_tile_idx, + batch_idx, + cfg, + ) + + element_idx = tidx * Int32(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD) + neg_inf = Float32(-Float32.inf) + lse_slots_per_lane = ceil_div(slots_per_rank, cfg.threads_per_warp) + smem_local_lse = cutlass.Array( + Float32, + 1, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + smem_local_scale = cutlass.Array( + Float32, + slots_per_rank, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + # One warp computes the rank-local softmax state for this row. All output + # threads consume the same scales, so repeating LSE loads and SFU work in + # every vector lane only adds instructions and register pressure. + if warp_idx == Int32(0): + lane_lse = cutlass.Array( + Float32, + lse_slots_per_lane, + space=cutlass.AddressSpace.rmem, + ) + local_lse_max = neg_inf + for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): + local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) + split_idx = cluster_rank * Int32(slots_per_rank) + local_slot_idx + active_slot = ( + query_is_valid + & (local_slot_idx < Int32(slots_per_rank)) + & (split_idx < Int32(actual_splits)) + & (split_idx < active_split_kv) + ) + lane_lse[lane_slot_i] = neg_inf + # Keep the workspace access inside the dynamic predicate. Padded + # ranks must not form a load from the unpadded producer allocation. + if active_slot: + lane_lse[lane_slot_i] = Float32( + acc_lse[row_in_tile, split_idx, query_tile_idx, batch_idx] + ) + local_lse_max = fmax_f32( + local_lse_max, + lane_lse[lane_slot_i], + ) + + local_lse_max = warp_reduce_max_f32(local_lse_max) + local_exp_frame = local_lse_max if local_lse_max != neg_inf else Float32(0.0) + lane_exp = cutlass.Array( + Float32, + lse_slots_per_lane, + space=cutlass.AddressSpace.rmem, + ) + local_sum_lse = Float32(0.0) + for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): + lane_exp[lane_slot_i] = Float32( + cute.math.exp2( + lane_lse[lane_slot_i] - local_exp_frame, + fastmath=True, + ) + ) + local_sum_lse += lane_exp[lane_slot_i] + local_sum_lse = warp_reduce_sum_f32(local_sum_lse) + local_lse_value = finalize_log2_sum_exp(local_exp_frame, local_sum_lse) + if lane_idx == Int32(0): + smem_local_lse[0] = local_lse_value + + for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): + local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) + if local_slot_idx < Int32(slots_per_rank): + smem_local_scale[local_slot_idx] = normalized_lse_weight( + lane_lse[lane_slot_i], local_lse_value + ) + + prims.barrier_cta_sync(SPLIT_REDUCTION_SCALE_BARRIER_ID) + + # Every thread owns one contiguous BF16 vec4 GMEM fragment and immediately + # converts it to FP32 registers. The shared scales directly normalize the + # rank-local result, so no second numerator array or reciprocal is needed. + local_output = cutlass.Array( + Float32, + PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + local_output[j] = Float32(0.0) + + acc_output_ptr = acc_output.iterator.raw_ptr() + for slot_i in cutlass.range_constexpr(slots_per_rank): + split_idx = cluster_rank * Int32(slots_per_rank) + Int32(slot_i) + active_slot = ( + query_is_valid + & (split_idx < Int32(actual_splits)) + & (split_idx < active_split_kv) + ) + if active_slot: + split_scale = Float32(smem_local_scale[slot_i]) + partial_offset = Int64( + row_in_tile * acc_output.stride[0] + + split_idx * acc_output.stride[1] + + element_idx * acc_output.stride[2] + + query_tile_idx * acc_output.stride[3] + + batch_idx * acc_output.stride[4] + ) + partial_output = ( + (acc_output_ptr + partial_offset) + .load( + count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + alignment=8, + ) + .to(Float32) + ) + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + local_output[j] += partial_output[j] * split_scale + + local_lse_value = smem_local_lse[0] + + # Small split counts need no cluster exchange. G1 still uses one CTA-local + # scale barrier so its LSE/SFU work is shared rather than repeated 128 ways. + if cutlass.const_expr(cluster_size == 1): + _store_parallel_reduction_result( + output, + lse, + local_output, + local_lse_value, + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + element_idx, + tidx, + batch_idx, + cu_seqlens_q, + ) + return + + # G2+ publishes normalized BF16 O and one FP32 LSE scalar per rank. Every rank + # writes the neutral (-inf, zero) state for padded or invalid rows and + # participates in both cluster barriers, so rank zero's DSMEM pointers + # cannot outlive peers. + smem_output = cutlass.Array( + acc_output.element_type, + PARALLEL_REDUCTION_HEAD_DIM, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + local_output_bf16 = cutlass.Array( + acc_output.element_type, + PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + local_output_bf16[j] = acc_output.element_type(local_output[j]) + (smem_output.data_ptr() + element_idx).store( + local_output_bf16.data_ptr().load( + count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + alignment=8, + ), + alignment=8, + ) + + prims.barrier_cta_sync(0) + # Rank zero immediately consumes each peer's published SMEM through DSMEM. + # Use the ordered arrival; the relaxed form does not order prior writes. + prims.barrier_cluster_arrive() + prims.barrier_cluster_wait() + + smem_merged_lse = cutlass.Array( + Float32, + 1, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + smem_peer_scale = cutlass.Array( + Float32, + cluster_size, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + if cluster_rank == Int32(0): + # One warp merges the per-rank LSE scalars. Keep each mapa rank uniform + # at its issue site, then let one lane retain that peer's value. + if warp_idx == Int32(0): + lane_peer_lse = neg_inf + for peer_rank_i in cutlass.range_constexpr(cluster_size): + peer_rank = Int32(peer_rank_i) + peer_lse_ptr = prims.mapa( + smem_local_lse.data_ptr(), + peer_rank, + ) + if lane_idx == peer_rank: + lane_peer_lse = peer_lse_ptr.load() + + merged_lse_max = warp_reduce_max_f32(lane_peer_lse) + merged_exp_frame = ( + merged_lse_max if merged_lse_max != neg_inf else Float32(0.0) + ) + lane_peer_exp = Float32( + cute.math.exp2( + lane_peer_lse - merged_exp_frame, + fastmath=True, + ) + ) + merged_sum_lse = warp_reduce_sum_f32(lane_peer_exp) + merged_lse = finalize_log2_sum_exp(merged_exp_frame, merged_sum_lse) + if lane_idx == Int32(0): + smem_merged_lse[0] = merged_lse + if lane_idx < Int32(cluster_size): + smem_peer_scale[lane_idx] = normalized_lse_weight( + lane_peer_lse, merged_lse + ) + + prims.barrier_cta_sync(SPLIT_REDUCTION_SCALE_BARRIER_ID) + + merged_output = cutlass.Array( + Float32, + PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + merged_output[j] = Float32(0.0) + + for peer_rank_i in cutlass.range_constexpr(cluster_size): + peer_rank = Int32(peer_rank_i) + peer_output = prims.mapa(smem_output.data_ptr(), peer_rank) + peer_scale = Float32(smem_peer_scale[peer_rank_i]) + peer_output_vals = ( + (peer_output + element_idx) + .load( + count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + alignment=8, + ) + .to(Float32) + ) + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + merged_output[j] += peer_output_vals[j] * peer_scale + + _store_parallel_reduction_result( + output, + lse, + merged_output, + smem_merged_lse[0], + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + element_idx, + tidx, + batch_idx, + cu_seqlens_q, + ) + + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py new file mode 100644 index 000000000000..5de512dea6d0 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py @@ -0,0 +1,298 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Split-KV reduction body for the throughput 2CTA MLA TS kernel.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 +from cutlass.experimental import primitives as prims + +from .config import ( + REDUCTION_THREADS_PER_ROW, + REDUCTION_VALUES_PER_THREAD, + REDUCTION_VECTOR_BYTES, +) +from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ..helpers.math import ceil_div +from ..helpers.mask import MaskType, mask_visible_k_length +from ..helpers.ops import ( + fmax_f32, + warp_reduce_max_f32, + warp_reduce_sum_f32, + vector_from_scalars, +) +from ..helpers.query import flat_query_row_state, query_batch_bounds +from .work_partition import ( + runtime_row_prefix_active_split_count, + runtime_split_kv_cap, +) + + +@cute.jit +def run_reduction_kernel( + kernel, + output, + lse, + acc_output, + acc_lse, + split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + cfg, + max_splits: cutlass.Constexpr[int], + rows_per_cta: cutlass.Constexpr[int], +): + """Combine consecutive logical split rows in one reference-reducer CTA. + + Each 64-thread row group owns one D512 output row. Its first warp computes + FP32 LSE rescale factors, both warps consume one contiguous 16-byte BF16 + fragment per thread, and the final accumulation remains FP32. The launch + is flattened over logical query rows, then each row is decomposed back into + its physical M128 workspace tile and row. Only the last CTA can contain + inactive row groups; those groups still synchronize but never access the + workspace or publish public output. + """ + reduction_group_idx, _, batch_idx = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + row_in_cta = tidx // Int32(REDUCTION_THREADS_PER_ROW) + row_thread_idx = tidx - row_in_cta * Int32(REDUCTION_THREADS_PER_ROW) + row_warp_idx = cute.arch.make_warp_uniform( + row_thread_idx // Int32(cfg.threads_per_warp) + ) + lane_idx = row_thread_idx % Int32(cfg.threads_per_warp) + + physical_tile_rows = Int32(cfg.mma_qk_tiler[0]) + logical_query_rows = Int32(kernel.num_heads * kernel.seq_len_q) + local_flat_query_row = reduction_group_idx * Int32(rows_per_cta) + row_in_cta + row_is_valid = local_flat_query_row < logical_query_rows + # Clamp the final CTA's inactive row groups before decomposing the physical + # workspace coordinate. ``rows_per_cta`` divides M128, so every CTA + # containing valid rows belongs to exactly one producer query tile. + safe_local_flat_query_row = cute.math.min( + local_flat_query_row, logical_query_rows - Int32(1) + ) + if cutlass.const_expr(kernel.num_heads * kernel.seq_len_q <= cfg.mma_qk_tiler[0]): + query_tile_idx = Int32(0) + row_in_tile = safe_local_flat_query_row + else: + query_tile_idx = safe_local_flat_query_row // physical_tile_rows + row_in_tile = safe_local_flat_query_row - query_tile_idx * physical_tile_rows + if cutlass.const_expr(cu_seqlens_q is None): + # Fixed Q storage is already flat in the producer's logical row order; + # bypass packed-offset clamping on this latency-sensitive reducer path. + storage_flat_query_row = safe_local_flat_query_row + logical_q_idx = storage_flat_query_row // Int32(kernel.num_heads) + mapped_query_is_valid = True + else: + ( + storage_flat_query_row, + _, + logical_q_idx, + _, + mapped_query_is_valid, + ) = flat_query_row_state( + row_in_tile, + query_tile_idx, + cfg.mma_qk_tiler[0], + kernel.num_heads, + kernel.seq_len_q, + cu_seqlens_q, + batch_idx, + ) + query_is_valid = row_is_valid and mapped_query_is_valid + public_flat_query_row = storage_flat_query_row + if cutlass.const_expr(cu_seqlens_q is None): + public_flat_query_row = public_flat_query_row + batch_idx * Int32( + kernel.seq_len_q * kernel.num_heads + ) + + # The scalar split count remains the grid/workspace capacity. A variable- + # split launch optionally contracts it with block_split_kvs[batch]. + if cutlass.const_expr( + kernel.static_split_kv is not None and not kernel.is_var_split_kv + ): + split_kv_cap = Int32(max_splits) + else: + split_kv_cap = runtime_split_kv_cap( + split_kv, + kernel.is_var_split_kv, + block_split_kvs, + batch_idx, + ) + # Producer splits are sized from the group's largest K domain, while each + # logical row consumes only the prefix containing its visible K tiles. + tile_k = cache_seqs[batch_idx] + row_k = tile_k + if cutlass.const_expr( + cfg.mask_type == MaskType.CAUSAL.value and kernel.seq_len_q > 1 + ): + _, logical_seq_len_q = query_batch_bounds( + cu_seqlens_q, + batch_idx, + kernel.seq_len_q, + ) + _, _, tile_last_logical_q_idx, _, _ = flat_query_row_state( + Int32(cfg.mma_qk_tiler[0] - 1), + query_tile_idx, + cfg.mma_qk_tiler[0], + kernel.num_heads, + kernel.seq_len_q, + cu_seqlens_q, + batch_idx, + ) + tile_k = mask_visible_k_length( + cfg.mask_type, + tile_k, + tile_last_logical_q_idx, + logical_seq_len_q, + ) + row_k = mask_visible_k_length( + cfg.mask_type, + cache_seqs[batch_idx], + logical_q_idx, + logical_seq_len_q, + ) + tile_k_tile_total = (tile_k + cfg.mma_qk_tiler[1] - 1) // cfg.mma_qk_tiler[1] + row_k_tile_total = (row_k + cfg.mma_qk_tiler[1] - 1) // cfg.mma_qk_tiler[1] + # A causal row consumes the prefix of configured-span ranges intersecting + # its visible K tiles. This must stay identical to the producer geometry. + local_split_kv = runtime_row_prefix_active_split_count( + row_k_tile_total, + tile_k_tile_total, + split_kv_cap, + ) + + smem_lse_scale = cutlass.Array( + kernel.lse_dtype, + rows_per_cta * max_splits, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + row_scale_offset = row_in_cta * Int32(max_splits) + + acc_lse_tile = acc_lse[row_in_tile, None, query_tile_idx, batch_idx] + if row_warp_idx == 0: + # The first warp for each row owns its log-sum-exp merge. It publishes + # one rescale factor per active split for the row's second warp too. + lse_per_thread = ceil_div(max_splits, cfg.threads_per_warp) + local_lse = cutlass.Array(kernel.lse_dtype, lse_per_thread) + lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = lane_idx + i * cfg.threads_per_warp + active_slot = query_is_valid & cute.elem_less( + split_kv_idx, + local_split_kv, + ) + local_lse[i] = -kernel.lse_dtype.inf + # Keep the workspace access inside the dynamic predicate. Split + # capacities need not be warp-aligned, so padded lanes must not + # form an address beyond the producer allocation. + if active_slot: + local_lse[i] = acc_lse_tile[split_kv_idx] + lse_max = fmax_f32(lse_max, local_lse[i]) + lse_max = warp_reduce_max_f32(lse_max) + lse_max = lse_max if lse_max != -kernel.lse_dtype.inf else 0.0 + sum_lse = kernel.lse_dtype(0.0) + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = warp_reduce_sum_f32(sum_lse) + has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype(0.0) + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if has_finite_mass + else -kernel.lse_dtype.inf + ) + if lane_idx == 0 and query_is_valid: + (lse.iterator.raw_ptr() + Int64(public_flat_query_row)).store(global_lse) + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = lane_idx + i * cfg.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[row_scale_offset + split_kv_idx] = ( + cute.math.exp2(local_lse[i] - global_lse, fastmath=True) + if has_finite_mass + else kernel.lse_dtype(0.0) + ) + + # Independent writer warps publish scales before any row consumes + # them. Invalid/padded rows participate so this remains a full CTA barrier. + prims.barrier_cta_sync(SPLIT_REDUCTION_SCALE_BARRIER_ID) + + element_idx = row_thread_idx * Int32(REDUCTION_VALUES_PER_THREAD) + acc_vec = vector_from_scalars( + ( + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + ), + dtype=Float32, + ) + acc_output_ptr = acc_output.iterator.raw_ptr() + if query_is_valid: + partial_elem_offset_base = ( + Int64(batch_idx) + * Int64(cute.size(acc_output.shape[3])) + * Int64(physical_tile_rows) + * Int64(split_kv) + * Int64(cfg.latent_dim) + + Int64(query_tile_idx) + * Int64(physical_tile_rows) + * Int64(split_kv) + * Int64(cfg.latent_dim) + + Int64(row_in_tile) * Int64(split_kv) * Int64(cfg.latent_dim) + + Int64(element_idx) + ) + partial_output_ptr = acc_output_ptr + partial_elem_offset_base + # S2 is the common two-wave 2CTA decode reducer. Materialize its two + # fixed partials directly; larger reducers retain the compact dynamic + # loop because fully unrolling S4 increases instruction pressure. + if cutlass.const_expr(max_splits == 2): + for i in cutlass.range_constexpr(max_splits): + if Int32(i) < local_split_kv: + partial_vec = ( + (partial_output_ptr + Int64(i * cfg.latent_dim)) + .load( + count=REDUCTION_VALUES_PER_THREAD, + alignment=REDUCTION_VECTOR_BYTES, + ) + .to(Float32) + ) + scale = Float32(smem_lse_scale[row_scale_offset + Int32(i)]) + acc_vec = acc_vec + partial_vec * scale + else: + for i in range(local_split_kv): + partial_vec = ( + (partial_output_ptr + Int64(i) * Int64(cfg.latent_dim)) + .load( + count=REDUCTION_VALUES_PER_THREAD, + alignment=REDUCTION_VECTOR_BYTES, + ) + .to(Float32) + ) + scale = Float32(smem_lse_scale[row_scale_offset + i]) + acc_vec = acc_vec + partial_vec * scale + + output_elem_offset = Int64(public_flat_query_row) * Int64( + cfg.latent_dim + ) + Int64(element_idx) + (output.iterator.raw_ptr() + output_elem_offset).store( + acc_vec.to(output.element_type), + alignment=REDUCTION_VECTOR_BYTES, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py new file mode 100644 index 000000000000..55523137fdea --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py @@ -0,0 +1,3598 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resource definitions for the throughput 2CTA MLA decode TS kernel. + +Resource classes matching the throughput 2CTA pipeline structure: + +GMEM/register resources (not pipelined) +---------------------------------------- +- PageOffsetWindowResource: one warp-wide, 32-page-ID register window + +SMEM resources (pipelined) +-------------------------- +- SmemQResource : TmaUmmaAsync(1 stage), LoadTma -> Mma +- SmemKVResource : TmaUmmaAsync(7 stages), LoadTma -> Mma + +SMEM/TMEM resources (pipelined) +------------------------------- +- SmemPResource : UmmaConsumerAsync(2 stages), SoftmaxTask -> MmaTask +- TmemSResource : UmmaProducerAsync(2 stages), MmaTask -> SoftmaxTask +- TmemCorrResource : Async(2 stages), SoftmaxTask -> CorrectionTask +- TmemOResource : UmmaProducerAsync(1 stage), MmaTask -> CorrectionTask + +GMEM (no pipeline) +------------------ +- GmemOResource : No pipeline, Correction -> GMEM +""" + +from dataclasses import dataclass, field +from typing import Any, ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims +from cutlass import Boolean, Float32, Int16, Int32, Int64 + +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + TaskLocalVariable, + WorkQueue, +) +from cutlass.experimental.task_scheduling.resources import consumer_work, producer_work +from cutlass.experimental.task_scheduling.enums import WorkAttr +from ...mask import kv_tile_needs_right_mask +from ...tensor_map import transform_ragged_coords + +from .config import ( + V_SMEM_K_BLOCK_TOKENS, + V_TMA_LATENT_ELEMENTS, + MlaDecodeConfig, +) +from ..helpers.constants import ( + BF16_OUTPUT_VECTOR_ELEMENTS, + EPILOGUE_COLUMN_GROUP_SHIFT, + EPILOGUE_ROW_MASK, + EPILOGUE_THREAD_TILE_MASK, + EPILOGUE_THREAD_TILE_THREADS, + FP8_OUTPUT_VECTOR_ELEMENTS, + PACKED_FP8_OUTPUT_REGS, + TCGEN05_32B_REGS_PER_LOAD, + TCGEN05_32B_SHAPE, + WARP_LANE_SHIFT, +) +from ..helpers.tile_scheduler import ( + MLAStaticTileScheduler, + MLAStaticTileSchedulerParams, + create_mla_static_tile_scheduler, + divmod_constexpr_power_of_two_or_fdd, +) +from ..helpers.math import ( + ceil_div, + mma_k_step_for_qkv, + mma_kind_for_qkv, + add_packed_f32x2, + fma_packed_f32x2, + mul_packed_f32x2, + output_dtype, + p_desc_layout, + p_desc_leading_byte_offset, + p_desc_stride_byte_offset, + qk_desc_layout, + qk_desc_layout_for_head_dim, + qk_desc_leading_byte_offset, + qk_desc_leading_byte_offset_for_head_dim, + qk_desc_stride_byte_offset, + qk_desc_stride_byte_offset_for_head_dim, + qkv_dtype, + qkv_major_k_stride_bytes_for, +) +from ..helpers.ops import ( + fp8_log2_quant_scale, + fp8_quant_scale_rcp, + fmax_f32, + pack_float4_to_fp8_e4m3, +) +from ..helpers.mask import MaskType, mask_visible_k_length +from ..helpers.query import ( + flat_query_row_state, + query_batch_bounds, + runtime_flat_query_tile_has_rows, +) +from .work_partition import ( + runtime_split_kv_cap, + runtime_split_tile_range, +) + + +def _install_task_local_specs(resource: object, specs: tuple[tuple, ...]) -> None: + """Install TaskLocalVariable fields declared by resource classes.""" + for spec in specs: + field_name, dtype, default, docs = spec[:4] + runtime_slot_name = spec[4] if len(spec) > 4 else None + object.__setattr__( + resource, + field_name, + TaskLocalVariable( + dtype=dtype, + default=default, + docs=docs, + runtime_slot_name=runtime_slot_name, + ), + ) + + +@dataclass(kw_only=True) +class HighThroughputMlaResource(MemoryResource): + """Base class that binds captured-schedule task-local variables.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = () + + def __post_init__(self) -> None: + _install_task_local_specs(self, self._task_local_specs) + + +# ===================================================================== +# WorkThrottleBarrierResource — Cluster-safe CLC scheduler pacing +# ===================================================================== + + +@dataclass(kw_only=True) +class WorkThrottleBarrierResource(MemoryResource): + """Pace CLC schedule token reuse against the leader CTA's active MMA task. + + The barrier has no payload. Its producer task already runs only on CTA 0, + so ordinary public pipeline ownership is sufficient for every stage and + producer-tail operation. + """ + + is_barrier: cutlass.Constexpr[bool] = True + + +# ===================================================================== +# MlaWorkQueue — Persistent tile scheduler for MLA decode +# ===================================================================== + + +class MlaTsWorkTileInfo: + """MLA work tile with scalar fields for TS persistent loop carry. + + Besides the scheduler coordinate, cache the per-tile K-domain values that + hot resource paths need. Each persistent warp loop computes the K domain + once and passes the derived indices to the page-offset, TMA, MMA, and + softmax bodies. + """ + + @cute.jit + def __init__( + self, + tile_idx, + is_valid, + k_len=0, + k_tile_count=0, + k_index_base=0, + ): + """Initialize the staged tile coordinate and cached K-domain metadata.""" + cluster_idx, seq_q_idx, batch_idx, split_kv_idx = tile_idx + self._cluster_idx = Int32(cluster_idx) + self._seq_q_idx = Int32(seq_q_idx) + self._batch_idx = Int32(batch_idx) + self._split_kv_idx = Int32(split_kv_idx) + self._is_valid = Boolean(is_valid) + self._k_len = Int32(k_len) + self._k_tile_count = Int32(k_tile_count) + self._k_index_base = Int32(k_index_base) + + @property + @cute.jit + def tile_idx(self): + """Return the scheduler tile coordinate tuple.""" + return ( + self._cluster_idx, + self._seq_q_idx, + self._batch_idx, + self._split_kv_idx, + ) + + @property + @cute.jit + def is_valid_tile(self): + """Return whether this tile participates in the current launch.""" + return self._is_valid + + @property + @cute.jit + def k_len(self): + """Return the request's graph-live K length for this tile.""" + return self._k_len + + @property + @cute.jit + def k_tile_count(self): + """Return the number of K tiles assigned to this work tile.""" + return self._k_tile_count + + @property + @cute.jit + def k_index_base(self): + """Return the first K tile index owned by this work tile.""" + return self._k_index_base + + @cute.jit + def update_from(self, other) -> None: + """Replace this tile info with another tile info object.""" + cluster_idx, seq_q_idx, batch_idx, split_kv_idx = other.tile_idx + self._cluster_idx = cluster_idx + self._seq_q_idx = seq_q_idx + self._batch_idx = batch_idx + self._split_kv_idx = split_kv_idx + self._is_valid = Boolean(other.is_valid_tile) + self._k_len = Int32(other.k_len) + self._k_tile_count = Int32(other.k_tile_count) + self._k_index_base = Int32(other.k_index_base) + + def __extract_mlir_values__(self): + """Extract scalar MLIR values for value-type lowering.""" + values = cutlass.extract_mlir_values(self._cluster_idx) + values += cutlass.extract_mlir_values(self._seq_q_idx) + values += cutlass.extract_mlir_values(self._batch_idx) + values += cutlass.extract_mlir_values(self._split_kv_idx) + values += cutlass.extract_mlir_values(self._is_valid) + values += cutlass.extract_mlir_values(self._k_len) + values += cutlass.extract_mlir_values(self._k_tile_count) + values += cutlass.extract_mlir_values(self._k_index_base) + return values + + def __new_from_mlir_values__(self, values): + """Rebuild a tile info object from lowered scalar MLIR values.""" + return MlaTsWorkTileInfo( + ( + cutlass.new_from_mlir_values(self._cluster_idx, [values[0]]), + cutlass.new_from_mlir_values(self._seq_q_idx, [values[1]]), + cutlass.new_from_mlir_values(self._batch_idx, [values[2]]), + cutlass.new_from_mlir_values(self._split_kv_idx, [values[3]]), + ), + cutlass.new_from_mlir_values(self._is_valid, [values[4]]), + cutlass.new_from_mlir_values(self._k_len, [values[5]]), + cutlass.new_from_mlir_values(self._k_tile_count, [values[6]]), + cutlass.new_from_mlir_values(self._k_index_base, [values[7]]), + ) + + +@dataclass(kw_only=True) +class MlaWorkQueue(WorkQueue): + """WorkQueue that preserves MLA coordinates for static or CLC scheduling. + + Static scheduling uses ``MLAStaticTileScheduler`` directly. BF16 CLC + scheduling uses the public CUTLASS ``WorkQueue`` implementation, then maps + its cluster response ``(cluster_rank, 0, linear_cluster)`` back to MLA's + ``(cluster_rank, s_idx, b_idx, split_kv_idx)`` coordinate. Both modes feed + the same K-domain cache and task-local value type. + + ``tile_sched_params`` is the ``MLAStaticTileSchedulerParams`` object created + by the kernel. + """ + + tile_sched_params: Any = None + cache_seqs: Any = None # cache_seqs tensor for k_tile_count + split_kv: Any = None # maximum split slots in the launch/workspace + block_split_kvs: Any = None # optional per-batch split caps + is_var_split_kv: cutlass.Constexpr[bool] = False + cfg: Any = None # MlaDecodeConfig for tile sizes + static_split_kv: cutlass.Constexpr = None + static_seq_len_k: cutlass.Constexpr = None + cu_seqlens_q: Any = None + logical_num_heads_q: cutlass.Constexpr[int] = 128 + logical_seq_len_q: cutlass.Constexpr[int] = 1 + static_problem_shape_b: cutlass.Constexpr[int] = None + static_problem_shape_s: cutlass.Constexpr[int] = None + use_clc_dynamic: cutlass.Constexpr[bool] = False + work_tile: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + skip_work_tile: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def __init__( + self, + tile_sched_params, + cache_seqs=None, + split_kv=None, + block_split_kvs=None, + is_var_split_kv=False, + cfg=None, + static_split_kv=None, + static_seq_len_k=None, + cu_seqlens_q=None, + logical_num_heads_q=128, + logical_seq_len_q=1, + static_problem_shape_b=None, + static_problem_shape_s=None, + use_clc_dynamic=False, + tile_scheduler_config=None, + **kwargs, + ): + self.use_clc_dynamic = use_clc_dynamic + if use_clc_dynamic: + WorkQueue.__init__( + self, + tile_scheduler_config=tile_scheduler_config, + **kwargs, + ) + else: + # The custom static scheduler does not use TileSchedulerConfig. + MemoryResource.__init__(self, **kwargs) + self.tile_scheduler_config = None + self.tile_sched_params = tile_sched_params + self.cache_seqs = cache_seqs + self.split_kv = split_kv + self.block_split_kvs = block_split_kvs + self.is_var_split_kv = is_var_split_kv + self.cfg = cfg + self.static_split_kv = static_split_kv + self.static_seq_len_k = static_seq_len_k + self.cu_seqlens_q = cu_seqlens_q + self.logical_num_heads_q = logical_num_heads_q + self.logical_seq_len_q = logical_seq_len_q + self.static_problem_shape_b = static_problem_shape_b + self.static_problem_shape_s = static_problem_shape_s + self.work_tile = TaskLocalVariable( + dtype=MlaTsWorkTileInfo, + default=MlaTsWorkTileInfo( + (Int32(0), Int32(0), Int32(0), Int32(0)), + Boolean(False), + Int32(0), + Int32(0), + Int32(0), + ), + docs="Current MLA persistent-scheduler work tile.", + ) + self.skip_work_tile = TaskLocalVariable( + dtype=Boolean, + default=Boolean(False), + docs="Whether the current work tile should skip skippable work.", + ) + + def create_tile_scheduler(self): + if cutlass.const_expr(self.use_clc_dynamic): + return WorkQueue.create_tile_scheduler(self) + return create_mla_static_tile_scheduler( + self.tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + + def _create_placeholder_tile_scheduler(self): + """Create a dead structural scheduler for the shared prologue. + + Real persistent scheduling is created per task in MlaTask. Keeping + this placeholder independent of grid_dim avoids shared-prologue nctaid + values being CSE'd into every warp-dispatch branch. + """ + blk = cute.arch.block_idx() + dummy_params = MLAStaticTileSchedulerParams( + self.tile_sched_params.is_persistent, + blk[0], + blk[0], + self.tile_sched_params.cluster_shape_mnk, + blk[0], + ) + return MLAStaticTileScheduler(dummy_params, blk[0], blk, blk) + + def create(self) -> None: + """Create pipeline (if any) and the scheduler used by all TS tasks.""" + if cutlass.const_expr(self.use_clc_dynamic): + WorkQueue.create(self) + return + # Call MemoryResource.create() (not WorkQueue.create) to avoid + # the tile_scheduler_config check in the parent. + MemoryResource.create(self) + self.tile_scheduler = self._create_placeholder_tile_scheduler() + + @cute.jit + def initial_work_tile_info(self): + work_tile = self.tile_scheduler.initial_work_tile_info() + if cutlass.const_expr(self.use_clc_dynamic): + return self._work_tile_from_clc_tile(work_tile) + return self.wrap_work_tile(work_tile) + + @cute.jit + def wrap_work_tile(self, work_tile): + return MlaTsWorkTileInfo(work_tile.tile_idx, work_tile.is_valid_tile) + + @cute.jit + def _work_tile_from_clc_tile(self, work_tile): + """Map one cluster-wide CLC response to MLA's four coordinates.""" + query_cluster_idx, _, split_batch_idx = work_tile.tile_idx + params = self.tile_sched_params + cluster_width = Int32(params.cluster_shape_mnk[0]) + s_idx = query_cluster_idx // cluster_width + cluster_idx = query_cluster_idx % cluster_width + split_kv_idx, b_idx = divmod_constexpr_power_of_two_or_fdd( + split_batch_idx, + self.static_problem_shape_b, + params.problem_shape_b_fdd, + ) + return self._make_work_tile_info( + (cluster_idx, s_idx, b_idx, split_kv_idx), + work_tile.is_valid_tile, + ) + + @cute.jit + def _make_work_tile_info( + self, + tile_idx, + is_valid, + ): + cfg = self.cfg + params = self.tile_sched_params + _, s_idx, b_idx, split_kv_idx = tile_idx + safe_b_idx = cute.math.min(b_idx, params.problem_shape_b - Int32(1)) + if cutlass.const_expr(self.static_seq_len_k is not None): + K = Int32(self.static_seq_len_k) + else: + K = Int32(self.cache_seqs[safe_b_idx]) + sequence_k_len = K + + # The launch uses the largest flat-Q tile count in the batch. Shorter + # variable-Q requests can therefore receive a whole tile containing + # no real query rows. Give that tile an empty K domain so every task + # skips it before any Q/K/V or partial-output access. + _, q_len = query_batch_bounds( + self.cu_seqlens_q, + safe_b_idx, + self.logical_seq_len_q, + ) + query_tile_has_rows = runtime_flat_query_tile_has_rows( + s_idx, + cfg.mma_qk_tiler[0], + self.logical_num_heads_q, + self.logical_seq_len_q, + self.cu_seqlens_q, + safe_b_idx, + ) + if cutlass.const_expr( + cfg.mask_type == MaskType.CAUSAL.value and self.logical_seq_len_q > 1 + ): + # Split partitioning owns the mask-visible CTA domain. The final + # physical row resolves to the tile's latest valid Q + # token. Row-causal softmax narrows earlier rows without changing + # producer split geometry. + _, _, logical_q_idx, _, _ = flat_query_row_state( + Int32(cfg.mma_qk_tiler[0] - 1), + s_idx, + cfg.mma_qk_tiler[0], + self.logical_num_heads_q, + self.logical_seq_len_q, + self.cu_seqlens_q, + safe_b_idx, + ) + K = mask_visible_k_length(cfg.mask_type, K, logical_q_idx, q_len) + K = K if query_tile_has_rows else Int32(0) + k_tile_total = (K + Int32(cfg.mma_qk_tiler[1] - 1)) // Int32( + cfg.mma_qk_tiler[1] + ) + if cutlass.const_expr(self.static_split_kv == 1): + # Single-split fixed profiles process the whole K domain in one CTA. + # Avoid the runtime split-KV ceil-div/min/max sequence in every + # persistent task branch. + k_index_base = Int32(0) + k_tile_count = k_tile_total + else: + # Grid/workspace geometry stays at ``split_kv`` while each batch's + # optional cap and valid K select the configured-span nonempty prefix. + if cutlass.const_expr( + self.static_split_kv is not None and not self.is_var_split_kv + ): + split_kv_cap = Int32(self.static_split_kv) + else: + split_kv_cap = runtime_split_kv_cap( + self.split_kv, + self.is_var_split_kv, + self.block_split_kvs, + safe_b_idx, + ) + k_index_base, k_tile_count = runtime_split_tile_range( + k_tile_total, + split_kv_cap, + split_kv_idx, + ) + return MlaTsWorkTileInfo( + tile_idx, + is_valid, + sequence_k_len, + k_tile_count, + k_index_base, + ) + + @cute.jit + def k_tile_count_for_tile(self, tile_idx): + """Return the dynamic K-loop bound required by stock Task.""" + + return self._make_work_tile_info(tile_idx, Boolean(True)).k_tile_count + + @cute.jit + def skip_work_tile_if(self, work_tile): + """Skip zero-K CLC tiles while retaining WorkQueue bookkeeping.""" + + return work_tile.k_tile_count <= Int32(0) + + @cute.jit + def _work_tile_from_linear_idx(self, current_work_linear_idx): + params = self.tile_sched_params + current_work_cluster_batch, cluster_idx = ( + current_work_linear_idx // params.cluster_shape_mnk[0], + current_work_linear_idx % params.cluster_shape_mnk[0], + ) + current_work_s_batch, s_idx = divmod_constexpr_power_of_two_or_fdd( + current_work_cluster_batch, + self.static_problem_shape_s, + params.problem_shape_s_fdd, + ) + current_work_b_batch, b_idx = divmod_constexpr_power_of_two_or_fdd( + current_work_s_batch, + self.static_problem_shape_b, + params.problem_shape_b_fdd, + ) + if cutlass.const_expr(self.static_split_kv == 1): + split_kv_idx = Int32(0) + num_blocks = ( + params.cluster_shape_mnk[0] + * params.problem_shape_s + * params.problem_shape_b + ) + else: + _, split_kv_idx = divmod( + current_work_b_batch, + params.split_kv_fdd, + ) + num_blocks = ( + params.cluster_shape_mnk[0] + * params.problem_shape_s + * params.problem_shape_b + * params.split_kv + ) + return self._make_work_tile_info( + (cluster_idx, s_idx, b_idx, split_kv_idx), + current_work_linear_idx < num_blocks, + ) + + @cute.jit + def _work_tile_from_block_idx(self, block_idx): + params = self.tile_sched_params + s_idx, b_idx = divmod_constexpr_power_of_two_or_fdd( + block_idx[1], + self.static_problem_shape_b, + params.problem_shape_b_fdd, + ) + return self._make_work_tile_info( + (block_idx[0], s_idx, b_idx, block_idx[2]), + Boolean(True), + ) + + @cute.jit + def _linear_idx_from_tile(self, tile_idx): + params = self.tile_sched_params + cluster_idx, s_idx, b_idx, split_kv_idx = tile_idx + return ( + ((split_kv_idx * params.problem_shape_b + b_idx) * params.problem_shape_s) + + s_idx + ) * params.cluster_shape_mnk[0] + cluster_idx + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=(work_tile, skip_work_tile)) + @cute.jit + def init_work_tile(self, stage_info: StageInfo): + """Seed captured schedules from the current custom MLA work tile.""" + return stage_info.work_tile, Boolean(False) + + @consumer_work(returns=work_tile) + @cute.jit + def advance_tile(self, stage_info: StageInfo): + """Advance CLC through its response pipeline; static is task-owned.""" + if cutlass.const_expr(self.use_clc_dynamic): + # Decode through the public scheduler/config surfaces so the + # task-local value keeps its MLA type and cached runtime K-domain + # fields without depending on WorkQueue's private response helpers. + assert self.tile_scheduler_config is not None + assert self.tile_scheduler_config.response_ptr is not None + assert self.tile_scheduler is not None + assert self.pipeline_config is not None + stage_response_ptr = self.tile_scheduler_config.response_ptr + if cutlass.const_expr(self.pipeline_config.num_stages > 1): + stage_response_ptr = stage_response_ptr + stage_info.stage_idx + work_tile = self.tile_scheduler.work_tile_info_from_clc_response( + stage_response_ptr + ) + return self._work_tile_from_clc_tile(work_tile) + return stage_info.work_tile + + +# ===================================================================== +# PageOffsetWindowResource — TMA-warp register page-table window +# ===================================================================== + + +@dataclass(kw_only=True) +class PageOffsetWindowResource(HighThroughputMlaResource): + """GMEM page table cached as one register per lane of the TMA warp.""" + + page_offsets: Any = None # GMEM page-offset tensor + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cached_k_pages: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + cached_v_pages: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + cached_next_v_pages: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + cached_window_page: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "cached_k_pages", + cutlass.Array, + None, + "Cached page ids for the current K tile.", + ), + ( + "cached_v_pages", + cutlass.Array, + None, + "Cached page ids for the delayed V tile.", + ), + ( + "cached_next_v_pages", + cutlass.Array, + None, + "Cached page ids for the next delayed V tile.", + ), + ( + "cached_window_page", + Int32, + Int32(0), + "Per-lane page id retained across one 32-entry page-table window.", + ), + ) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + cached_window_page, + ), + ) + @cute.jit + def init_read_state(self, stage_info: StageInfo): + """Create cached page-index arrays used by staged KV TMA loads.""" + del stage_info + cfg = self.cfg + return ( + cutlass.Array( + Int32, + cfg.pages_per_k_cta, + space=cutlass.AddressSpace.rmem, + ), + cutlass.Array( + Int32, + cfg.pages_per_v_tile, + space=cutlass.AddressSpace.rmem, + ), + cutlass.Array( + Int32, + cfg.pages_per_v_tile, + space=cutlass.AddressSpace.rmem, + ), + Int32(0), + ) + + @consumer_work( + returns=( + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + cached_window_page, + ) + ) + @cute.jit + def read_page_offset_window( + self, + stage_info: StageInfo, + *, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + cached_window_page, + init_v_cache: cutlass.Constexpr[bool] = False, + ): + """Load and reuse one warp-wide 32-page-ID window. + + The TMA warp owns one page ID per lane. It refreshes the register + window after ``32 / pages_per_k_tile`` logical K tiles, when the next + group of 32 page-table entries is needed, then uses indexed warp + shuffles to assemble the IDs for the current K/V tile. The refresh + cadence is derived entirely from page and tile geometry. + """ + cfg = self.cfg + work_tile = stage_info.work_tile + blk_coord = work_tile.tile_idx + local_k_index = Int32(stage_info.loop_offset) + global_k_index = work_tile.k_index_base + local_k_index + pages_per_k_tile = cutlass.const_expr(cfg.pages_per_k_tile) + page_window_tiles = cutlass.const_expr(32 // pages_per_k_tile) + page_window_mask = cutlass.const_expr(page_window_tiles - 1) + page_offsets_batch = self.page_offsets[None, blk_coord[2]] + lane_idx = cute.arch.thread_idx()[0] & Int32(31) + + if (local_k_index & Int32(page_window_mask)) == Int32(0): + logical_page_idx = global_k_index * Int32(pages_per_k_tile) + lane_idx + bounded_page_idx = cute.math.min( + logical_page_idx, Int32(page_offsets_batch.shape[0] - 1) + ) + cached_window_page = Int32(page_offsets_batch[bounded_page_idx]) + + page_lane_base = (local_k_index & Int32(page_window_mask)) * Int32( + pages_per_k_tile + ) + cta_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + cached_k = cached_k_pages + cached_v = cached_v_pages + cached_next_v = cached_next_v_pages + + if cutlass.const_expr(init_v_cache): + for pk in cutlass.range_constexpr(cfg.pages_per_v_tile): + cached_v[pk] = Int32(0) + else: + for pk in cutlass.range_constexpr(cfg.pages_per_v_tile): + cached_v[pk] = cached_next_v[pk] + + for pk in cutlass.range_constexpr(cfg.pages_per_k_cta): + source_lane = ( + page_lane_base + cta_v * Int32(cfg.pages_per_k_cta) + Int32(pk) + if cfg.pages_per_k_tile > 1 + else page_lane_base + ) + cached_k[pk] = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_window_page, + offset=source_lane, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + + for pk in cutlass.range_constexpr(cfg.pages_per_v_tile): + source_lane = ( + page_lane_base + if cfg.pages_per_v_tile == 1 + else page_lane_base + Int32(pk) + ) + cached_next_v[pk] = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_window_page, + offset=source_lane, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + return cached_k, cached_v, cached_next_v, cached_window_page + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(cached_k_pages, cached_v_pages, cached_next_v_pages), + ) + @cute.jit + def forward_page_ids( + self, + stage_info: StageInfo, + *, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + ): + """Forward the latest delayed-V page IDs after the domain loop.""" + del stage_info + return cached_k_pages, cached_v_pages, cached_next_v_pages + + +# ===================================================================== +# SmemQResource — Q SMEM buffer with TmaUmmaAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemQResource(HighThroughputMlaResource): + """SMEM Q buffer (latent + rope). Producer: LoadTma (TMA). Consumer: Mma. + + Pipeline: TmaUmmaAsync, 1 stage. + Q is loaded once (LoopFirstIter) and released once (LoopLastIter). + """ + + smem_q_latent: Any = None # SMEM Q latent array + smem_q_rope: Any = None # SMEM Q rope array + tma_desc_q_latent: Any = None + tma_desc_q_rope: Any = None + cu_seqlens_q: Any = None + logical_num_heads_q: cutlass.Constexpr[int] = 128 + logical_seq_len_q: cutlass.Constexpr[int] = 1 + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cta_rank: Any = field(init=False, default=None) + is_leader: Any = field(init=False, default=None) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize CTA-local Q-load state for the high-throughput path.""" + del stage_info + self.cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + self.is_leader = self.cta_rank == 0 + + @cute.jit + def _query_tma_coords( + self, + dim_coord, + local_flat_query_row, + batch_idx, + storage_flat_query_row, + query_row_extent, + ): + """Return dense or ragged TMA coordinates for one Q dimension slice.""" + if cutlass.const_expr(self.cu_seqlens_q is not None): + return transform_ragged_coords( + (dim_coord, storage_flat_query_row), + ragged_dim_idx=1, + ragged_box_size=self.cfg.mma_qk_tiler[0] // self.cfg.num_mma_ctas, + ragged_extent=query_row_extent, + ) + return dim_coord, local_flat_query_row, batch_idx + + @producer_work + @cute.jit + def tma_load(self, stage_info: StageInfo) -> None: + """TMA load Q latent + Q rope into SMEM (all sub-tiles in one commit).""" + cfg = self.cfg + work_tile = stage_info.work_tile + blk_coord = work_tile.tile_idx + + cta_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + coord_m = cta_v * Int32(cfg.mma_qk_tiler[0] // cfg.num_mma_ctas) + physical_tile_rows = Int32(cfg.mma_qk_tiler[0]) + flat_query_row = Int32(blk_coord[1]) * physical_tile_rows + coord_m + batch_idx = Int32(blk_coord[2]) + storage_flat_query_row = flat_query_row + query_row_extent = Int32(cfg.mma_qk_tiler[0] // cfg.num_mma_ctas) + if cutlass.const_expr(self.cu_seqlens_q is not None): + q_start, q_len = query_batch_bounds( + self.cu_seqlens_q, + batch_idx, + self.logical_seq_len_q, + ) + storage_flat_query_row = ( + q_start * Int32(self.logical_num_heads_q) + flat_query_row + ) + query_row_extent = q_len * Int32(self.logical_num_heads_q) - flat_query_row + mask_q = Int16(Int32(1) << cta_v) + q_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) + + if prims.elect_sync(): + # Load Q latent sub-tiles + q_latent_stage_elems = cutlass.const_expr( + cfg.mma_qk_tiler[0] // cfg.num_mma_ctas * cfg.mma_qk_tiler_k + ) + for i in cutlass.range(cfg.iterations_qk_latent): + coord_kl = cutlass.Int32(i * cfg.mma_qk_tiler_k) + q_latent_coords = self._query_tma_coords( + coord_kl, + flat_query_row, + batch_idx, + storage_flat_query_row, + query_row_extent, + ) + q_smem_arr = cutlass.Array( + self.smem_q_latent.data_ptr(i * q_latent_stage_elems), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + q_smem_arr, + self.tma_desc_q_latent, + q_latent_coords, + q_mbar_arr, + [], + multicast_mask=mask_q, + group=prims.CTAGroup.CTA_2, + ) + # Load Q rope sub-tiles + for i in cutlass.range(cfg.iterations_qk_rope): + coord_kr = cutlass.Int32(i * cfg.mma_qk_tiler_k) + q_rope_coords = self._query_tma_coords( + coord_kr, + flat_query_row, + batch_idx, + storage_flat_query_row, + query_row_extent, + ) + qr_smem_arr = cutlass.Array( + self.smem_q_rope.data_ptr(), dtype=qkv_dtype(cfg) + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + qr_smem_arr, + self.tma_desc_q_rope, + q_rope_coords, + q_mbar_arr, + [], + multicast_mask=mask_q, + group=prims.CTAGroup.CTA_2, + ) + + @consumer_work + @cute.jit + def q_desc(self, stage_info: StageInfo) -> None: + """Schedule marker after the Q SMEM stage is waited.""" + del stage_info + + +# ===================================================================== +# SmemKVResource — K/V SMEM buffer with TmaUmmaAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemKVResource(HighThroughputMlaResource): + """SMEM K/V buffer. Producer: LoadTma (TMA). Consumer: Mma. + + Pipeline: TmaUmmaAsync, 7 stages. + Per logical k-tile: K sub-tiles are loaded before delayed V sub-tiles. + """ + + smem_kv: Any = None + tma_desc_c_latent: Any = None + tma_desc_c_rope: Any = None + tma_desc_c_transpose: Any = None + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cta_rank: Any = field(init=False, default=None) + is_leader: Any = field(init=False, default=None) + desc_k_base: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + desc_v_base: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("desc_k_base", Int64, Int64(0), "SMEM descriptor for staged K."), + ("desc_v_base", Int64, Int64(0), "SMEM descriptor for staged V."), + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize CTA-local KV-load descriptors and leader state.""" + del stage_info + self.cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + self.is_leader = self.cta_rank == 0 + + @producer_work + @cute.jit + def tma_load( + self, + stage_info: StageInfo, + *, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + is_v: cutlass.Constexpr[bool], + subtile_idx: cutlass.Constexpr[int], + use_next_v_pages: cutlass.Constexpr[bool] = False, + ) -> None: + """TMA load one K or delayed-V sub-tile into the shared KV ring.""" + cfg = self.cfg + stage_idx = stage_info.stage_idx + + kc_page_smem_elems = cfg.kc_page_tile_size * cfg.mma_qk_tiler_k + kv_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) + kv_stage_stride_elems = cutlass.const_expr(cfg.smem_k_stage_elems) + + pages_per_k_cta = cfg.pages_per_k_cta + + cta_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + coord_n_k = (cta_v * Int32(cfg.mma_qk_tiler[1] // cfg.num_mma_ctas)) % Int32( + cfg.page_size + ) + mask_k = Int16(Int32(1) << cta_v) + + is_v_subtile = is_v + k_call = subtile_idx + + if cutlass.const_expr( + not is_v_subtile and k_call < cfg.iterations_qk_latent_stages + ): + cached_k = cached_k_pages + k_subtile_smem_elems = cutlass.const_expr( + cfg.mma_qk_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_tiler_k + ) + for stage_subtile_idx in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + logical_k_call = k_call * cfg.kv_subtiles_per_stage + stage_subtile_idx + coord_kcl = cutlass.Int32(logical_k_call * cfg.mma_qk_tiler_k) + for pk in cutlass.range_constexpr(pages_per_k_cta): + if prims.elect_sync(): + kcl_smem = cutlass.Array( + self.smem_kv.data_ptr( + stage_idx * kv_stage_stride_elems + + stage_subtile_idx * k_subtile_smem_elems + + pk * kc_page_smem_elems + ), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + kcl_smem, + self.tma_desc_c_latent, + (coord_kcl, coord_n_k, cutlass.Int32(cached_k[pk])), + kv_mbar_arr, + [], + multicast_mask=mask_k, + group=prims.CTAGroup.CTA_2, + ) + + elif cutlass.const_expr(not is_v_subtile and k_call < cfg.iterations_qk_stages): + rope_idx = k_call - cfg.iterations_qk_latent_stages + cached_k = cached_k_pages + + coord_kcr = cutlass.Int32(rope_idx * cfg.mma_qk_rope_tiler[2]) + rope_page_smem_elems = cfg.kc_page_tile_size * cfg.mma_qk_rope_tiler[2] + rope_tile_smem_elems = ( + cfg.mma_qk_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_rope_tiler[2] + ) + for pk in cutlass.range_constexpr(pages_per_k_cta): + if prims.elect_sync(): + kcr_smem = cutlass.Array( + self.smem_kv.data_ptr( + stage_idx * kv_stage_stride_elems + + pk * rope_page_smem_elems + ), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + kcr_smem, + self.tma_desc_c_rope, + (coord_kcr, coord_n_k, cutlass.Int32(cached_k[pk])), + kv_mbar_arr, + [], + multicast_mask=mask_k, + group=prims.CTAGroup.CTA_2, + ) + if cutlass.const_expr(cfg.kv_subtiles_per_stage > 1): + # The BF16 stage contains two legal K64 transfers. + # Duplicate the final K64 RoPE slice into the unused + # half so its mbarrier transaction count matches the + # common K128 stage contract. + kcr_smem_dup = cutlass.Array( + self.smem_kv.data_ptr( + stage_idx * kv_stage_stride_elems + + rope_tile_smem_elems + + pk * rope_page_smem_elems + ), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + kcr_smem_dup, + self.tma_desc_c_rope, + (coord_kcr, coord_n_k, cutlass.Int32(cached_k[pk])), + kv_mbar_arr, + [], + multicast_mask=mask_k, + group=prims.CTAGroup.CTA_2, + ) + + else: + pv_n_per_cta = cutlass.const_expr(cfg.mma_pv_tiler[1] // cfg.num_mma_ctas) + coord_n_v = cta_v * pv_n_per_cta + mask_v = Int16(Int32(1) << cta_v) + v_tma_copy_smem_elems = cfg.v_tma_token_count * V_TMA_LATENT_ELEMENTS + v_subtile_smem_elems = cutlass.const_expr( + cfg.mma_pv_tiler[1] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] + ) + + pages_per_v_subtile = cfg.pages_per_v_subtile + cached_v = ( + cached_next_v_pages + if cutlass.const_expr(use_next_v_pages) + else cached_v_pages + ) + + # A physical V stage contains two adjacent K32 slices for one + # D256 output panel. Across four stages this is the same + # head-dimension-128/token-partition-2 decomposition used by the + # 2CTA BF16 schedule: (D0,K0:64), (D0,K64:128), then D256. + pv_j = subtile_idx // cfg.kv_subtiles_per_stage + token_partition = subtile_idx % cfg.kv_subtiles_per_stage + for stage_subtile_idx in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + pv_i = token_partition * cfg.kv_subtiles_per_stage + stage_subtile_idx + coord_k_v = cutlass.Int32((pv_i * cfg.mma_pv_tiler[2]) % cfg.page_size) + coord_nj = coord_n_v + cutlass.Int32(pv_j * cfg.mma_pv_tiler[1]) + + for pk in cutlass.range_constexpr(pages_per_v_subtile): + k_idx_i = cached_v[ + pk + pv_i // cfg.v_subtiles_per_page * pages_per_v_subtile + ] + if prims.elect_sync(): + # Keep both 64-wide V panels contiguous within each + # K32 slice; the next K32 slice follows the complete + # first slice in this physical stage. + stage_subtile_base = stage_subtile_idx * v_subtile_smem_elems + v_page_offset = pk * v_tma_copy_smem_elems + v_second_panel_offset = ( + pages_per_v_subtile * v_tma_copy_smem_elems + v_page_offset + ) + v_smem_0 = cutlass.Array( + self.smem_kv.data_ptr( + stage_idx * kv_stage_stride_elems + + stage_subtile_base + + v_page_offset + ), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + v_smem_0, + self.tma_desc_c_transpose, + (coord_nj, coord_k_v, k_idx_i), + kv_mbar_arr, + [], + multicast_mask=mask_v, + group=prims.CTAGroup.CTA_2, + ) + v_smem_1 = cutlass.Array( + self.smem_kv.data_ptr( + stage_idx * kv_stage_stride_elems + + stage_subtile_base + + v_second_panel_offset + ), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + v_smem_1, + self.tma_desc_c_transpose, + ( + coord_nj + V_TMA_LATENT_ELEMENTS, + coord_k_v, + k_idx_i, + ), + kv_mbar_arr, + [], + multicast_mask=mask_v, + group=prims.CTAGroup.CTA_2, + ) + + @consumer_work(returns=desc_k_base) + @cute.jit + def k_desc(self, stage_info: StageInfo, *, k_subtile_idx: cutlass.Constexpr[int]): + """Build the SMEM descriptor consumed by QK MMA.""" + cfg = self.cfg + stage_idx = stage_info.stage_idx + + kc_copy_elems = cutlass.const_expr(cfg.smem_k_stage_elems) + tile_rows = cutlass.const_expr(cfg.mma_qk_tiler[1] // cfg.num_mma_ctas) + leading_byte_offset = cutlass.const_expr(qk_desc_leading_byte_offset(cfg)) + stride_byte_offset = cutlass.const_expr(qk_desc_stride_byte_offset(cfg)) + layout = cutlass.const_expr(qk_desc_layout(cfg)) + if cutlass.const_expr( + cfg.is_fp8_qkv() + and k_subtile_idx >= cfg.iterations_qk_latent + and k_subtile_idx < cfg.iterations_qk + ): + leading_byte_offset = cutlass.const_expr( + qk_desc_leading_byte_offset_for_head_dim( + cfg, tile_rows, cfg.mma_qk_rope_tiler[2] + ) + ) + stride_byte_offset = cutlass.const_expr( + qk_desc_stride_byte_offset_for_head_dim(cfg, cfg.mma_qk_rope_tiler[2]) + ) + layout = cutlass.const_expr( + qk_desc_layout_for_head_dim(cfg, cfg.mma_qk_rope_tiler[2]) + ) + + sk_ptr = self.smem_kv.data_ptr(stage_idx * kc_copy_elems) + desc = Int64( + prims.Tcgen05SmemDesc.build( + start_address=sk_ptr.toint(Int32), + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=layout, + ) + ) + return desc + + @consumer_work(returns=desc_v_base) + @cute.jit + def v_desc(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int]): + """Build the SMEM descriptor consumed by PV MMA.""" + del v_subtile_idx + cfg = self.cfg + stage_idx = stage_info.stage_idx + svc_copy_elems = cutlass.const_expr(cfg.smem_k_stage_elems) + leading_byte_offset = cutlass.const_expr(4096) + stride_byte_offset = cutlass.const_expr(1024) + layout = cutlass.const_expr(2) + if cutlass.const_expr(cfg.is_fp8_qkv()): + leading_byte_offset = cutlass.const_expr( + V_SMEM_K_BLOCK_TOKENS * V_TMA_LATENT_ELEMENTS * cfg.qkv_dtype_bytes + ) + stride_byte_offset = cutlass.const_expr( + qkv_major_k_stride_bytes_for(cfg, cfg.mma_pv_tiler[2]) + ) + layout = cutlass.const_expr( + qk_desc_layout_for_head_dim(cfg, cfg.mma_pv_tiler[2]) + ) + + svc_ptr = self.smem_kv.data_ptr(stage_idx * svc_copy_elems) + desc = Int64( + prims.Tcgen05SmemDesc.build( + start_address=svc_ptr.toint(Int32), + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=layout, + ) + ) + return desc + + +# ===================================================================== +# SmemKResource — FP8 K SMEM buffer with one pipeline stage per K tile +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemKResource(HighThroughputMlaResource): + """SMEM K buffer for the FP8 split-MMA path. + + Producer: LoadTma. Consumer: MmaQkTask. A single producer stage contains + all latent K sub-tiles plus the RoPE K sub-tile for one logical K tile. + """ + + smem_k: Any = None + page_offsets: Any = None + tma_desc_c_latent: Any = None + tma_desc_c_rope: Any = None + logical_seq_len_q: cutlass.Constexpr[int] = 1 + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + desc_k_base: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("desc_k_base", Int64, Int64(0), "SMEM descriptor for staged K."), + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize per-warp state used by K TMA loads.""" + del stage_info + + @producer_work + @cute.jit + def tma_load_direct(self, stage_info: StageInfo) -> None: + """TMA load one full K tile using page offsets read directly from GMEM.""" + cfg = self.cfg + stage_idx = stage_info.stage_idx + stage_base = stage_idx * cfg.smem_k_stage_elems + work_tile = stage_info.work_tile + blk_coord = work_tile.tile_idx + k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) + + page_row_idx = blk_coord[2] + page_offsets_batch = self.page_offsets[None, page_row_idx] + kv_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) + pages_per_k_cta = cfg.pages_per_k_cta + kc_page_smem_elems = cfg.kc_page_tile_size * cfg.mma_qk_tiler_k + cta_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + coord_n_k = (cta_v * Int32(cfg.mma_qk_tiler[1] // cfg.num_mma_ctas)) % Int32( + cfg.page_size + ) + mask_k = Int16(Int32(1) << cta_v) + k_latent_subtile_elems = cutlass.const_expr( + cfg.mma_qk_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_tiler_k + ) + # Use page 0 for masked fragments beyond a compact page table. A direct + # bounds predicate avoids carrying runtime K-length division or a full + # page-ID register array through this producer warp. + + for k_call in cutlass.range_constexpr(cfg.iterations_qk_latent): + coord_kcl = cutlass.Int32(k_call * cfg.mma_qk_tiler_k) + subtile_base = stage_base + k_call * k_latent_subtile_elems + for pk in cutlass.range_constexpr(pages_per_k_cta): + logical_page_idx = ( + k_index + if cfg.pages_per_k_tile == 1 + else (k_index * Int32(cfg.num_mma_ctas) + cta_v) + * Int32(pages_per_k_cta) + + Int32(pk) + ) + page_idx = Int32(0) + if cute.elem_less(logical_page_idx, page_offsets_batch.shape[0]): + page_idx = page_offsets_batch[logical_page_idx] + if prims.elect_sync(): + kcl_smem = cutlass.Array( + self.smem_k.data_ptr(subtile_base + pk * kc_page_smem_elems), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + kcl_smem, + self.tma_desc_c_latent, + (coord_kcl, coord_n_k, cutlass.Int32(page_idx)), + kv_mbar_arr, + [], + multicast_mask=mask_k, + group=prims.CTAGroup.CTA_2, + ) + + rope_stage_base = stage_base + k_latent_subtile_elems * cfg.iterations_qk_latent + rope_page_smem_elems = cfg.kc_page_tile_size * cfg.mma_qk_rope_tiler[2] + k_rope_subtile_elems = cutlass.const_expr( + cfg.mma_qk_rope_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_rope_tiler[2] + ) + for rope_idx in cutlass.range_constexpr(cfg.iterations_qk_rope): + coord_kcr = cutlass.Int32(rope_idx * cfg.mma_qk_rope_tiler[2]) + subtile_base = rope_stage_base + rope_idx * k_rope_subtile_elems + for pk in cutlass.range_constexpr(pages_per_k_cta): + logical_page_idx = ( + k_index + if cfg.pages_per_k_tile == 1 + else (k_index * Int32(cfg.num_mma_ctas) + cta_v) + * Int32(pages_per_k_cta) + + Int32(pk) + ) + page_idx = Int32(0) + if cute.elem_less(logical_page_idx, page_offsets_batch.shape[0]): + page_idx = page_offsets_batch[logical_page_idx] + if prims.elect_sync(): + kcr_smem = cutlass.Array( + self.smem_k.data_ptr(subtile_base + pk * rope_page_smem_elems), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + kcr_smem, + self.tma_desc_c_rope, + (coord_kcr, coord_n_k, cutlass.Int32(page_idx)), + kv_mbar_arr, + [], + multicast_mask=mask_k, + group=prims.CTAGroup.CTA_2, + ) + + @consumer_work(returns=desc_k_base) + @cute.jit + def k_desc(self, stage_info: StageInfo, *, k_subtile_idx: cutlass.Constexpr[int]): + """Build the K SMEM descriptor for the current QK sub-MMA.""" + cfg = self.cfg + stage_idx = stage_info.stage_idx + + k_latent_subtile_elems = cutlass.const_expr( + cfg.mma_qk_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_tiler_k + ) + k_rope_subtile_elems = cutlass.const_expr( + cfg.mma_qk_rope_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_rope_tiler[2] + ) + subtile_offset = stage_idx * cfg.smem_k_stage_elems + tile_rows = cutlass.const_expr(cfg.mma_qk_tiler[1] // cfg.num_mma_ctas) + leading_byte_offset = cutlass.const_expr(qk_desc_leading_byte_offset(cfg)) + stride_byte_offset = cutlass.const_expr(qk_desc_stride_byte_offset(cfg)) + layout = cutlass.const_expr(qk_desc_layout(cfg)) + + if cutlass.const_expr(k_subtile_idx < cfg.iterations_qk_latent): + subtile_offset += k_subtile_idx * k_latent_subtile_elems + else: + rope_idx = k_subtile_idx - cfg.iterations_qk_latent + subtile_offset += ( + k_latent_subtile_elems * cfg.iterations_qk_latent + + rope_idx * k_rope_subtile_elems + ) + leading_byte_offset = cutlass.const_expr( + qk_desc_leading_byte_offset_for_head_dim( + cfg, tile_rows, cfg.mma_qk_rope_tiler[2] + ) + ) + stride_byte_offset = cutlass.const_expr( + qk_desc_stride_byte_offset_for_head_dim(cfg, cfg.mma_qk_rope_tiler[2]) + ) + layout = cutlass.const_expr( + qk_desc_layout_for_head_dim(cfg, cfg.mma_qk_rope_tiler[2]) + ) + + sk_ptr = self.smem_k.data_ptr(subtile_offset) + desc = Int64( + prims.Tcgen05SmemDesc.build( + start_address=sk_ptr.toint(Int32), + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=layout, + ) + ) + return desc + + +# ===================================================================== +# SmemVResource — FP8 V SMEM buffer with one pipeline stage per V tile +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemVResource(HighThroughputMlaResource): + """SMEM V buffer for the FP8 split-MMA path. + + Producer: LoadTma. Consumer: MmaPvTask. A single stage contains every V + sub-tile needed by the PV MMA for one logical K tile. + """ + + smem_v: Any = None + page_offsets: Any = None + tma_desc_c_transpose: Any = None + logical_seq_len_q: cutlass.Constexpr[int] = 1 + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + desc_v_base: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("desc_v_base", Int64, Int64(0), "SMEM descriptor for staged V."), + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize per-warp state used by V TMA loads.""" + del stage_info + + @producer_work + @cute.jit + def tma_load_direct(self, stage_info: StageInfo) -> None: + """TMA load one full V tile using page offsets read directly from GMEM.""" + cfg = self.cfg + stage_idx = stage_info.stage_idx + stage_base = stage_idx * cfg.smem_v_stage_elems + work_tile = stage_info.work_tile + blk_coord = work_tile.tile_idx + k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) + page_row_idx = blk_coord[2] + page_offsets_batch = self.page_offsets[None, page_row_idx] + + v_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) + cta_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + pv_n_per_cta = cutlass.const_expr(cfg.mma_pv_tiler[1] // cfg.num_mma_ctas) + coord_n_v = cta_v * pv_n_per_cta + mask_v = Int16(Int32(1) << cta_v) + pages_per_v_tile = cfg.pages_per_v_tile + v_smem_panel_elems = V_SMEM_K_BLOCK_TOKENS * V_TMA_LATENT_ELEMENTS + v_smem_k_block_elems = cfg.num_mma_ctas * v_smem_panel_elems + svc_copy_elems = cutlass.const_expr( + cfg.mma_pv_tiler[1] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] + ) + # As for K, out-of-table fragments use page 0 and are masked by the + # runtime sequence length. p128 still shares one page ID across CTAs. + + for v_call in cutlass.range_constexpr( + cfg.iterations_pv_k * cfg.iterations_pv_n + ): + pv_i = v_call // cfg.iterations_pv_n + pv_j = v_call % cfg.iterations_pv_n + coord_nj = coord_n_v + cutlass.Int32(pv_j * cfg.mma_pv_tiler[1]) + subtile_base = stage_base + v_call * svc_copy_elems + + # Assemble the fixed K32 SMEM blocks consumed by tcgen05 from + # page-bounded TMA copies. Physical pages only select the GMEM + # page/coordinate; they never change the SMEM descriptor layout. + for copy_idx in cutlass.range_constexpr(cfg.v_tma_copies_per_subtile): + local_token_offset = copy_idx * cfg.v_tma_token_count + tile_token_offset = pv_i * cfg.mma_pv_tiler[2] + local_token_offset + page_offset = tile_token_offset // cfg.page_size + coord_k_v = Int32(tile_token_offset % cfg.page_size) + logical_page_idx = ( + k_index + if pages_per_v_tile == 1 + else k_index * Int32(pages_per_v_tile) + Int32(page_offset) + ) + page_idx = Int32(0) + if cute.elem_less(logical_page_idx, page_offsets_batch.shape[0]): + page_idx = page_offsets_batch[logical_page_idx] + if prims.elect_sync(): + smem_k_block_idx = local_token_offset // V_SMEM_K_BLOCK_TOKENS + token_offset_in_k_block = local_token_offset % V_SMEM_K_BLOCK_TOKENS + v_copy_offset = ( + smem_k_block_idx * v_smem_k_block_elems + + token_offset_in_k_block * V_TMA_LATENT_ELEMENTS + ) + v_smem_0 = cutlass.Array( + self.smem_v.data_ptr(subtile_base + v_copy_offset), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + v_smem_0, + self.tma_desc_c_transpose, + (coord_nj, coord_k_v, page_idx), + v_mbar_arr, + [], + multicast_mask=mask_v, + group=prims.CTAGroup.CTA_2, + ) + v_smem_1 = cutlass.Array( + self.smem_v.data_ptr( + subtile_base + v_copy_offset + v_smem_panel_elems + ), + dtype=qkv_dtype(cfg), + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + v_smem_1, + self.tma_desc_c_transpose, + ( + coord_nj + V_TMA_LATENT_ELEMENTS, + coord_k_v, + page_idx, + ), + v_mbar_arr, + [], + multicast_mask=mask_v, + group=prims.CTAGroup.CTA_2, + ) + + @consumer_work(returns=desc_v_base) + @cute.jit + def v_desc(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int]): + """Build the V SMEM descriptor for the current PV sub-MMA.""" + cfg = self.cfg + svc_copy_elems = cutlass.const_expr( + cfg.mma_pv_tiler[1] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] + ) + subtile_offset = ( + stage_info.stage_idx * cfg.smem_v_stage_elems + + v_subtile_idx * svc_copy_elems + ) + leading_byte_offset = cutlass.const_expr( + V_SMEM_K_BLOCK_TOKENS * V_TMA_LATENT_ELEMENTS * cfg.qkv_dtype_bytes + ) + stride_byte_offset = cutlass.const_expr( + qkv_major_k_stride_bytes_for(cfg, cfg.mma_pv_tiler[2]) + ) + layout = cutlass.const_expr( + qk_desc_layout_for_head_dim(cfg, cfg.mma_pv_tiler[2]) + ) + svc_ptr = self.smem_v.data_ptr(subtile_offset) + desc = Int64( + prims.Tcgen05SmemDesc.build( + start_address=svc_ptr.toint(Int32), + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=layout, + ) + ) + return desc + + @consumer_work(returns=desc_v_base) + @cute.jit + def v_desc_n_major( + self, + stage_info: StageInfo, + *, + pv_n_idx: cutlass.Constexpr[int], + pv_k_idx: cutlass.Constexpr[int], + ): + """Build the V descriptor when PV commits one output-N slice at a time.""" + cfg = self.cfg + pv_j = cutlass.const_expr(pv_n_idx) + pv_i = cutlass.const_expr(pv_k_idx) + v_call = pv_i * cfg.iterations_pv_n + pv_j + svc_copy_elems = cutlass.const_expr( + cfg.mma_pv_tiler[1] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] + ) + subtile_offset = stage_info.stage_idx * cfg.smem_v_stage_elems + ( + v_call * svc_copy_elems + ) + leading_byte_offset = cutlass.const_expr( + V_SMEM_K_BLOCK_TOKENS * V_TMA_LATENT_ELEMENTS * cfg.qkv_dtype_bytes + ) + stride_byte_offset = cutlass.const_expr( + qkv_major_k_stride_bytes_for(cfg, cfg.mma_pv_tiler[2]) + ) + layout = cutlass.const_expr( + qk_desc_layout_for_head_dim(cfg, cfg.mma_pv_tiler[2]) + ) + svc_ptr = self.smem_v.data_ptr(subtile_offset) + desc = Int64( + prims.Tcgen05SmemDesc.build( + start_address=svc_ptr.toint(Int32), + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=layout, + ) + ) + return desc + + +# ===================================================================== +# TmemSResource — S scores in TMEM, UmmaProducerAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemSResource(HighThroughputMlaResource): + """TMEM S scores. Producer: MmaTask (QK MMA). Consumer: SoftmaxTask. + + Pipeline: UmmaProducerAsync, 2 stages. + Producer call_idx 0..iterations_qk-1: issue QK MMA for each K sub-tile. + Consumer: load S, apply mask, compute softmax, and stage P for PV MMA. + """ + + tmem_base_addr: Any = None # TMEM base address (from alloc) + smem_q_latent: Any = None # SMEM Q pointers for descriptor building + smem_q_rope: Any = None + smem_p: Any = None # SMEM P array + smem_exchange: Any = None # SMEM array for cross-warp max exchange + softmax_scale_log2: Any = None # softmax_scale * log2(e) + cache_seqs: Any = None # per-batch valid K length + cu_seqlens_q: Any = None # cumulative compact-Q offsets, or None for fixed Q + split_kv: Any = None # per-work-tile split count + logical_num_heads_q: cutlass.Constexpr[int] = 128 + logical_seq_len_q: cutlass.Constexpr[int] = 1 + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + tiled_mma_qk: Any = None + cta_rank: Any = field(init=False, default=None) + is_leader: Any = field(init=False, default=None) + row_max_state: Any = field(init=False, default=None) + row_sum_state: Any = field(init=False, default=None) + qk_acc_regs: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + row_max: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + row_sum: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + row_sum_out: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + row_max_new: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + correction_factor_out: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + no_correction_out: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + qk_acc_regs_odd: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + row_max_odd: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + row_sum_odd: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + row_sum_out_odd: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + row_max_new_odd: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + correction_factor_out_odd: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + no_correction_out_odd: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("qk_acc_regs", cutlass.Array, None, "Registers holding softmax P values."), + ("row_max", Float32, Float32(-Float32.inf), "Running row maximum."), + ("row_sum", Float32, Float32(0), "Running row sum."), + ("row_sum_out", Float32, Float32(0), "Row sum published to correction."), + ("row_max_new", Float32, Float32(0), "Updated row maximum."), + ( + "correction_factor_out", + Float32, + Float32(0), + "Correction factor for the previous O tile.", + ), + ( + "no_correction_out", + Int32, + Int32(0), + "Whether O correction may be skipped.", + ), + ( + "qk_acc_regs_odd", + cutlass.Array, + None, + "Odd-lane registers holding softmax P values.", + ), + ("row_max_odd", Float32, Float32(-Float32.inf), "Odd-lane row maximum."), + ("row_sum_odd", Float32, Float32(0), "Odd-lane row sum."), + ( + "row_sum_out_odd", + Float32, + Float32(0), + "Odd-lane row sum published to correction.", + ), + ("row_max_new_odd", Float32, Float32(0), "Odd-lane updated row maximum."), + ( + "correction_factor_out_odd", + Float32, + Float32(0), + "Odd-lane correction factor for the previous O tile.", + ), + ( + "no_correction_out_odd", + Int32, + Int32(0), + "Whether odd-lane O correction may be skipped.", + ), + ) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ), + ) + @cute.jit + def init_softmax_state(self, stage_info: StageInfo): + """Create softmax accumulator registers and row-stat state.""" + del stage_info + self.cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + self.is_leader = self.cta_rank == 0 + self.row_max_state = Float32(-Float32.inf) + self.row_sum_state = Float32(0) + return ( + cutlass.Array( + Float32, + 64, + space=cutlass.AddressSpace.rmem, + ), + Float32(-Float32.inf), + Float32(0), + Float32(0), + Float32(0), + Float32(0), + Int32(0), + ) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + qk_acc_regs_odd, + row_max_odd, + row_sum_odd, + row_sum_out_odd, + row_max_new_odd, + correction_factor_out_odd, + no_correction_out_odd, + ), + ) + @cute.jit + def init_softmax_state_odd(self, stage_info: StageInfo): + """Create odd-lane softmax accumulator registers and row-stat state.""" + del stage_info + self.cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + self.is_leader = self.cta_rank == 0 + self.row_max_state = Float32(-Float32.inf) + self.row_sum_state = Float32(0) + return ( + cutlass.Array( + Float32, + 64, + space=cutlass.AddressSpace.rmem, + ), + Float32(-Float32.inf), + Float32(0), + Float32(0), + Float32(0), + Float32(0), + Int32(0), + ) + + @producer_work + @cute.jit + def qk_mma( + self, + stage_info: StageInfo, + *, + desc_k_base, + k_subtile_idx: cutlass.Constexpr[int], + ) -> None: + """Issue one QK MMA sub-tile (K latent or K rope). + Only leader CTA issues MMA (2CTA UMMA principle). + + Descriptor computation is hoisted OUTSIDE the leader-CTA gate so + ptxas keeps values in uniform registers. + Only the actual tcgen05_mma is gated by elect_sync + leader check. + """ + cfg = self.cfg + call_idx = k_subtile_idx + + # Hoist descriptor computation outside leader-CTA gate to preserve + # uniform register allocation (avoids R2UR demote/promote). + tmem_s_addr = self.tmem_base_addr + 64 * stage_info.stage_idx + tmem_s_ptr = prims.make_tmem_ptr(tmem_s_addr, Float32) + + idesc_qk = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + n_dim=cfg.mma_qk_tiler[1], + m_dim=cfg.mma_qk_tiler[0], + ) + mma_kind = mma_kind_for_qkv(cfg) + cta_group = prims.CTAGroup.CTA_2 + k_block_count = cutlass.const_expr( + ceil_div(cfg.mma_qk_tiler_k, mma_k_step_for_qkv(cfg)) + ) + is_leader_cta = ( + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) == 0 + ) + + desc_k = desc_k_base + + if cutlass.const_expr(call_idx < cfg.iterations_qk_latent_stages): + qc_copy_elems = cutlass.const_expr( + cfg.mma_qk_tiler[0] // cfg.num_mma_ctas * cfg.mma_qk_tiler_k + ) + kc_copy_elems = cutlass.const_expr( + cfg.mma_qk_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_tiler_k + ) + desc_k_delta = cutlass.const_expr(kc_copy_elems * cfg.qkv_dtype_bytes // 16) + for stage_subtile_idx in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + logical_call_idx = ( + call_idx * cfg.kv_subtiles_per_stage + stage_subtile_idx + ) + q_smem_ptr = self.smem_q_latent.data_ptr( + logical_call_idx * qc_copy_elems + ) + desc_q = Int64( + prims.Tcgen05SmemDesc.build( + start_address=q_smem_ptr.toint(Int32), + leading_byte_offset=qk_desc_leading_byte_offset(cfg), + stride_byte_offset=qk_desc_stride_byte_offset(cfg), + layout=qk_desc_layout(cfg), + ) + ) + desc_k_subtile = desc_k + stage_subtile_idx * desc_k_delta + if is_leader_cta: + for k_block in cutlass.range_constexpr(k_block_count): + scale_d = cutlass.const_expr( + logical_call_idx > 0 or k_block > 0 + ) + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + cta_group, + tmem_s_ptr, + desc_q + k_block * 2, + desc_k_subtile + k_block * 2, + idesc_qk, + Boolean(scale_d), + ) + + elif cutlass.const_expr(call_idx < cfg.iterations_qk_stages): + # K rope: build Q rope descriptor unconditionally + rope_idx = call_idx - cfg.iterations_qk_latent_stages + qc_copy_elems = cutlass.const_expr( + cfg.mma_qk_tiler[0] // cfg.num_mma_ctas * cfg.mma_qk_rope_tiler[2] + ) + q_rope_smem_ptr = self.smem_q_rope.data_ptr(rope_idx * qc_copy_elems) + q_rope_rows = cutlass.const_expr(cfg.mma_qk_tiler[0] // cfg.num_mma_ctas) + q_rope_dim = cutlass.const_expr(cfg.mma_qk_rope_tiler[2]) + desc_qr = Int64( + prims.Tcgen05SmemDesc.build( + start_address=q_rope_smem_ptr.toint(Int32), + leading_byte_offset=qk_desc_leading_byte_offset_for_head_dim( + cfg, q_rope_rows, q_rope_dim + ), + stride_byte_offset=qk_desc_stride_byte_offset_for_head_dim( + cfg, q_rope_dim + ), + layout=qk_desc_layout_for_head_dim(cfg, q_rope_dim), + ) + ) + if is_leader_cta: + for k_block in cutlass.range_constexpr( + cfg.rope_dim // mma_k_step_for_qkv(cfg) + ): + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + cta_group, + tmem_s_ptr, + desc_qr + k_block * 2, + desc_k + k_block * 2, + idesc_qk, + Boolean(True), + ) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=(row_sum, row_sum_out)) + @cute.jit + def finish_row_sum( + self, + stage_info: StageInfo, + *, + qk_acc_regs, + row_sum, + correction_factor_out, + ): + """Finish row-sum reduction after P is published to SMEM.""" + del stage_info + row_sum = row_sum * correction_factor_out + row_sum_vec = (Float32(0), Float32(0)) + for i in cutlass.range_constexpr(0, 64, 2): + row_sum_vec = add_packed_f32x2( + row_sum_vec, + (qk_acc_regs[i], qk_acc_regs[i + 1]), + ) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + self.row_sum_state = row_sum + return row_sum, row_sum + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, returns=(row_sum_odd, row_sum_out_odd) + ) + @cute.jit + def finish_row_sum_odd( + self, + stage_info: StageInfo, + *, + qk_acc_regs_odd, + row_sum_odd, + correction_factor_out_odd, + ): + """Finish odd-lane row-sum reduction after P is published to SMEM.""" + del stage_info + row_sum = row_sum_odd * correction_factor_out_odd + row_sum_vec = (Float32(0), Float32(0)) + for i in cutlass.range_constexpr(0, 64, 2): + row_sum_vec = add_packed_f32x2( + row_sum_vec, + (qk_acc_regs_odd[i], qk_acc_regs_odd[i + 1]), + ) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + self.row_sum_state = row_sum + return row_sum, row_sum + + @cute.jit + def _load_s_impl( + self, + stage_info: StageInfo, + *, + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ): + """Load S from TMEM and compute the local row max.""" + del row_sum_out, row_max_new, correction_factor_out, no_correction_out + cfg = self.cfg + tidx = cute.arch.thread_idx()[0] + num_compute_threads = cfg.num_compute_warps * cfg.threads_per_warp + local_tidx = tidx % num_compute_threads + stage_idx = stage_info.stage_idx + + work_tile = stage_info.work_tile + K = Int32(work_tile.k_len) + k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) + tile_offset_k = k_index * Int32(cfg.mma_qk_tiler[1]) + needs_row_causal_mask = cutlass.const_expr( + cfg.mask_type == MaskType.CAUSAL.value and self.logical_seq_len_q > 1 + ) + if cutlass.const_expr(needs_row_causal_mask): + batch_idx = Int32(work_tile.tile_idx[2]) + _, logical_seq_len_q = query_batch_bounds( + self.cu_seqlens_q, + batch_idx, + self.logical_seq_len_q, + ) + first_flat_query_row = Int32(work_tile.tile_idx[1]) * Int32( + cfg.mma_qk_tiler[0] + ) + # A row r sees key k iff H * (k - K + SQ) <= r. Apply the + # equivalent boundary test to the first row in this physical tile + # so non-power-of-two H never needs a quotient on the softmax path. + first_mask_flat_row = Int32(self.logical_num_heads_q) * ( + tile_offset_k + + Int32(cfg.mma_qk_tiler[1]) + - K + + logical_seq_len_q + - Int32(1) + ) + group_needs_mask = first_flat_query_row < first_mask_flat_row + else: + group_needs_mask = kv_tile_needs_right_mask( + tile_offset_k, + Int32(cfg.mma_qk_tiler[1]), + K, + ) + + neg_inf = Float32(-Float32.inf) + row_max_tile = row_max + warp_id = local_tidx >> 5 + tmem_warp_row_id = self.tmem_base_addr + warp_id * TCGEN05_32B_REGS_PER_LOAD + stage_offset = stage_idx * 64 + tmem_raw_addr = (tmem_warp_row_id << 16) | stage_offset + t2r_shape = TCGEN05_32B_SHAPE + for load_idx in cutlass.range_constexpr(2): + curr_addr = tmem_raw_addr + load_idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + loaded = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + + if group_needs_mask: + if cutlass.const_expr(needs_row_causal_mask): + # Clamp padded physical tail rows to the request's last real + # row. Their Q/output accesses remain independently + # predicated, while this keeps the masking arithmetic safe. + row_in_tile = self.cta_rank * Int32( + cfg.mma_qk_tiler[0] // cfg.num_mma_ctas + ) + (local_tidx & Int32(EPILOGUE_ROW_MASK)) + flat_query_row = ( + Int32(work_tile.tile_idx[1]) * Int32(cfg.mma_qk_tiler[0]) + + row_in_tile + ) + last_flat_query_row = cute.math.max( + logical_seq_len_q * Int32(self.logical_num_heads_q) - Int32(1), + Int32(0), + ) + flat_query_row = cute.math.min( + flat_query_row, + last_flat_query_row, + ) + tidx_col = ( + local_tidx >> EPILOGUE_COLUMN_GROUP_SHIFT + ) << EPILOGUE_COLUMN_GROUP_SHIFT + for i in cutlass.range_constexpr(64): + token_idx = tile_offset_k + tidx_col + Int32(i) + if cutlass.const_expr(needs_row_causal_mask): + mask_flat_row = Int32(self.logical_num_heads_q) * ( + token_idx - K + logical_seq_len_q + ) + token_is_visible = flat_query_row >= mask_flat_row + else: + token_is_visible = token_idx < K + qk_acc_regs[i] = qk_acc_regs[i] if token_is_visible else neg_inf + + max0 = neg_inf + max1 = neg_inf + max2 = neg_inf + max3 = neg_inf + for i in cutlass.range_constexpr(16): + max0 = fmax_f32(max0, qk_acc_regs[i]) + max1 = fmax_f32(max1, qk_acc_regs[i + 16]) + max2 = fmax_f32(max2, qk_acc_regs[i + 32]) + max3 = fmax_f32(max3, qk_acc_regs[i + 48]) + row_max_tile = fmax_f32( + row_max_tile, + fmax_f32(fmax_f32(max0, max1), fmax_f32(max2, max3)), + ) + cute.arch.fence_view_async_tmem_load() + return ( + qk_acc_regs, + row_max, + row_sum, + row_sum, + row_max_tile, + Float32(1), + Int32(1), + ) + + @consumer_work( + returns=( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ), + ) + @cute.jit + def load_s( + self, + stage_info: StageInfo, + *, + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ): + """Load S for the even softmax group.""" + return self._load_s_impl( + stage_info, + qk_acc_regs=qk_acc_regs, + row_max=row_max, + row_sum=row_sum, + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + ) + + @consumer_work( + returns=( + qk_acc_regs_odd, + row_max_odd, + row_sum_odd, + row_sum_out_odd, + row_max_new_odd, + correction_factor_out_odd, + no_correction_out_odd, + ), + ) + @cute.jit + def load_s_odd( + self, + stage_info: StageInfo, + *, + qk_acc_regs_odd, + row_max_odd, + row_sum_odd, + row_sum_out_odd, + row_max_new_odd, + correction_factor_out_odd, + no_correction_out_odd, + ): + """Load S for the odd softmax group.""" + return self._load_s_impl( + stage_info, + qk_acc_regs=qk_acc_regs_odd, + row_max=row_max_odd, + row_sum=row_sum_odd, + row_sum_out=row_sum_out_odd, + row_max_new=row_max_new_odd, + correction_factor_out=correction_factor_out_odd, + no_correction_out=no_correction_out_odd, + ) + + @cute.jit + def _finish_softmax_impl( + self, + stage_info: StageInfo, + *, + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + softmax_group_id: cutlass.Constexpr[int] = 0, + ): + """Finish row-max exchange, online correction, and P exponentiation.""" + del row_sum_out, correction_factor_out, no_correction_out + cfg = self.cfg + tidx = cute.arch.thread_idx()[0] + num_compute_threads = cfg.num_compute_warps * cfg.threads_per_warp + local_tidx = tidx % num_compute_threads + neg_inf = Float32(-Float32.inf) + row_max_prev = row_max + row_sum_prev = row_sum + row_max_tile = row_max_new + + group_exchange_base = Int32(softmax_group_id * num_compute_threads) + self.smem_exchange[group_exchange_base + local_tidx] = row_max_tile + prims.barrier_cta_sync( + cfg.softmax_sync_bar_id + softmax_group_id, + thread_count=cfg.softmax_sync_threads, + ) + peer_idx = (local_tidx + 64) % num_compute_threads + row_max_tile = fmax_f32( + row_max_tile, self.smem_exchange[group_exchange_base + peer_idx] + ) + # The exchange buffer is reused on the next KV iteration. Keep every + # peer read ahead of any warp's next write to the same slot. + prims.barrier_cta_sync( + cfg.softmax_sync_bar_id + softmax_group_id, + thread_count=cfg.softmax_sync_threads, + ) + + if cutlass.const_expr(cfg.use_fp8_dual_softmax_schedule): + stage_idx = stage_info.stage_idx + + def load_peer_state(): + """Load the peer softmax group's correction state from TMEM.""" + + peer_stage_idx = (stage_idx + Int32(1)) % Int32(cfg.p_cor_stage) + corr_col_offset = cfg.correction_factor_offset + peer_stage_idx * 4 + peer_warp_id = local_tidx >> 5 + peer_tmem_row = ( + self.tmem_base_addr + peer_warp_id * TCGEN05_32B_REGS_PER_LOAD + ) + peer_tmem_addr = (peer_tmem_row << 16) | corr_col_offset + peer_tmem_ptr = prims.make_tmem_ptr(peer_tmem_addr, Float32) + peer_corr = prims.tcgen05_ld(TCGEN05_32B_SHAPE, peer_tmem_ptr, num=2) + return peer_corr[1], peer_corr[0] + + if cutlass.const_expr(softmax_group_id == 1): + if stage_info.loop_offset != Int32(0): + prims.barrier_cta_sync( + cfg.softmax_order_bar_1_id, + thread_count=2 * num_compute_threads, + ) + cute.arch.fence_acq_rel_cta() + row_max_prev, row_sum_prev = load_peer_state() + else: + if stage_info.loop_offset != Int32(0): + prims.barrier_cta_sync( + cfg.softmax_order_bar_0_id, + thread_count=2 * num_compute_threads, + ) + cute.arch.fence_acq_rel_cta() + row_max_prev, row_sum_prev = load_peer_state() + + row_max_new = fmax_f32(row_max_prev, row_max_tile) + row_has_values = row_max_new != neg_inf + safe_row_max_prev = row_max_prev if row_has_values else Float32(0) + safe_row_max_new = row_max_new if row_has_values else Float32(0) + # Exact max equality makes the correction scale exactly one. Keep that + # lane on the identity value and avoid issuing exp2 altogether. + max_changed = safe_row_max_prev != safe_row_max_new + correction_factor = Float32(1) + if max_changed: + correction_factor = cute.math.exp2( + (safe_row_max_prev - safe_row_max_new) * self.softmax_scale_log2, + fastmath=True, + ) + no_correction = Int32(not max_changed) + + fma_b = self.softmax_scale_log2 + fma_c = Float32(0) - safe_row_max_new * self.softmax_scale_log2 + if cutlass.const_expr(cfg.is_fp8_qkv()): + # Match the 448-scaled E4M3 P convention used by the reference + # output and the 1CTA implementation. + fma_c = fma_c + fp8_log2_quant_scale() + for i in cutlass.range_constexpr(0, 64, 2): + fma_result = fma_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (fma_b, fma_b), + (fma_c, fma_c), + ) + qk_acc_regs[i] = cute.math.exp2(fma_result[0], fastmath=True) + qk_acc_regs[i + 1] = cute.math.exp2(fma_result[1], fastmath=True) + + self.row_max_state = row_max_new + return ( + qk_acc_regs, + row_max_new, + row_sum_prev, + row_sum_prev, + row_max_new, + correction_factor, + no_correction, + ) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ), + ) + @cute.jit + def finish_softmax( + self, + stage_info: StageInfo, + *, + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ): + """Finish softmax for the even group after S release.""" + return self._finish_softmax_impl( + stage_info, + qk_acc_regs=qk_acc_regs, + row_max=row_max, + row_sum=row_sum, + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + softmax_group_id=0, + ) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + qk_acc_regs_odd, + row_max_odd, + row_sum_odd, + row_sum_out_odd, + row_max_new_odd, + correction_factor_out_odd, + no_correction_out_odd, + ), + ) + @cute.jit + def finish_softmax_odd( + self, + stage_info: StageInfo, + *, + qk_acc_regs_odd, + row_max_odd, + row_sum_odd, + row_sum_out_odd, + row_max_new_odd, + correction_factor_out_odd, + no_correction_out_odd, + ): + """Finish softmax for the odd group after S release.""" + return self._finish_softmax_impl( + stage_info, + qk_acc_regs=qk_acc_regs_odd, + row_max=row_max_odd, + row_sum=row_sum_odd, + row_sum_out=row_sum_out_odd, + row_max_new=row_max_new_odd, + correction_factor_out=correction_factor_out_odd, + no_correction_out=no_correction_out_odd, + softmax_group_id=1, + ) + + +# ===================================================================== +# SmemPResource — P in SMEM, UmmaConsumerAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemPResource(HighThroughputMlaResource): + """SMEM P buffer. Producer: SoftmaxTask. Consumer: MmaTask (PV MMA). + + Pipeline: UmmaConsumerAsync, 2 stages. + """ + + smem_p: Any = None + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cta_rank: Any = field(init=False, default=None) + desc_p_base: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("desc_p_base", Int64, Int64(0), "SMEM descriptor for staged P."), + ) + + @cute.jit + def _store_p_impl(self, stage_info: StageInfo, *, qk_acc_regs) -> None: + """Store softmax P tile to SMEM for PV MMA. + + Converts P values to FP16 and stores to SMEM with correct swizzle layout. + """ + cfg = self.cfg + stage_idx = stage_info.stage_idx + num_mma_ctas = cfg.num_mma_ctas + + tidx = cute.arch.thread_idx()[0] + num_compute_threads = cfg.num_compute_warps * cfg.threads_per_warp + local_tidx = tidx % num_compute_threads + lane_idx = local_tidx & 31 + warp_idx = local_tidx >> 5 + + # Convert F32 -> QKV dtype into a vectorized local buffer. + qkv_element_dtype = qkv_dtype(cfg) + s_regs = cutlass.Array(qkv_element_dtype, 64, space=cutlass.AddressSpace.rmem) + s_regs.store(qk_acc_regs.load(0, 64).to(qkv_element_dtype), 0) + + # Compute SMEM P stage offset + sp_stage_stride_elems = cutlass.const_expr( + cfg.mma_pv_tiler[0] + // num_mma_ctas + * cfg.mma_pv_tiler[2] + * cfg.iterations_pv_k + ) + smem_p_base_bytes = ( + self.smem_p.data_ptr().toint(cutlass.Int32) + + stage_idx * sp_stage_stride_elems * cfg.qkv_dtype_bytes + ) + + if cutlass.const_expr(cfg.is_fp8_qkv()): + # FP8 P layout: + # S<2,4,3> o ((64,32),1,2,(2,2)):((64,1),0,32,(4096,8192)). + # Each universal SMEM copy is 128 bits, i.e. 16 E4M3 elements. + # Lane pairs write opposite halves of the 128B swizzled row. The + # signed strides below walk the four 16B blocks inside that row. + m = ((lane_idx >> 1) & 3) * 128 + base = cutlass.Int32(0) + base = base + (lane_idx & 1) * 64 + base = base + (lane_idx >> 3) * 512 + base = base + (warp_idx & 1) * 2048 + base = base + (warp_idx >> 1) * 4096 + swizzle_xor = m ^ ((m & 384) >> 3) + off_base = base + swizzle_xor + + stride_a = cutlass.Int32(16) + if (m & 128) != 0: + stride_a = cutlass.Int32(-16) + stride_b = cutlass.Int32(32) + if (m & 256) != 0: + stride_b = cutlass.Int32(-32) + + dst_blk_offs = ( + cutlass.Int32(0), + stride_a, + stride_b, + stride_a + stride_b, + ) + for blk in cutlass.range_constexpr(4): + src_blk_base = blk * 16 + dst_blk_addr = smem_p_base_bytes + off_base + dst_blk_offs[blk] + vec = s_regs.load(src_blk_base, 16) + smem_ptr = cutlass.inttoptr(dst_blk_addr, 3, qkv_element_dtype) + smem_ptr.store(vec, alignment=16) + else: + # BF16 P layout: + # S<2,4,3> o ((64,16),1,2,(4,2)):((32,1),0,16,(2048,8192)). + # Each universal SMEM copy is 128 bits, i.e. 8 BF16 elements. + # The two K slices live 2048 elements apart in the staged P tile. + # Per-block strides mirror the FP8 swizzle at half the byte width. + m = ((lane_idx >> 1) & 3) * 64 + base = cutlass.Int32(0) + base = base + (lane_idx & 1) * 32 + base = base + (lane_idx >> 3) * 256 + base = base + (warp_idx & 1) * 1024 + base = base + (warp_idx >> 1) * 4096 + swizzle_xor = m ^ ((m & 192) >> 3) + off_base = base + swizzle_xor + + stride_a = cutlass.Int32(8) + if (m & 64) != 0: + stride_a = cutlass.Int32(-8) + stride_b = cutlass.Int32(16) + if (m & 128) != 0: + stride_b = cutlass.Int32(-16) + + dst_blk_offs = ( + cutlass.Int32(0), + stride_a, + stride_b, + stride_a + stride_b, + ) + for k in cutlass.range_constexpr(2): + k_base = off_base + k * 2048 + src_k_base = k * 32 + for blk in cutlass.range_constexpr(4): + src_blk_base = src_k_base + blk * 8 + dst_blk_addr = ( + smem_p_base_bytes + + (k_base + dst_blk_offs[blk]) * cfg.qkv_dtype_bytes + ) + vec = s_regs.load(src_blk_base, 8) + smem_ptr = cutlass.inttoptr(dst_blk_addr, 3, qkv_element_dtype) + smem_ptr.store(vec, alignment=16) + + # Fence between SMEM store and MMA read + prims.fence_proxy( + kind=prims.Proxy.ASYNC_SHARED, + space=prims.SharedSpace.shared_cta, + ) + + @producer_work + @cute.jit + def store_p(self, stage_info: StageInfo, *, qk_acc_regs) -> None: + """Store P from the even softmax group.""" + self._store_p_impl(stage_info, qk_acc_regs=qk_acc_regs) + + @producer_work + @cute.jit + def store_p_odd(self, stage_info: StageInfo, *, qk_acc_regs_odd) -> None: + """Store P from the odd softmax group.""" + self._store_p_impl(stage_info, qk_acc_regs=qk_acc_regs_odd) + + @consumer_work(returns=desc_p_base) + @cute.jit + def p_desc(self, stage_info: StageInfo): + """MMA warp builds P SMEM descriptor for PV MMA.""" + cfg = self.cfg + sp_stage_stride_elems = cutlass.const_expr( + cfg.mma_pv_tiler[0] + // cfg.num_mma_ctas + * cfg.mma_pv_tiler[2] + * cfg.iterations_pv_k + ) + sp_ptr = self.smem_p.data_ptr(stage_info.stage_idx * sp_stage_stride_elems) + desc_p = Int64( + prims.Tcgen05SmemDesc.build( + start_address=sp_ptr.toint(Int32), + leading_byte_offset=p_desc_leading_byte_offset(cfg), + stride_byte_offset=p_desc_stride_byte_offset(cfg), + layout=p_desc_layout(cfg), + ) + ) + return desc_p + + +# ===================================================================== +# TmemCorrResource — Correction factors via TMEM, Async pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemCorrResource(HighThroughputMlaResource): + """Correction factors in TMEM. Producer: SoftmaxTask. Consumer: CorrectionTask. + + Pipeline: Async, 2 stages. + Carries (row_sum, row_max, correction_factor, no_correction) per thread. + """ + + tmem_base_addr: Any = None + smem_exchange: Any = None + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cta_rank: Any = field(init=False, default=None) + final_row_stats: Any = field(init=False, default=None) + row_sum: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + row_max: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + epilogue_row_sum: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + epilogue_row_max: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + correction_factor: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + no_correction: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("row_sum", Float32, Float32(0), "Final running row sum."), + ("row_max", Float32, Float32(0), "Final running row max."), + ("epilogue_row_sum", Float32, Float32(0), "Final exchanged row sum."), + ("epilogue_row_max", Float32, Float32(0), "Final row max for LSE."), + ( + "correction_factor", + Float32, + Float32(0), + "Correction factor for the previous O tile.", + ), + ("no_correction", Int32, Int32(0), "Whether O correction may be skipped."), + ) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(row_sum, row_max, correction_factor, no_correction), + ) + @cute.jit + def init_load_state(self, stage_info: StageInfo): + """Create row-stat variables consumed by correction and epilogue code.""" + del stage_info + self.cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + self.final_row_stats = cutlass.Array( + Float32, + 2, + space=cutlass.AddressSpace.rmem, + ) + return Float32(0), Float32(0), Float32(0), Int32(0) + + @cute.jit + def _store_corr_impl( + self, + stage_info: StageInfo, + *, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + softmax_group_id: cutlass.Constexpr[int] = 0, + arrive_peer: cutlass.Constexpr[bool] = True, + ) -> None: + """Store correction factors [row_sum, row_max, correction, no_correction] to TMEM.""" + cfg = self.cfg + stage_idx = stage_info.stage_idx + + tidx = cute.arch.thread_idx()[0] + num_compute_threads = cfg.num_compute_warps * cfg.threads_per_warp + local_tidx = tidx % num_compute_threads + warp_id = local_tidx >> 5 + + col_offset = cfg.correction_factor_offset + stage_idx * 4 + tmem_warp_row_id = self.tmem_base_addr + warp_id * TCGEN05_32B_REGS_PER_LOAD + # tcgen05 addresses pack the TMEM row into the high 16 bits. Each warp + # owns one correction row and four adjacent columns for sum/max/scale. + tmem_raw_addr = (tmem_warp_row_id << 16) | col_offset + tmem_ptr_arr = prims.make_tmem_ptr(tmem_raw_addr, Float32) + + correction_regs = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + correction_regs[0] = row_sum_out + correction_regs[1] = row_max_new + correction_regs[2] = correction_factor_out + correction_regs[3] = prims.mov_b32(no_correction_out, target_type=Float32) + + prims.tcgen05_st( + TCGEN05_32B_SHAPE, + tmem_ptr_arr, + correction_regs[0:4], + ) + cute.arch.fence_view_async_tmem_store() + + if cutlass.const_expr(cfg.use_fp8_dual_softmax_schedule and arrive_peer): + # Dual-softmax pipes are ordered one loop tile apart: after one pipe + # publishes correction state, it releases the peer pipe for the next + # loop tile if that peer still has work. + has_peer_next = stage_info.loop_offset + Int32(1) < stage_info.loop_end + if has_peer_next: + prims.barrier_cta_arrive( + ( + cfg.softmax_order_bar_0_id + if cutlass.const_expr(softmax_group_id == 1) + else cfg.softmax_order_bar_1_id + ), + 2 * num_compute_threads, + ) + + @producer_work + @cute.jit + def store_corr( + self, + stage_info: StageInfo, + *, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ) -> None: + """Store correction metadata for the even softmax group.""" + self._store_corr_impl( + stage_info, + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + softmax_group_id=0, + arrive_peer=True, + ) + + @producer_work + @cute.jit + def store_corr_odd( + self, + stage_info: StageInfo, + *, + row_sum_out_odd, + row_max_new_odd, + correction_factor_out_odd, + no_correction_out_odd, + ) -> None: + """Store correction metadata for the odd softmax group.""" + self._store_corr_impl( + stage_info, + row_sum_out=row_sum_out_odd, + row_max_new=row_max_new_odd, + correction_factor_out=correction_factor_out_odd, + no_correction_out=no_correction_out_odd, + softmax_group_id=1, + arrive_peer=True, + ) + + @consumer_work(returns=(row_sum, row_max, correction_factor, no_correction)) + @cute.jit + def load_corr(self, stage_info: StageInfo): + """Load correction factors from TMEM.""" + cfg = self.cfg + stage_idx = stage_info.stage_idx + + tidx = cute.arch.thread_idx()[0] + # Use local tidx within 4-warp correction group (matching bare-metal) + local_tidx = tidx % (cfg.num_compute_warps * cfg.threads_per_warp) + warp_id = local_tidx >> 5 + + col_offset = cfg.correction_factor_offset + stage_idx * 4 + tmem_warp_row_id = self.tmem_base_addr + warp_id * TCGEN05_32B_REGS_PER_LOAD + # Load from the same packed row/column address used by store_corr so the + # correction consumer sees the row statistics for its current stage. + tmem_raw_addr = (tmem_warp_row_id << 16) | col_offset + tmem_ptr_arr = prims.make_tmem_ptr(tmem_raw_addr, Float32) + + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + loaded = prims.tcgen05_ld( + TCGEN05_32B_SHAPE, + tmem_ptr_arr, + num=4, + ) + + self.final_row_stats[0] = loaded[0] + self.final_row_stats[1] = loaded[1] + return loaded[0], loaded[1], loaded[2], loaded[3].bitcast(Int32) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, returns=(epilogue_row_sum, epilogue_row_max) + ) + @cute.jit + def prepare_epilogue_slice_store(self, stage_info: StageInfo): + """Exchange final row statistics once before per-slice O stores.""" + del stage_info + cfg = self.cfg + tidx = cute.arch.thread_idx()[0] + row_sum = self.final_row_stats[0] + row_max = self.final_row_stats[1] + + num_compute_threads = cfg.num_compute_warps * cfg.threads_per_warp + local_tidx = tidx % num_compute_threads + smem_ex_ptr = cutlass.inttoptr( + self.smem_exchange + local_tidx * 4, + 3, + Float32, + ) + smem_ex_ptr.store(row_sum) + prims.barrier_cta_sync( + cfg.epilogue_sync_bar_id, thread_count=cfg.epilogue_sync_threads + ) + # The two CTAs in the cluster own complementary halves of the 2CTA row. + # Exchanging row sums through SMEM gives both epilogue slices the same + # denominator while preserving each CTA's local row max for LSE. + peer_idx = (local_tidx + 64) % num_compute_threads + peer_ptr = cutlass.inttoptr( + self.smem_exchange + peer_idx * 4, + 3, + Float32, + ) + row_sum = row_sum + peer_ptr.load() + prims.barrier_cta_sync( + cfg.epilogue_sync_bar_id, thread_count=cfg.epilogue_sync_threads + ) + return row_sum, row_max + + +# ===================================================================== +# TmemOResource — O accumulator in TMEM, UmmaProducerAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemOResource(HighThroughputMlaResource): + """TMEM O accumulator. Producer: Mma (PV MMA). Consumer: Correction. + + Pipeline: UmmaProducerAsync, 1 stage. + """ + + tmem_base_addr: Any = None + tmem_corr_ref: Any = None # Reference to TmemCorrResource for correction data + smem_p: Any = None # SMEM P for PV MMA + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cta_rank: Any = field(init=False, default=None) + is_leader: Any = field(init=False, default=None) + + @producer_work + @cute.jit + def pv_mma( + self, + stage_info: StageInfo, + *, + desc_p_base, + desc_v_base, + v_subtile_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """Issue one PV MMA sub-tile. + Only leader CTA issues MMA (2CTA UMMA principle). + + Descriptor computation is hoisted OUTSIDE the leader-CTA gate so + ptxas keeps values in uniform registers. + """ + cfg = self.cfg + + idesc_pv = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + n_dim=cfg.mma_pv_tiler[1], + m_dim=cfg.mma_pv_tiler[0], + b_major=1, + ) + mma_kind = mma_kind_for_qkv(cfg) + cta_group = prims.CTAGroup.CTA_2 + is_leader_cta = ( + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) == 0 + ) + + sp_copy_elems = cutlass.const_expr( + cfg.mma_pv_tiler[0] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] + ) + desc_p_delta = cutlass.const_expr(sp_copy_elems // 8) + if cutlass.const_expr(cfg.is_fp8_qkv()): + desc_p_delta = cutlass.const_expr(256) + sv_copy_elems = cutlass.const_expr( + cfg.mma_pv_tiler[1] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] + ) + desc_v_delta = cutlass.const_expr(sv_copy_elems * cfg.qkv_dtype_bytes // 16) + + pv_k_block_count = cutlass.const_expr( + cfg.mma_pv_tiler[2] // mma_k_step_for_qkv(cfg) + ) + + # Match the producer's physical V-stage decomposition. Each stage + # carries two adjacent K32 slices for one D256 output panel. + pv_j = v_subtile_idx // cfg.kv_subtiles_per_stage + token_partition = v_subtile_idx % cfg.kv_subtiles_per_stage + tmem_o_base_addr = self.tmem_base_addr + cfg.tmem_o_offset + tmem_o_addr = tmem_o_base_addr + pv_j * 128 + + for stage_subtile_idx in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + pv_i = token_partition * cfg.kv_subtiles_per_stage + stage_subtile_idx + desc_p = desc_p_base + pv_i * desc_p_delta + desc_v = desc_v_base + stage_subtile_idx * desc_v_delta + + # Clear O for each output-N panel on its first PV K slice, then + # accumulate the remaining slices and subsequent sequence tiles. + scale_d_pv = Boolean(True) + if cutlass.const_expr(pv_i == 0): + if cutlass.const_expr(is_tail): + scale_d_pv = Boolean( + stage_info.loop_end > Int32(stage_info.loop_start) + ) + else: + scale_d_pv = Boolean( + stage_info.loop_offset != Int32(stage_info.loop_start) + ) + if is_leader_cta: + for k_block in cutlass.range_constexpr(pv_k_block_count): + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + cta_group, + prims.make_tmem_ptr(tmem_o_addr, Float32), + desc_p + k_block * 2, + desc_v + k_block * (256 if cfg.is_fp8_qkv() else 128), + idesc_pv, + Boolean(scale_d_pv), + ) + scale_d_pv = Boolean(True) + + @producer_work + @cute.jit + def pv_mma_n_major( + self, + stage_info: StageInfo, + *, + desc_p_base, + desc_v_base, + pv_n_idx: cutlass.Constexpr[int], + pv_k_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """Issue PV MMA in output-N-major order for per-slice O tokens.""" + cfg = self.cfg + + pv_j = pv_n_idx + pv_i = pv_k_idx + tmem_o_base_addr = self.tmem_base_addr + cfg.tmem_o_offset + tmem_o_addr = tmem_o_base_addr + pv_j * 128 + + idesc_pv = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + n_dim=cfg.mma_pv_tiler[1], + m_dim=cfg.mma_pv_tiler[0], + b_major=1, + ) + mma_kind = mma_kind_for_qkv(cfg) + cta_group = prims.CTAGroup.CTA_2 + is_leader_cta = ( + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) == 0 + ) + + sp_copy_elems = cutlass.const_expr( + cfg.mma_pv_tiler[0] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] + ) + desc_p_delta = cutlass.const_expr(sp_copy_elems // 8) + if cutlass.const_expr(cfg.is_fp8_qkv()): + desc_p_delta = cutlass.const_expr(256) + desc_p = desc_p_base + pv_i * desc_p_delta + + pv_k_block_count = cutlass.const_expr( + cfg.mma_pv_tiler[2] // mma_k_step_for_qkv(cfg) + ) + scale_d_pv = Boolean(True) + if cutlass.const_expr(pv_i == 0): + if cutlass.const_expr(is_tail): + scale_d_pv = Boolean(stage_info.loop_end > Int32(stage_info.loop_start)) + else: + scale_d_pv = Boolean( + stage_info.loop_offset != Int32(stage_info.loop_start) + ) + + if is_leader_cta: + for k_block in cutlass.range_constexpr(pv_k_block_count): + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + cta_group, + prims.make_tmem_ptr(tmem_o_addr, Float32), + desc_p + k_block * 2, + desc_v_base + k_block * (256 if cfg.is_fp8_qkv() else 128), + idesc_pv, + Boolean(scale_d_pv), + ) + scale_d_pv = Boolean(True) + + @consumer_work + @cute.jit + def rescale_o( + self, stage_info: StageInfo, *, correction_factor, no_correction + ) -> None: + """Rescale O in-place in TMEM by correction_factor.""" + cfg = self.cfg + tidx = cute.arch.thread_idx()[0] + + # Use local tidx within 4-warp correction group (matching bare-metal) + local_tidx = tidx % (cfg.num_compute_warps * cfg.threads_per_warp) + tmem_warp_row_id = ( + self.tmem_base_addr + + (local_tidx >> WARP_LANE_SHIFT) * TCGEN05_32B_REGS_PER_LOAD + ) + tmem_raw_addr = (tmem_warp_row_id << 16) | cfg.tmem_o_offset + + t2r_shape = TCGEN05_32B_SHAPE + num_tmem_ops = 4 # 4 loads x 32 = 128 elements per iter_n + + skip_correction = prims.vote_sync( + cute.arch.FULL_MASK, + no_correction == 1, + prims.VoteSync.ALL, + ) + + if not skip_correction: + for iter_n in cutlass.range_constexpr(cfg.iterations_pv_n): + tmem_addr_offset = iter_n * (cfg.mma_pv_tiler[1] // cfg.warps_in_n) + for idx in cutlass.range_constexpr(num_tmem_ops): + curr_addr = ( + tmem_raw_addr + + tmem_addr_offset + + idx * TCGEN05_32B_REGS_PER_LOAD + ) + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + chunk = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + scaled = chunk * cutlass.full_like(chunk, correction_factor) + prims.tcgen05_st(t2r_shape, tmem_ptr, scaled) + + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + + @consumer_work + @cute.jit + def rescale_o_slice( + self, + stage_info: StageInfo, + *, + correction_factor, + no_correction, + iter_n: cutlass.Constexpr[int], + ) -> None: + """Rescale one PV output-N slice after its O token is ready.""" + cfg = self.cfg + tidx = cute.arch.thread_idx()[0] + local_tidx = tidx % (cfg.num_compute_warps * cfg.threads_per_warp) + tmem_warp_row_id = ( + self.tmem_base_addr + + (local_tidx >> WARP_LANE_SHIFT) * TCGEN05_32B_REGS_PER_LOAD + ) + tmem_raw_addr = (tmem_warp_row_id << 16) | ( + cfg.tmem_o_offset + iter_n * (cfg.mma_pv_tiler[1] // cfg.warps_in_n) + ) + + t2r_shape = TCGEN05_32B_SHAPE + num_tmem_ops = 4 + skip_correction = prims.vote_sync( + cute.arch.FULL_MASK, + no_correction == 1, + prims.VoteSync.ALL, + ) + if not skip_correction: + for idx in cutlass.range_constexpr(num_tmem_ops): + curr_addr = tmem_raw_addr + idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + chunk = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + scaled = chunk * cutlass.full_like(chunk, correction_factor) + prims.tcgen05_st(t2r_shape, tmem_ptr, scaled) + + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + + +# ===================================================================== +# GmemOResource — Output to GMEM (no pipeline) +# ===================================================================== + + +@dataclass(kw_only=True) +class GmemOResource(HighThroughputMlaResource): + """GMEM output. Producer: Correction (epilogue store). No pipeline.""" + + output: Any = None + partial_output: Any = None + lse: Any = None + partial_lse: Any = None + tmem_o_ref: Any = None # Reference to TmemOResource for tmem_base_addr + tmem_corr_ref: Any = None # Reference to TmemCorrResource for correction data + output_scale: Any = None + softmax_scale_log2: Any = None + smem_exchange: Any = None # SMEM for row_sum exchange (as Int32 base addr) + split_kv: Any = None + cu_seqlens_q: Any = None + logical_num_heads_q: cutlass.Constexpr[int] = 128 + logical_seq_len_q: cutlass.Constexpr[int] = 1 + cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + + @cute.jit + def _query_row_state(self, row_in_tile, query_tile_idx, batch_idx): + """Map one physical flat-tile row to public storage.""" + return flat_query_row_state( + row_in_tile, + query_tile_idx, + self.cfg.mma_qk_tiler[0], + self.logical_num_heads_q, + self.logical_seq_len_q, + self.cu_seqlens_q, + batch_idx, + ) + + @producer_work + @cute.jit + def epilogue_store(self, stage_info: StageInfo) -> None: + """Full epilogue: load O from TMEM, normalize by row_sum, store to GMEM. + + Steps: + 1. Exchange row_sum across warp pairs via SMEM + 2. Load O from TMEM + 3. Normalize by output_scale / row_sum + 4. Store to GMEM + 5. Compute and store LSE + """ + cfg = self.cfg + tidx = cute.arch.thread_idx()[0] + # Use work_tile's blk_coord (updated by persistent loop) instead of + # self.blk_coord (captured at construction time) so that each work + # tile writes to the correct output location. + blk_coord = stage_info.work_tile.tile_idx + + row_sum = self.tmem_corr_ref.final_row_stats[0] + row_max = self.tmem_corr_ref.final_row_stats[1] + + # Exchange row_sum between warp pairs (0,2) and (1,3) via SMEM + # Use local thread index within the 4-warp correction group, + # matching bare-metal: tidx % (num_compute_warps * threads_per_warp) + num_compute_threads = cfg.num_compute_warps * cfg.threads_per_warp + local_tidx = tidx % num_compute_threads + smem_ex_ptr = cutlass.inttoptr( + self.smem_exchange + local_tidx * 4, + 3, + Float32, + ) + smem_ex_ptr.store(row_sum) + prims.barrier_cta_sync( + cfg.epilogue_sync_bar_id, thread_count=cfg.epilogue_sync_threads + ) + peer_idx = (local_tidx + 64) % num_compute_threads + peer_ptr = cutlass.inttoptr( + self.smem_exchange + peer_idx * 4, + 3, + Float32, + ) + row_sum = row_sum + peer_ptr.load() + prims.barrier_cta_sync( + cfg.epilogue_sync_bar_id, thread_count=cfg.epilogue_sync_threads + ) + + # TMEM address for O — use local tidx within 4-warp correction group + tmem_base_addr = self.tmem_o_ref.tmem_base_addr + tmem_warp_row_id = ( + tmem_base_addr + (local_tidx >> WARP_LANE_SHIFT) * TCGEN05_32B_REGS_PER_LOAD + ) + tmem_raw_addr = (tmem_warp_row_id << 16) | cfg.tmem_o_offset + + t2r_shape = TCGEN05_32B_SHAPE + num_tmem_loads = 4 # 4 x 32 = 128 elements per iter_n + + # Per-thread O indexing — use local tidx + tile_h = cfg.mma_pv_tiler[0] // cfg.num_mma_ctas # 64 + tile_d = cfg.mma_pv_tiler[1] # 256 + tidx_g = local_tidx & EPILOGUE_THREAD_TILE_MASK + g_i = tidx_g & EPILOGUE_ROW_MASK + g_j = (tidx_g >> EPILOGUE_COLUMN_GROUP_SHIFT) * EPILOGUE_THREAD_TILE_THREADS + + # Public O/LSE remain in logical coordinates. Split-KV partials retain + # physical flat-tile coordinates until final reduction. + logical_num_heads_q = Int32(self.logical_num_heads_q) + physical_tile_rows = Int32(cfg.mma_qk_tiler[0]) + D = cfg.latent_dim + head_tile_idx = blk_coord[0] + seq_q_idx = blk_coord[1] + batch_idx = blk_coord[2] + split_kv_idx = blk_coord[3] + row_in_tile = head_tile_idx * tile_h + g_i + ( + storage_flat_query_row, + logical_head_idx, + logical_q_idx, + _, + query_is_valid, + ) = self._query_row_state(row_in_tile, seq_q_idx, batch_idx) + + # Fully masked split rows can occur when physical tail rows or earlier + # causal query rows have no visible K values in this split. Store zero + # O and -inf LSE so split-KV reduction gives those rows zero weight. + row_has_values = row_sum > Float32(0) + safe_row_sum = row_sum if row_has_values else Float32(1) + norm_scale = self.output_scale * cute.math.rcp(safe_row_sum, approx=True) + + for iter_n in cutlass.range_constexpr(cfg.iterations_pv_n): + # Load O from TMEM + qk_acc_regs = cutlass.Array(Float32, 128, space=cutlass.AddressSpace.rmem) + tmem_raw_addr_n = tmem_raw_addr + ( + iter_n * (cfg.mma_pv_tiler[1] // cfg.warps_in_n) + ) + for load_idx in cutlass.range_constexpr(num_tmem_loads): + curr_addr = tmem_raw_addr_n + load_idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + loaded = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + + # Normalize: O = O * output_scale / row_sum + for i in cutlass.range_constexpr(0, 128, 2): + scaled = mul_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (norm_scale, norm_scale), + ) + qk_acc_regs[i] = scaled[0] + qk_acc_regs[i + 1] = scaled[1] + + # Store O to GMEM + if row_in_tile < physical_tile_rows and query_is_valid: + if cutlass.const_expr(self.partial_output is not None): + # Split-KV partial O uses BF16 workspace storage. LSE and + # the eventual cross-split accumulation remain FP32. + S_q = ( + cutlass.Int32(self.partial_output.shape[3]) + if self.partial_output is not None + else Int32(1) + ) + o_base_ptr = ( + self.partial_output.iterator.raw_ptr() + + Int64(row_in_tile) * Int64(self.split_kv) * Int64(D) + + Int64(split_kv_idx) * Int64(D) + + Int64(seq_q_idx) + * Int64(self.split_kv) + * Int64(physical_tile_rows) + * Int64(D) + + Int64(batch_idx) + * Int64(physical_tile_rows) + * Int64(self.split_kv) + * Int64(S_q) + * Int64(D) + ) + output_base = o_base_ptr + iter_n * tile_d + g_j + for load_idx in cutlass.range_constexpr(num_tmem_loads): + for j in cutlass.range_constexpr(4): + offset = ( + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * BF16_OUTPUT_VECTOR_ELEMENTS + ) + vec_f32 = qk_acc_regs.load( + offset, BF16_OUTPUT_VECTOR_ELEMENTS + ) + vec_partial = vec_f32.to(cutlass.BFloat16) + ( + output_base + + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * BF16_OUTPUT_VECTOR_ELEMENTS + ).nvvm_store_ext( + vec_partial, + evict="noallocate", + ) + else: + # 16-bit output (split_kv == 1, direct output) + if cutlass.const_expr(self.cu_seqlens_q is not None): + o_base_ptr = self.output.iterator.raw_ptr() + Int64( + storage_flat_query_row + ) * Int64(D) + else: + S_q = ( + cutlass.Int32(self.output.shape[2]) + if self.output is not None + else Int32(1) + ) + o_base_ptr = ( + self.output.iterator.raw_ptr() + + Int64(logical_head_idx) * Int64(D) + + Int64(logical_q_idx) + * Int64(logical_num_heads_q) + * Int64(D) + + Int64(batch_idx) + * Int64(logical_num_heads_q) + * Int64(D) + * Int64(S_q) + ) + output_base = o_base_ptr + iter_n * tile_d + g_j + for load_idx in cutlass.range_constexpr(num_tmem_loads): + for j in cutlass.range_constexpr(2): + offset = ( + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * FP8_OUTPUT_VECTOR_ELEMENTS + ) + vec_f32 = qk_acc_regs.load( + offset, FP8_OUTPUT_VECTOR_ELEMENTS + ) + if cutlass.const_expr(cfg.use_fp8_output == 1): + packed_o = cutlass.Array( + Int32, + PACKED_FP8_OUTPUT_REGS, + space=cutlass.AddressSpace.rmem, + ) + for pack_idx in cutlass.range_constexpr( + PACKED_FP8_OUTPUT_REGS + ): + pack_offset = pack_idx * PACKED_FP8_OUTPUT_REGS + packed_o[pack_idx] = pack_float4_to_fp8_e4m3( + vec_f32[pack_offset], + vec_f32[pack_offset + 1], + vec_f32[pack_offset + 2], + vec_f32[pack_offset + 3], + ) + raw_ptr = cutlass.inttoptr( + ( + output_base + + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * FP8_OUTPUT_VECTOR_ELEMENTS + ).toint(Int64), + mem_space=1, + dtype=Int32, + ) + raw_ptr.store( + packed_o.load(0, PACKED_FP8_OUTPUT_REGS), + alignment=16, + ) + else: + vec_o = vec_f32.to(output_dtype(self.cfg)) + ( + output_base + + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * FP8_OUTPUT_VECTOR_ELEMENTS + ).nvvm_store_ext( + vec_o, + evict="noallocate", + ) + + # Compute and store LSE in the same row-sum domain used by P. + lse_row_sum = row_sum + if cutlass.const_expr(cfg.is_fp8_qkv()): + lse_row_sum = lse_row_sum * fp8_quant_scale_rcp() + lse = ( + cute.math.log2(lse_row_sum, fastmath=True) + + self.softmax_scale_log2 * row_max + if row_has_values + else Float32(-Float32.inf) + ) + + # Use local_tidx (0..127 within correction warpgroup) for LSE + # indexing, not global tidx (which is 128..255 for correction warps). + lse_tidx = local_tidx + if lse_tidx < tile_h: + lse_row_in_tile = head_tile_idx * tile_h + lse_tidx + ( + storage_flat_lse_row, + logical_lse_head_idx, + logical_lse_q_idx, + _, + lse_query_is_valid, + ) = self._query_row_state(lse_row_in_tile, seq_q_idx, batch_idx) + if lse_row_in_tile < physical_tile_rows and lse_query_is_valid: + if cutlass.const_expr(self.partial_lse is not None): + S_q = ( + cutlass.Int32(self.partial_lse.shape[2]) + if self.partial_lse is not None + else Int32(1) + ) + lse_base_ptr = ( + self.partial_lse.iterator.raw_ptr() + + Int64(lse_row_in_tile) * Int64(self.split_kv) + + Int64(split_kv_idx) + + Int64(seq_q_idx) + * Int64(physical_tile_rows) + * Int64(self.split_kv) + + Int64(batch_idx) + * Int64(physical_tile_rows) + * Int64(self.split_kv) + * Int64(S_q) + ) + lse_base_ptr.store(lse) + elif cutlass.const_expr(self.lse is not None): + if cutlass.const_expr(self.cu_seqlens_q is not None): + lse_base_ptr = ( + self.lse.iterator.raw_ptr() + storage_flat_lse_row + ) + else: + S_q = ( + cutlass.Int32(self.lse.shape[1]) + if self.lse is not None + else Int32(1) + ) + lse_base_ptr = ( + self.lse.iterator.raw_ptr() + + Int64(logical_lse_head_idx) + + Int64(logical_lse_q_idx) * Int64(logical_num_heads_q) + + Int64(batch_idx) * Int64(logical_num_heads_q) * Int64(S_q) + ) + lse_base_ptr.store(lse) + + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + + @producer_work + @cute.jit + def epilogue_store_slice( + self, + stage_info: StageInfo, + *, + row_sum, + row_max, + iter_n: cutlass.Constexpr[int], + ) -> None: + """Store one output-N slice after its O pipeline token is ready.""" + cfg = self.cfg + tidx = cute.arch.thread_idx()[0] + blk_coord = stage_info.work_tile.tile_idx + + num_compute_threads = cfg.num_compute_warps * cfg.threads_per_warp + local_tidx = tidx % num_compute_threads + + tmem_base_addr = self.tmem_o_ref.tmem_base_addr + tmem_warp_row_id = ( + tmem_base_addr + (local_tidx >> WARP_LANE_SHIFT) * TCGEN05_32B_REGS_PER_LOAD + ) + tmem_raw_addr = (tmem_warp_row_id << 16) | ( + cfg.tmem_o_offset + iter_n * (cfg.mma_pv_tiler[1] // cfg.warps_in_n) + ) + + t2r_shape = TCGEN05_32B_SHAPE + num_tmem_loads = 4 + qk_acc_regs = cutlass.Array(Float32, 128, space=cutlass.AddressSpace.rmem) + for load_idx in cutlass.range_constexpr(num_tmem_loads): + curr_addr = tmem_raw_addr + load_idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + loaded = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + + row_has_values = row_sum > Float32(0) + safe_row_sum = row_sum if row_has_values else Float32(1) + norm_scale = self.output_scale * cute.math.rcp(safe_row_sum, approx=True) + for i in cutlass.range_constexpr(0, 128, 2): + scaled = mul_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (norm_scale, norm_scale), + ) + qk_acc_regs[i] = scaled[0] + qk_acc_regs[i + 1] = scaled[1] + + tile_h = cfg.mma_pv_tiler[0] // cfg.num_mma_ctas + tile_d = cfg.mma_pv_tiler[1] + tidx_g = local_tidx & EPILOGUE_THREAD_TILE_MASK + g_i = tidx_g & EPILOGUE_ROW_MASK + g_j = (tidx_g >> EPILOGUE_COLUMN_GROUP_SHIFT) * EPILOGUE_THREAD_TILE_THREADS + + logical_num_heads_q = Int32(self.logical_num_heads_q) + physical_tile_rows = Int32(cfg.mma_qk_tiler[0]) + D = cfg.latent_dim + head_tile_idx = blk_coord[0] + seq_q_idx = blk_coord[1] + batch_idx = blk_coord[2] + split_kv_idx = blk_coord[3] + row_in_tile = head_tile_idx * tile_h + g_i + ( + storage_flat_query_row, + logical_head_idx, + logical_q_idx, + _, + query_is_valid, + ) = self._query_row_state(row_in_tile, seq_q_idx, batch_idx) + + if row_in_tile < physical_tile_rows and query_is_valid: + if cutlass.const_expr(self.partial_output is not None): + S_q = ( + cutlass.Int32(self.partial_output.shape[3]) + if self.partial_output is not None + else Int32(1) + ) + o_base_ptr = ( + self.partial_output.iterator.raw_ptr() + + Int64(row_in_tile) * Int64(self.split_kv) * Int64(D) + + Int64(split_kv_idx) * Int64(D) + + Int64(seq_q_idx) + * Int64(self.split_kv) + * Int64(physical_tile_rows) + * Int64(D) + + Int64(batch_idx) + * Int64(physical_tile_rows) + * Int64(self.split_kv) + * Int64(S_q) + * Int64(D) + ) + output_base = o_base_ptr + iter_n * tile_d + g_j + for load_idx in cutlass.range_constexpr(num_tmem_loads): + for j in cutlass.range_constexpr(4): + offset = ( + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * BF16_OUTPUT_VECTOR_ELEMENTS + ) + vec_f32 = qk_acc_regs.load(offset, BF16_OUTPUT_VECTOR_ELEMENTS) + vec_partial = vec_f32.to(cutlass.BFloat16) + ( + output_base + + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * BF16_OUTPUT_VECTOR_ELEMENTS + ).nvvm_store_ext( + vec_partial, + evict="noallocate", + ) + else: + if cutlass.const_expr(self.cu_seqlens_q is not None): + o_base_ptr = self.output.iterator.raw_ptr() + Int64( + storage_flat_query_row + ) * Int64(D) + else: + S_q = ( + cutlass.Int32(self.output.shape[2]) + if self.output is not None + else Int32(1) + ) + o_base_ptr = ( + self.output.iterator.raw_ptr() + + Int64(logical_head_idx) * Int64(D) + + Int64(logical_q_idx) * Int64(logical_num_heads_q) * Int64(D) + + Int64(batch_idx) + * Int64(logical_num_heads_q) + * Int64(D) + * Int64(S_q) + ) + output_base = o_base_ptr + iter_n * tile_d + g_j + for load_idx in cutlass.range_constexpr(num_tmem_loads): + for j in cutlass.range_constexpr(2): + offset = ( + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * FP8_OUTPUT_VECTOR_ELEMENTS + ) + vec_f32 = qk_acc_regs.load(offset, FP8_OUTPUT_VECTOR_ELEMENTS) + if cutlass.const_expr(cfg.use_fp8_output == 1): + packed_o = cutlass.Array( + Int32, + PACKED_FP8_OUTPUT_REGS, + space=cutlass.AddressSpace.rmem, + ) + for pack_idx in cutlass.range_constexpr( + PACKED_FP8_OUTPUT_REGS + ): + pack_offset = pack_idx * PACKED_FP8_OUTPUT_REGS + packed_o[pack_idx] = pack_float4_to_fp8_e4m3( + vec_f32[pack_offset], + vec_f32[pack_offset + 1], + vec_f32[pack_offset + 2], + vec_f32[pack_offset + 3], + ) + raw_ptr = cutlass.inttoptr( + ( + output_base + + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * FP8_OUTPUT_VECTOR_ELEMENTS + ).toint(Int64), + mem_space=1, + dtype=Int32, + ) + raw_ptr.store( + packed_o.load(0, PACKED_FP8_OUTPUT_REGS), + alignment=16, + ) + else: + vec_o = vec_f32.to(output_dtype(self.cfg)) + ( + output_base + + load_idx * TCGEN05_32B_REGS_PER_LOAD + + j * FP8_OUTPUT_VECTOR_ELEMENTS + ).nvvm_store_ext( + vec_o, + evict="noallocate", + ) + + if cutlass.const_expr(iter_n == 0): + lse_row_sum = row_sum + if cutlass.const_expr(cfg.is_fp8_qkv()): + lse_row_sum = lse_row_sum * fp8_quant_scale_rcp() + lse = ( + cute.math.log2(lse_row_sum, fastmath=True) + + self.softmax_scale_log2 * row_max + if row_has_values + else Float32(-Float32.inf) + ) + lse_tidx = local_tidx + if lse_tidx < tile_h: + lse_row_in_tile = head_tile_idx * tile_h + lse_tidx + ( + storage_flat_lse_row, + logical_lse_head_idx, + logical_lse_q_idx, + _, + lse_query_is_valid, + ) = self._query_row_state(lse_row_in_tile, seq_q_idx, batch_idx) + if lse_row_in_tile < physical_tile_rows and lse_query_is_valid: + if cutlass.const_expr(self.partial_lse is not None): + S_q = ( + cutlass.Int32(self.partial_lse.shape[2]) + if self.partial_lse is not None + else Int32(1) + ) + lse_base_ptr = ( + self.partial_lse.iterator.raw_ptr() + + Int64(lse_row_in_tile) * Int64(self.split_kv) + + Int64(split_kv_idx) + + Int64(seq_q_idx) + * Int64(physical_tile_rows) + * Int64(self.split_kv) + + Int64(batch_idx) + * Int64(physical_tile_rows) + * Int64(self.split_kv) + * Int64(S_q) + ) + lse_base_ptr.store(lse) + elif cutlass.const_expr(self.lse is not None): + if cutlass.const_expr(self.cu_seqlens_q is not None): + lse_base_ptr = ( + self.lse.iterator.raw_ptr() + storage_flat_lse_row + ) + else: + S_q = ( + cutlass.Int32(self.lse.shape[1]) + if self.lse is not None + else Int32(1) + ) + lse_base_ptr = ( + self.lse.iterator.raw_ptr() + + Int64(logical_lse_head_idx) + + Int64(logical_lse_q_idx) * Int64(logical_num_heads_q) + + Int64(batch_idx) + * Int64(logical_num_heads_q) + * Int64(S_q) + ) + lse_base_ptr.store(lse) + + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py new file mode 100644 index 000000000000..dba6349c05d1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py @@ -0,0 +1,1298 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task definitions for the throughput 2CTA MLA decode TS kernel. + +The selected graph depends on dtype and scheduler policy. BF16 uses a 12-warp +combined TMA/MMA graph, with an additional scheduler task under CLC. FP8 uses a +16-warp split K/V and QK/PV graph with two softmax groups. + +Domain semantics: domain = k_tile_count (total k-tiles to process). +Tasks with domain_start=1 handle the first k-tile in HEAD, then LOOP +covers k-tiles 1..N-1, and TAIL handles cleanup. + +Stagger pattern (K loads 1 ahead of V, QK MMA 1 ahead of PV MMA): + k-tile 0: load K[0], QK MMA -> S[0] + k-tile 1: load K[1]+V[0], PV MMA(P[0],V[0])->O[0], QK MMA->S[1] + ... + k-tile N-1: load K[N-1]+V[N-2], PV MMA(P[N-2],V[N-2])->O[N-2], QK MMA->S[N-1] + tail: load V[N-1], PV MMA(P[N-1],V[N-1])->O[N-1] + +The BF16 register roles are: + LoadTmaTask (warp 9, 1 warp, 96 regs; owns page-ID window) + MmaTask (warp 8, 1 warp, 96 regs) + SoftmaxTask (warps 0-3, 4 warps, 192 regs) + CorrectionTask (warps 4-7, 4 warps, 208 regs) + PaddingTask (warps 10-11, 96 regs; non-CLC alignment) + SchedulerTask (warp 11, 96 regs; CLC only) +""" + +from collections.abc import Callable + +import cutlass +import cutlass.cute as cute + +from cutlass.experimental.task_scheduling.memory import ResourceContext +from cutlass.experimental.task_scheduling.schedule_builder import ( + domain_loop, + schedule, + work_tile_loop, +) +from cutlass.experimental.task_scheduling.resources import StageInfo +from cutlass.experimental.task_scheduling.task import Task + +from .resources import ( + MlaWorkQueue, + WorkThrottleBarrierResource, + PageOffsetWindowResource, + SmemQResource, + SmemKResource, + SmemKVResource, + SmemVResource, + SmemPResource, + TmemSResource, + TmemCorrResource, + TmemOResource, + GmemOResource, +) +from ..helpers.schedule import ( + captured_loop_bounds, + staged_kv_tma_load, + staged_pv_mma, + staged_pv_mma_v_tile, + staged_pv_mma_v_tile_per_n, + staged_qk_mma, + staged_qk_mma_k_tile, + work_queue_tail, +) + + +def _fixed_lane_work_tile_bounds(fixed_lane, cumulative_k_parity, actual_domain): + """Model fixed-lane work-tile bounds for CPU contract tests. + + Generated code keeps this arithmetic inline because returning staged DSL + values through a plain-Python tuple breaks staged-frontend state threading. + """ + + mapped_start = (fixed_lane + cumulative_k_parity) % 2 + lane_iterations = (actual_domain - mapped_start + 1) // 2 + fixed_loop_end = fixed_lane + lane_iterations * 2 + return mapped_start, fixed_loop_end + + +def _capture_clc_work_tile_body( + work_queue, + body: Callable[..., None], + non_skippable_prelude: Callable[[], object] | None = None, + *, + use_clc_dynamic: bool = False, +) -> None: + """Capture one MLA tile with data work skippable and WQ progress mandatory. + + Pure register-state initializers may run in ``non_skippable_prelude`` so + values they create dominate the separately guarded HEAD, LOOP, and TAIL + regions emitted by the stock skipped-tile executor. The prelude must not + issue memory operations or advance pipeline state. + """ + + def run_body(): + prelude_state = ( + non_skippable_prelude() if non_skippable_prelude is not None else None + ) + if non_skippable_prelude is None: + body() + else: + body(prelude_state) + + if work_queue is not None and use_clc_dynamic: + with work_tile_loop( + work_queue, + skip_if=MlaWorkQueue.skip_work_tile_if, + ) as work_tiles: + prelude_state = ( + non_skippable_prelude() if non_skippable_prelude is not None else None + ) + with work_tiles.skippable(): + if non_skippable_prelude is None: + body() + else: + body(prelude_state) + work_queue_tail(work_queue, advance_label="advance_tile") + return + + run_body() + work_queue_tail(work_queue, advance_label="advance_tile") + + +class MlaClcTask(Task): + """Stock Task persistent loop with an MLA-specific dynamic K domain.""" + + @cute.jit + def get_domain(self, tile_coord): + """Recompute the loop bound required by the stock Task public API.""" + + assert isinstance(self.work_queue, MlaWorkQueue) + return self.work_queue.k_tile_count_for_tile(tile_coord) + + +class MlaTask(Task): + """Task subclass that recomputes MLA k-domain per persistent work tile.""" + + @cute.jit + def _run_one_mla_work_tile( + self, + work_tile, + context: ResourceContext | None = None, + ) -> None: + """Run a persistent work tile using the cached split-KV domain.""" + + # WorkQueue decomposes the MLA persistent tile and caches the K-domain. + # Keep task bodies on that cached value so page-offset/TMA/MMA/softmax + # paths do not each rebuild the same split-KV arithmetic. + self.domain = work_tile.k_tile_count + self._run_task_body_impl(work_tile, context=context) + + @cute.jit + def _drain_mla_work_tile_tails(self) -> None: + """Drain producer tails after a persistent work-tile body completes.""" + + for resource in self.dst_resources: + if cutlass.const_expr( + resource.pipeline_config is not None + and resource is not self.work_queue + and not self._is_fork_secondary(resource) + ): + if cutlass.const_expr( + resource.pipeline_config.producer_acquire_interleave_stride > 1 + or resource.pipeline_config.producer_commit_interleave_stride > 1 + ): + # Interleaved producers own lane-specific pipeline states. + # The generic producer tail still drains at resource + # granularity, so calling it here can wait on a peer lane's + # physical stages. Consumer wait/release drains the live + # lane for these score/P resources. + pass + else: + self._producer_tail(resource) + if cutlass.const_expr( + self.work_queue is not None + and self.work_queue in self.dst_resources + and self.work_queue.pipeline_config is not None + ): + self.work_queue.producer_tail() + + @cute.jit + def _run_task_body_persistent( + self, + context: ResourceContext | None = None, + ) -> None: + """Schedule one or more persistent work tiles for this task instance.""" + + params = self.work_queue.tile_sched_params + + if cutlass.const_expr(not params.is_persistent): + work_tile = self.work_queue._work_tile_from_block_idx(cute.arch.block_idx()) + self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) + self._run_pre_work_loop_entries(work_tile, context) + # Runtime K/Q metadata can make a statically launched split empty. + # Keep the CTA on the ordinary initialized-pipeline path, but skip + # its captured HEAD/LOOP/TAIL data work when the domain is zero. + if work_tile.k_tile_count > cutlass.Int32(0): + self._run_one_mla_work_tile(work_tile, context) + self._run_post_work_loop_entries(work_tile, context) + self._drain_mla_work_tile_tails() + return + + current_work_linear_idx = cute.arch.block_idx()[0] + num_blocks = ( + params.cluster_shape_mnk[0] + * params.problem_shape_s + * params.problem_shape_b + * params.split_kv + ) + work_tile = self.work_queue._work_tile_from_linear_idx(current_work_linear_idx) + self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) + + self._run_pre_work_loop_entries(work_tile, context) + while current_work_linear_idx < num_blocks: + work_tile.update_from( + self.work_queue._work_tile_from_linear_idx(current_work_linear_idx) + ) + self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) + + # Variable K and causal Q visibility can leave individual logical + # splits empty. Skip them and continue grid-striding rather than + # running a captured HEAD/TAIL sequence with domain zero. + if work_tile.k_tile_count > cutlass.Int32(0): + self._run_one_mla_work_tile(work_tile, context) + + # Each warp branch advances from the same scalar tile id, keeping + # the persistent loop state compact across task bodies. + current_work_linear_idx += cute.size(cute.arch.grid_dim()) + # self.dummy keeps the captured persistent loop body live even when + # a specialized task instance has no visible local result. + self.dummy = True + work_tile.update_from( + self.work_queue._work_tile_from_linear_idx(current_work_linear_idx) + ) + self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) + self._run_post_work_loop_entries(work_tile, context) + self._drain_mla_work_tile_tails() + + +class MlaInterleavedTask(MlaTask): + """Persistent task that carries its interleave lane across work tiles.""" + + def __init__(self, *args, **kwargs) -> None: + """Create staged parity and index-remapping state for this task.""" + + super().__init__(*args, **kwargs) + self._cumulative_k_parity = cutlass.Int32(0) + self._mapped_domain_start = cutlass.Int32(self.domain_start) + self._actual_domain = cutlass.Int32(0) + self._fixed_loop_end = cutlass.Int32(self.domain_start) + + @cute.jit + def _run_one_mla_work_tile( + self, + work_tile, + context: ResourceContext | None = None, + ) -> None: + """Run one tile by mapping its local offsets onto this fixed lane.""" + + fixed_lane = cutlass.Int32(self.domain_start) + self._actual_domain = work_tile.k_tile_count + self._mapped_domain_start = ( + fixed_lane + self._cumulative_k_parity + ) % cutlass.Int32(2) + lane_iterations = ( + self._actual_domain - self._mapped_domain_start + cutlass.Int32(1) + ) // cutlass.Int32(2) + self._fixed_loop_end = fixed_lane + lane_iterations * cutlass.Int32(2) + self.domain = self._fixed_loop_end + self._run_task_body_impl(work_tile, context=context) + self._cumulative_k_parity = ( + self._cumulative_k_parity + self._actual_domain + ) % cutlass.Int32(2) + + @cute.jit + def _create_stage_info( + self, + resource, + idx, + work_tile=None, + is_producer=None, + resolved_domain=None, + label=None, + schedule_stage=None, + routing_slot=None, + context: ResourceContext | None = None, + ) -> StageInfo: + """Return base pipeline state with the work tile's actual K offset.""" + + base_info = Task._create_stage_info( + self, + resource, + idx, + work_tile, + is_producer, + resolved_domain, + label, + schedule_stage, + routing_slot, + context=context, + ) + actual_loop_offset = self._mapped_domain_start + ( + cutlass.Int32(base_info.loop_offset) - cutlass.Int32(self.domain_start) + ) + return StageInfo( + loop_offset=actual_loop_offset, + loop_start=self._mapped_domain_start, + loop_end=self._actual_domain, + loop_step=base_info.loop_step, + stage_idx=base_info.stage_idx, + label=base_info.label, + barrier=base_info.barrier, + work_tile=base_info.work_tile, + num_active_stages=base_info.num_active_stages, + context=base_info.context, + task_cache=base_info.task_cache, + ) + + +def create_load_tma_task( + page_offset_window: PageOffsetWindowResource, + smem_q: SmemQResource, + smem_kv: SmemKVResource, + iterations_qk: int = 9, + iterations_pv: int = 8, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the TMA load task (warp 9, 1 warp, 96 regs). + + domain_start=1: HEAD handles k-tile[0], LOOP handles k-tiles 1..N-1. + + HEAD: consume page offsets[0], load Q, load K[0]. + LOOP: consume page offsets[n], load K[n], load V[n-1]. + TAIL: load V[last]. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 1) + use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + + def load_tma_prelude(page_offset_window, smem_q, smem_kv): + """Create register and descriptor state before any dynamic skip guard.""" + cached_page_state = page_offset_window.init_read_state() + smem_q.init_load_state() + smem_kv.init_load_state() + return cached_page_state + + def load_tma_body(page_offset_window, smem_q, smem_kv, cached_page_state): + """Load Q once and load K/V tiles with the K-before-V cadence.""" + cached_k_pages, cached_v_pages, cached_next_v_pages, cached_window_page = ( + cached_page_state + ) + + # HEAD: Q is independent of page IDs. Enqueue it before the TMA warp + # refreshes its first coalesced 32-entry page-table window. + smem_q.acquire() + smem_q.tma_load() + smem_q.commit() + cached_k_pages, cached_v_pages, cached_next_v_pages, cached_window_page = ( + page_offset_window.read_page_offset_window( + cached_k_pages=cached_k_pages, + cached_v_pages=cached_v_pages, + cached_next_v_pages=cached_next_v_pages, + cached_window_page=cached_window_page, + init_v_cache=True, + ) + ) + staged_kv_tma_load( + smem_kv, + iterations_qk, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + is_v=False, + ) + + with domain_loop(loop_start, loop_end, loop_step): + # LOOP: cache K[n]/V[n] offsets, load K[n], then deferred V[n-1]. + cached_k_pages, cached_v_pages, cached_next_v_pages, cached_window_page = ( + page_offset_window.read_page_offset_window( + cached_k_pages=cached_k_pages, + cached_v_pages=cached_v_pages, + cached_next_v_pages=cached_next_v_pages, + cached_window_page=cached_window_page, + ) + ) + # K sub-tiles then deferred V sub-tiles, each with its own local + # sub-tile index; ``is_v`` tells the loader which path to take. + staged_kv_tma_load( + smem_kv, + iterations_qk, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + is_v=False, + ) + staged_kv_tma_load( + smem_kv, + iterations_pv, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + is_v=True, + ) + + # TAIL: V[last] uses the next-V offsets cached by the final wait. + cached_k_pages, cached_v_pages, cached_next_v_pages = ( + page_offset_window.forward_page_ids( + cached_k_pages=cached_k_pages, + cached_v_pages=cached_v_pages, + cached_next_v_pages=cached_next_v_pages, + ) + ) + staged_kv_tma_load( + smem_kv, + iterations_pv, + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + is_v=True, + use_next_v_pages=True, + ) + + @schedule + def load_tma_schedule(page_offset_window, smem_q, smem_kv, work_queue=None): + """Capture one active TMA tile and unconditional queue progress.""" + + _capture_clc_work_tile_body( + work_queue, + lambda cached_page_state: load_tma_body( + page_offset_window, + smem_q, + smem_kv, + cached_page_state, + ), + lambda: load_tma_prelude(page_offset_window, smem_q, smem_kv), + use_clc_dynamic=use_clc_dynamic, + ) + + if work_queue is None: + captured_schedule = load_tma_schedule(page_offset_window, smem_q, smem_kv) + else: + captured_schedule = load_tma_schedule( + page_offset_window, + smem_q, + smem_kv, + work_queue, + ) + + src = [page_offset_window] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[smem_q, smem_kv], + warp_idx=9, + num_warps=1, + schedule=captured_schedule, + name="LoadTmaTask", + **task_kwargs, + ) + + +def create_load_k_task( + smem_q: SmemQResource, + smem_k: SmemKResource, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the FP8 Q/K TMA task (warp 9). + + Q is loaded once before the loop. K uses one whole-tile pipeline stage per + logical K tile and reads page offsets directly from GMEM. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + + @schedule + def load_k_schedule(smem_q, smem_k, work_queue=None): + """Load Q once and then publish one K stage per loop tile.""" + smem_q.init_load_state() + smem_k.init_load_state() + smem_q.acquire() + smem_q.tma_load() + smem_q.commit() + + with domain_loop(loop_start, loop_end, loop_step): + smem_k.acquire() + smem_k.tma_load_direct() + smem_k.commit() + work_queue_tail(work_queue, advance_label="advance_tile") + + schedule_result = ( + load_k_schedule(smem_q, smem_k) + if work_queue is None + else load_k_schedule(smem_q, smem_k, work_queue) + ) + + src = [work_queue] if work_queue is not None else [] + return task_class( + src_resources=src, + dst_resources=[smem_q, smem_k], + warp_idx=9, + num_warps=1, + schedule=schedule_result, + name="LoadKTask", + **task_kwargs, + ) + + +def create_load_v_task( + smem_v: SmemVResource, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the FP8 V TMA task (warp 10).""" + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + + @schedule + def load_v_schedule(smem_v, work_queue=None): + """Publish one V stage per logical K tile.""" + smem_v.init_load_state() + with domain_loop(loop_start, loop_end, loop_step): + smem_v.acquire() + smem_v.tma_load_direct() + smem_v.commit() + work_queue_tail(work_queue, advance_label="advance_tile") + + schedule_result = ( + load_v_schedule(smem_v) + if work_queue is None + else load_v_schedule(smem_v, work_queue) + ) + + src = [work_queue] if work_queue is not None else [] + return task_class( + src_resources=src, + dst_resources=[smem_v], + warp_idx=10, + num_warps=1, + schedule=schedule_result, + name="LoadVTask", + **task_kwargs, + ) + + +def create_mma_task( + smem_q: SmemQResource, + smem_kv: SmemKVResource, + smem_p: SmemPResource, + tmem_s: TmemSResource, + tmem_o: TmemOResource, + iterations_qk: int = 9, + iterations_pv: int = 8, + work_queue: MlaWorkQueue = None, + work_throttle: WorkThrottleBarrierResource = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the MMA task (warp 8, 1 warp, 96 regs). + + domain_start=1: HEAD handles first QK MMA, LOOP handles PV+QK pairs. + run_only_on_cta_id=0 keeps the schedule on the CTA-pair leader, matching + the 2CTA UMMA contract. + + HEAD: consume Q, QK MMA for k-tile[0] -> S[0]. + LOOP: PV MMA for k-tile[n-1] -> O[n-1], then QK MMA for k-tile[n] -> S[n]. + TAIL: PV MMA for last k-tile -> O[last], release Q. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 1) + use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + + def mma_body( + smem_q, + smem_kv, + smem_p, + tmem_s, + tmem_o, + work_throttle=None, + ): + """Run QK one tile ahead of PV and keep UMMA on the leader CTA.""" + # HEAD: wait Q once and compute S[0] from K[0]. + smem_q.wait() + if work_throttle is not None: + # The leader MMA reaching Q proves that this cluster has started + # the current tile. Permit the scheduler to prepare one more work + # ID without relying on CTA-private resource-ownership fields. + work_throttle.try_acquire() + work_throttle.acquire() + work_throttle.commit() + smem_q.q_desc() + staged_qk_mma(smem_kv, tmem_s, iterations_qk) + + with domain_loop(loop_start, loop_end, loop_step): + # LOOP: compute S[n] before PV[n-1]. This gives softmax the QK + # instruction window to produce P[n-1] before PV consumes it. + staged_qk_mma(smem_kv, tmem_s, iterations_qk) + staged_pv_mma( + smem_kv, + smem_p, + tmem_o, + iterations_pv, + ) + + # TAIL: finish PV[last], then release the Q descriptor stage. + staged_pv_mma( + smem_kv, + smem_p, + tmem_o, + iterations_pv, + is_tail=True, + ) + smem_q.release() + + @schedule + def mma_schedule( + smem_q, + smem_kv, + smem_p, + tmem_s, + tmem_o, + work_queue=None, + work_throttle=None, + ): + """Capture one active MMA tile and unconditional queue progress.""" + + _capture_clc_work_tile_body( + work_queue, + lambda: mma_body( + smem_q, + smem_kv, + smem_p, + tmem_s, + tmem_o, + work_throttle, + ), + use_clc_dynamic=use_clc_dynamic, + ) + + if work_queue is None: + captured_schedule = mma_schedule(smem_q, smem_kv, smem_p, tmem_s, tmem_o) + elif work_throttle is None: + captured_schedule = mma_schedule( + smem_q, + smem_kv, + smem_p, + tmem_s, + tmem_o, + work_queue, + ) + else: + captured_schedule = mma_schedule( + smem_q, + smem_kv, + smem_p, + tmem_s, + tmem_o, + work_queue, + work_throttle, + ) + + src = [smem_q, smem_kv, smem_p] + if work_queue is not None: + src.append(work_queue) + dst = [tmem_s, tmem_o] + if work_throttle is not None: + dst.append(work_throttle) + return task_class( + src_resources=src, + dst_resources=dst, + warp_idx=8, + num_warps=1, + schedule=captured_schedule, + name="MmaTask", + run_only_on_cta_id=0, + **task_kwargs, + ) + + +def create_mma_qk_task( + smem_q: SmemQResource, + smem_kv: SmemKVResource, + tmem_s: TmemSResource, + iterations_qk: int = 9, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the FP8 QK-only MMA task (warp 8). + + QK stays one k-tile ahead of PV. Keeping QK and PV on separate warps avoids + serializing the two UMMA issue streams while preserving the existing K-before-V + TMA cadence and one-softmax schedule. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 1) + + @schedule + def mma_qk_schedule(smem_q, smem_kv, tmem_s, work_queue=None): + """Consume Q/K stages and publish S for every k-tile.""" + smem_q.wait() + smem_q.q_desc() + staged_qk_mma(smem_kv, tmem_s, iterations_qk) + + with domain_loop(loop_start, loop_end, loop_step): + staged_qk_mma(smem_kv, tmem_s, iterations_qk) + + smem_q.release() + work_queue_tail(work_queue, advance_label="advance_tile") + + schedule_result = ( + mma_qk_schedule(smem_q, smem_kv, tmem_s) + if work_queue is None + else mma_qk_schedule(smem_q, smem_kv, tmem_s, work_queue) + ) + + src = [smem_q, smem_kv] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_s], + warp_idx=8, + num_warps=1, + schedule=schedule_result, + name="MmaQkTask", + run_only_on_cta_id=0, + **task_kwargs, + ) + + +def create_mma_pv_task( + smem_kv: SmemKVResource, + smem_p: SmemPResource, + tmem_o: TmemOResource, + iterations_pv: int = 8, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the FP8 PV-only MMA task (warp 11). + + The PV task consumes P[n-1]/V[n-1] while QK produces S[n], then handles the + final P/V tile in TAIL. This matches the existing correction schedule. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 1) + + @schedule + def mma_pv_schedule(smem_kv, smem_p, tmem_o, work_queue=None): + """Consume delayed V and P stages and accumulate O.""" + with domain_loop(loop_start, loop_end, loop_step): + staged_pv_mma(smem_kv, smem_p, tmem_o, iterations_pv) + + staged_pv_mma( + smem_kv, + smem_p, + tmem_o, + iterations_pv, + is_tail=True, + ) + work_queue_tail(work_queue, advance_label="advance_tile") + + schedule_result = ( + mma_pv_schedule(smem_kv, smem_p, tmem_o) + if work_queue is None + else mma_pv_schedule(smem_kv, smem_p, tmem_o, work_queue) + ) + + src = [smem_kv, smem_p] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_o], + warp_idx=11, + num_warps=1, + schedule=schedule_result, + name="MmaPvTask", + run_only_on_cta_id=0, + **task_kwargs, + ) + + +def create_mma_qk_direct_task( + smem_q: SmemQResource, + smem_k: SmemKResource, + tmem_s: TmemSResource, + iterations_qk: int = 5, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the FP8 QK task with one K stage per domain-loop iteration.""" + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + + @schedule + def mma_qk_direct_schedule(smem_q, smem_k, tmem_s, work_queue=None): + """Consume Q once and publish one S stage per K tile.""" + smem_q.wait() + smem_q.q_desc() + with domain_loop(loop_start, loop_end, loop_step): + staged_qk_mma_k_tile(smem_k, tmem_s, iterations_qk) + smem_q.release() + work_queue_tail(work_queue, advance_label="advance_tile") + + schedule_result = ( + mma_qk_direct_schedule(smem_q, smem_k, tmem_s) + if work_queue is None + else mma_qk_direct_schedule(smem_q, smem_k, tmem_s, work_queue) + ) + + src = [smem_q, smem_k] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_s], + warp_idx=8, + num_warps=1, + schedule=schedule_result, + name="MmaQkTask", + run_only_on_cta_id=0, + **task_kwargs, + ) + + +def create_mma_pv_direct_task( + smem_v: SmemVResource, + smem_p: SmemPResource, + tmem_o: TmemOResource, + iterations_pv: int = 4, + iterations_pv_k: int = 2, + iterations_pv_n: int = 2, + per_n_o_pipeline: bool = False, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the FP8 PV task with one V stage per domain-loop iteration.""" + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + + @schedule + def mma_pv_direct_schedule(smem_v, smem_p, tmem_o, work_queue=None): + """Consume one P/V pair and publish one O stage per K tile.""" + with domain_loop(loop_start, loop_end, loop_step): + if per_n_o_pipeline: + staged_pv_mma_v_tile_per_n( + smem_v, + smem_p, + tmem_o, + iterations_pv_k=iterations_pv_k, + iterations_pv_n=iterations_pv_n, + ) + else: + staged_pv_mma_v_tile(smem_v, smem_p, tmem_o, iterations_pv) + work_queue_tail(work_queue, advance_label="advance_tile") + + schedule_result = ( + mma_pv_direct_schedule(smem_v, smem_p, tmem_o) + if work_queue is None + else mma_pv_direct_schedule(smem_v, smem_p, tmem_o, work_queue) + ) + + src = [smem_v, smem_p] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_o], + warp_idx=11, + num_warps=1, + schedule=schedule_result, + name="MmaPvTask", + run_only_on_cta_id=0, + **task_kwargs, + ) + + +def create_softmax_task( + tmem_s: TmemSResource, + tmem_corr: TmemCorrResource, + smem_p: SmemPResource, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + warp_idx: int = 0, + name: str = "SoftmaxTask", + softmax_group_id: int = 0, + **task_kwargs, +) -> Task: + """Create the Softmax task (warps 0-3, 4 warps, 192 regs). + + domain_start=0: processes all k_tile_count S tiles. + + LOOP: consume S, compute softmax, produce correction factors + P. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + + def softmax_prelude(tmem_s): + """Create softmax register arrays before the dynamic skip guard.""" + init_softmax_state = ( + tmem_s.init_softmax_state_odd + if softmax_group_id == 1 + else tmem_s.init_softmax_state + ) + return init_softmax_state() + + def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): + """Consume S tiles, materialize P, and publish correction factors.""" + load_s = tmem_s.load_s_odd if softmax_group_id == 1 else tmem_s.load_s + finish_softmax = ( + tmem_s.finish_softmax_odd + if softmax_group_id == 1 + else tmem_s.finish_softmax + ) + finish_row_sum = ( + tmem_s.finish_row_sum_odd + if softmax_group_id == 1 + else tmem_s.finish_row_sum + ) + store_corr = ( + tmem_corr.store_corr_odd if softmax_group_id == 1 else tmem_corr.store_corr + ) + store_p = smem_p.store_p_odd if softmax_group_id == 1 else smem_p.store_p + ( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ) = softmax_state + + with domain_loop(loop_start, loop_end, loop_step): + tmem_s.wait() + if softmax_group_id == 1: + ( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ) = load_s( + qk_acc_regs_odd=qk_acc_regs, + row_max_odd=row_max, + row_sum_odd=row_sum, + row_sum_out_odd=row_sum_out, + row_max_new_odd=row_max_new, + correction_factor_out_odd=correction_factor_out, + no_correction_out_odd=no_correction_out, + ) + else: + ( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ) = load_s( + qk_acc_regs=qk_acc_regs, + row_max=row_max, + row_sum=row_sum, + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + ) + tmem_s.release() + if softmax_group_id == 1: + ( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ) = finish_softmax( + qk_acc_regs_odd=qk_acc_regs, + row_max_odd=row_max, + row_sum_odd=row_sum, + row_sum_out_odd=row_sum_out, + row_max_new_odd=row_max_new, + correction_factor_out_odd=correction_factor_out, + no_correction_out_odd=no_correction_out, + ) + else: + ( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ) = finish_softmax( + qk_acc_regs=qk_acc_regs, + row_max=row_max, + row_sum=row_sum, + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + ) + smem_p.acquire() + if softmax_group_id == 1: + store_p(qk_acc_regs_odd=qk_acc_regs) + else: + store_p(qk_acc_regs=qk_acc_regs) + smem_p.commit() + if softmax_group_id == 1: + row_sum, row_sum_out = finish_row_sum( + qk_acc_regs_odd=qk_acc_regs, + row_sum_odd=row_sum, + correction_factor_out_odd=correction_factor_out, + ) + else: + row_sum, row_sum_out = finish_row_sum( + qk_acc_regs=qk_acc_regs, + row_sum=row_sum, + correction_factor_out=correction_factor_out, + ) + tmem_corr.acquire() + if softmax_group_id == 1: + store_corr( + row_sum_out_odd=row_sum_out, + row_max_new_odd=row_max_new, + correction_factor_out_odd=correction_factor_out, + no_correction_out_odd=no_correction_out, + ) + else: + store_corr( + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + ) + tmem_corr.commit() + + @schedule + def softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue=None): + """Capture one active softmax tile and unconditional queue progress.""" + + _capture_clc_work_tile_body( + work_queue, + lambda softmax_state: softmax_body( + tmem_s, + tmem_corr, + smem_p, + softmax_state, + ), + lambda: softmax_prelude(tmem_s), + use_clc_dynamic=use_clc_dynamic, + ) + + schedule_result = ( + softmax_schedule(tmem_s, tmem_corr, smem_p) + if work_queue is None + else softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue) + ) + + src = [tmem_s] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_corr, smem_p], + warp_idx=warp_idx, + num_warps=4, + schedule=schedule_result, + name=name, + **task_kwargs, + ) + + +def create_correction_task( + tmem_corr: TmemCorrResource, + tmem_o: TmemOResource, + gmem_o: GmemOResource, + iterations_pv_n: int = 1, + per_n_o_pipeline: bool = False, + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the Correction task (warps 4-7, 4 warps, 208 regs). + + domain_start=1: HEAD handles initial correction (no O yet), LOOP + handles correction+O pairs, TAIL handles final O + epilogue. + + TMEM visibility is ensured by kernel-level named barrier sync before + task_manager.run(), so o_init pipeline is no longer needed here. + + HEAD: consume Corr[0] (initial max/sum, no O rescaling). + LOOP: consume Corr[n] + O[n-1], rescale accumulated O. + TAIL: consume O[last], final epilogue store O + LSE to GMEM. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 1) + use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + + def correction_prelude(tmem_corr): + """Create correction and epilogue register state before skip guards.""" + return tmem_corr.init_load_state() + + def correction_body(tmem_corr, tmem_o, gmem_o, correction_state): + """Apply online-softmax correction and store the final O/LSE result.""" + row_sum, row_max, correction_factor, no_correction = correction_state + + # HEAD: consume the first correction factors. There is no O tile yet. + tmem_corr.wait() + row_sum, row_max, correction_factor, no_correction = tmem_corr.load_corr() + tmem_corr.release() + + with domain_loop(loop_start, loop_end, loop_step): + # LOOP: consume Corr[n] and rescale O[n-1]. + tmem_corr.wait() + row_sum, row_max, correction_factor, no_correction = tmem_corr.load_corr() + tmem_corr.release() + if per_n_o_pipeline: + for iter_n in range(iterations_pv_n): + tmem_o.wait() + tmem_o.rescale_o_slice( + correction_factor=correction_factor, + no_correction=no_correction, + iter_n=iter_n, + ) + tmem_o.release() + else: + tmem_o.wait() + tmem_o.rescale_o( + correction_factor=correction_factor, + no_correction=no_correction, + ) + tmem_o.release() + + # TAIL: consume final O. Do not call rescale_o here; the correction was + # already applied in LOOP and the last loop correction value would be + # stale for a second application. epilogue_store also writes LSE. + if per_n_o_pipeline: + epilogue_row_sum, epilogue_row_max = ( + tmem_corr.prepare_epilogue_slice_store() + ) + for iter_n in range(iterations_pv_n): + tmem_o.wait() + gmem_o.epilogue_store_slice( + row_sum=epilogue_row_sum, + row_max=epilogue_row_max, + iter_n=iter_n, + ) + tmem_o.release() + else: + tmem_o.wait() + gmem_o.epilogue_store() + tmem_o.release() + + @schedule + def correction_schedule(tmem_corr, tmem_o, gmem_o, work_queue=None): + """Capture one active correction tile and unconditional queue progress.""" + + _capture_clc_work_tile_body( + work_queue, + lambda correction_state: correction_body( + tmem_corr, + tmem_o, + gmem_o, + correction_state, + ), + lambda: correction_prelude(tmem_corr), + use_clc_dynamic=use_clc_dynamic, + ) + + captured_schedule = ( + correction_schedule(tmem_corr, tmem_o, gmem_o) + if work_queue is None + else correction_schedule(tmem_corr, tmem_o, gmem_o, work_queue) + ) + + src = [tmem_corr, tmem_o] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[gmem_o], + warp_idx=4, + num_warps=4, + schedule=captured_schedule, + name="CorrectionTask", + **task_kwargs, + ) + + +def create_padding_task( + work_queue: MlaWorkQueue = None, + task_class: type = MlaTask, + warp_idx: int = 11, + num_warps: int = 1, + **task_kwargs, +) -> Task: + """Create a padding task for unused warps in producer warpgroup 2. + + Empty task for warp-group alignment. + """ + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + + @schedule + def padding_schedule(work_queue=None): + """Reserve unused producer warps without issuing kernel work.""" + + def padding_body(): + with domain_loop(loop_start, loop_end, loop_step): + pass + + _capture_clc_work_tile_body( + work_queue, + padding_body, + use_clc_dynamic=use_clc_dynamic, + ) + + captured_schedule = ( + padding_schedule() if work_queue is None else padding_schedule(work_queue) + ) + src = [work_queue] if work_queue is not None else [] + return task_class( + src_resources=src, + dst_resources=[], + warp_idx=warp_idx, + num_warps=num_warps, + schedule=captured_schedule, + name="PaddingTask", + **task_kwargs, + ) + + +def create_scheduler_task( + work_queue: MlaWorkQueue, + work_throttle: WorkThrottleBarrierResource = None, + task_class: type = MlaTask, + **task_kwargs, +) -> Task: + """Create the BF16 cluster-wide CLC scheduler on warp 11.""" + + @schedule + def scheduler_schedule(work_queue, work_throttle=None): + """Fetch and distribute the next logical MLA cluster tile.""" + + with work_tile_loop( + work_queue, + skip_if=MlaWorkQueue.skip_work_tile_if, + ) as work_tiles: + with work_tiles.skippable(): + with domain_loop(0, 0, 1): + pass + if work_throttle is not None: + work_throttle.wait() + work_throttle.release() + work_queue.acquire() + work_queue.fetch_work_tile() + work_queue.commit() + work_queue_tail(work_queue, advance_label="advance_tile") + + captured_schedule = ( + scheduler_schedule(work_queue) + if work_throttle is None + else scheduler_schedule(work_queue, work_throttle) + ) + src = [work_queue] + if work_throttle is not None: + src.append(work_throttle) + return task_class( + src_resources=src, + dst_resources=[work_queue], + warp_idx=11, + num_warps=1, + schedule=captured_schedule, + name="SchedulerTask", + run_only_on_cta_id=0, + **task_kwargs, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/work_partition.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/work_partition.py new file mode 100644 index 000000000000..a9448f32c431 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/work_partition.py @@ -0,0 +1,153 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime nonempty-prefix split-KV helpers for throughput 2CTA MLA.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + + +def _validate_partition_inputs(k_tile_total: int, split_kv_cap: int) -> None: + """Validate host-side split partition inputs.""" + + if k_tile_total < 0: + raise ValueError("k_tile_total must be non-negative") + if split_kv_cap <= 0: + raise ValueError("split_kv_cap must be positive") + + +def active_split_count(k_tile_total: int, split_kv_cap: int) -> int: + """Return the nonempty prefix under the configured-span partition.""" + + _validate_partition_inputs(k_tile_total, split_kv_cap) + if k_tile_total == 0: + return 0 + tiles_per_split = (k_tile_total + split_kv_cap - 1) // split_kv_cap + return (k_tile_total + tiles_per_split - 1) // tiles_per_split + + +def split_tile_range( + k_tile_total: int, + split_kv_cap: int, + split_idx: int, +) -> tuple[int, int]: + """Return the configured-span ``(start, count)`` for one split slot.""" + + _validate_partition_inputs(k_tile_total, split_kv_cap) + if split_idx < 0: + raise ValueError("split_idx must be non-negative") + tiles_per_split = (k_tile_total + split_kv_cap - 1) // split_kv_cap + start = split_idx * tiles_per_split + count = min(tiles_per_split, max(k_tile_total - start, 0)) + return (start, count) if count else (k_tile_total, 0) + + +def row_prefix_active_split_count( + row_k_tile_total: int, + group_k_tile_total: int, + split_kv_cap: int, +) -> int: + """Return how many configured-span splits intersect a row's K prefix.""" + + _validate_partition_inputs(group_k_tile_total, split_kv_cap) + if row_k_tile_total < 0: + raise ValueError("row_k_tile_total must be non-negative") + row_k_tile_total = min(row_k_tile_total, group_k_tile_total) + if row_k_tile_total == 0: + return 0 + + tiles_per_split = (group_k_tile_total + split_kv_cap - 1) // split_kv_cap + # ``row_k_tile_total`` is clipped to the group's K prefix, so its rounded + # split count cannot exceed the group's active split count. Avoid deriving + # that second quotient on the hot reducer path. + return (row_k_tile_total + tiles_per_split - 1) // tiles_per_split + + +@cute.jit +def runtime_split_kv_cap( + max_split_kv, + is_var_split_kv: cutlass.Constexpr[bool], + block_split_kvs, + batch_idx, +): + """Return a positive per-batch cap bounded by launch/workspace capacity.""" + + max_split_kv = cute.math.max(Int32(max_split_kv), Int32(1)) + split_kv_cap = max_split_kv + if cutlass.const_expr(is_var_split_kv): + split_kv_cap = Int32(block_split_kvs[batch_idx]) + return cute.math.max( + cute.math.min(split_kv_cap, max_split_kv), + Int32(1), + ) + + +@cute.jit +def runtime_active_split_count(k_tile_total, split_kv_cap): + """Device form of :func:`active_split_count`.""" + + k_tile_total = cute.math.max(Int32(k_tile_total), Int32(0)) + split_kv_cap = cute.math.max(Int32(split_kv_cap), Int32(1)) + tiles_per_split = cute.math.max( + (k_tile_total + split_kv_cap - Int32(1)) // split_kv_cap, + Int32(1), + ) + return (k_tile_total + tiles_per_split - Int32(1)) // tiles_per_split + + +@cute.jit +def runtime_split_tile_range(k_tile_total, split_kv_cap, split_idx): + """Device form of :func:`split_tile_range`.""" + + k_tile_total = cute.math.max(Int32(k_tile_total), Int32(0)) + split_idx = Int32(split_idx) + split_kv_cap = cute.math.max(Int32(split_kv_cap), Int32(1)) + tiles_per_split = cute.math.max( + (k_tile_total + split_kv_cap - Int32(1)) // split_kv_cap, + Int32(1), + ) + start = split_idx * tiles_per_split + count = cute.math.min( + tiles_per_split, + cute.math.max(k_tile_total - start, Int32(0)), + ) + start = start if count > Int32(0) else k_tile_total + return start, count + + +@cute.jit +def runtime_row_prefix_active_split_count( + row_k_tile_total, + group_k_tile_total, + split_kv_cap, +): + """Device form of :func:`row_prefix_active_split_count`.""" + + group_k_tile_total = cute.math.max(Int32(group_k_tile_total), Int32(0)) + row_k_tile_total = cute.math.max( + cute.math.min(Int32(row_k_tile_total), group_k_tile_total), + Int32(0), + ) + split_kv_cap = cute.math.max(Int32(split_kv_cap), Int32(1)) + tiles_per_split = cute.math.max( + (group_k_tile_total + split_kv_cap - Int32(1)) // split_kv_cap, + Int32(1), + ) + row_active_splits = ( + row_k_tile_total + tiles_per_split - Int32(1) + ) // tiles_per_split + # The row K domain was clipped to the group's prefix, making the extra + # group-active quotient and min redundant. + return row_active_splits diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/__init__.py new file mode 100644 index 000000000000..d63fca5389f8 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Throughput-latency 1CTA policy for task-scheduled MLA decode.""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py new file mode 100644 index 000000000000..dedc950131d2 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py @@ -0,0 +1,1637 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration traits for the throughput-latency 1CTA MLA TS path. + +This module contains Python-side trait state only. Shape policy is kept in the +explicit kernel selector so the executable schedule stays focused on one +concrete profile at a time. + +``make_throughput_latency_mla_config`` is the public factory: it validates user shapes, +page size, and explicit profile values before deriving task, pipeline, and +workspace traits. Invalid inputs fail with ``ValueError`` rather than falling +through to DSL division or GMEM reduction layout errors. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from math import ceil +from typing import Any + +from ....split_kv_mode_policy import select_split_kv_modes +from ..helpers.constants import MAX_MLA_SPLITS_KV, SUPPORTED_MLA_PAGE_SIZES +from ..helpers.mask import MaskType, normalize_mask_type +from ..helpers.query import FlatQueryTileLayout + + +# 1CTA profiles are specialized for these Q/head tiles. Tile sizes 8/16 use +# the swaps-MMA-AB schedule; tile sizes 32/64 are supported for explicit or +# keeps-MMA-AB profiles. +SUPPORTED_TILE_SIZE_Q = (8, 16, 32, 64) + +# SM100A exposes 227 KiB of dynamic SMEM. Keep a small TS metadata reservation +# when deciding whether the automatic cluster-reduction scratch can fit. +SM100A_SMEM_CAPACITY_BYTES = 232448 +ESTIMATED_TS_BARRIER_BYTES = 512 +MAX_CLUSTER_SMEM_DATA_BYTES = SM100A_SMEM_CAPACITY_BYTES - ESTIMATED_TS_BARRIER_BYTES +# CUDA's nonportable cluster launch attribute accepts at most 32 CTAs. This is +# a launch-geometry limit, not a problem-shape policy. +MAX_CLUSTER_SIZE = 32 + + +def align_up(value: int, alignment: int) -> int: + """Return ``value`` rounded up to ``alignment`` bytes.""" + return ((value + alignment - 1) // alignment) * alignment + + +def estimated_throughput_latency_smem_data_bytes(cfg: "MlaConfig") -> int: + """Estimate steady SMEM data allocation for the 1CTA resource set.""" + stensor = cfg.stensor_align + total = 0 + total += align_up(cfg.smem_q_tile_bytes * cfg.q_stages, stensor) + total += align_up(cfg.smem_kv_tile_bytes * cfg.kv_stages, stensor) + total += align_up( + cfg.page_offsets_stages * cfg.page_offsets_entries_per_stage * 4, + 128, + ) + total += 2 * align_up(cfg.p_smem_tile_bytes, stensor) + total += align_up(16, 8) # P0/P1 ordering barrier. + total += 2 * align_up(cfg.softmax_scratch_bytes, 16) + total += align_up(cfg.o_smem_tile_bytes, stensor) + total += align_up(cfg.corr_scratch_bytes, 16) + if cfg.cluster_reduction_smem_bytes: + total += align_up(cfg.cluster_reduction_smem_bytes, 64) + total += align_up(8, 8) + return total + + +@dataclass(frozen=True) +class MlaConfig: + """Kernel traits for one throughput-latency 1CTA MLA schedule. + + The dataclass groups shape dimensions, tile sizes, CTA decomposition, + resource capacities, warp/register budgets, and feature switches consumed by + the captured task graph. Fields are plain Python values so they can be passed + as ``Constexpr`` into JIT resources and tasks. + """ + + # DS MLA dimensions. Dense MLA uses 512 latent channels plus 64 RoPE + # channels, so the QK head dimension defaults to 576. + batch_size: int = 1 + num_heads_q: int = 16 + seq_len_q: int = 4 + seq_len_kv: int = 4096 + logical_num_heads_q: int = 16 + logical_seq_len_q: int = 4 + # Causal is bottom-right aligned for speculative decode. Dense keeps only + # the ordinary per-batch KV-tail predicate. + mask_type: str = MaskType.CAUSAL.value + head_dim_qk: int = 576 + head_dim_v: int = 512 + latent_dim: int = 512 + rope_dim: int = 64 + + # Data types. + qkv_dtype: str = "bf16" + o_dtype: str = "bf16" + qkv_dtype_bytes: int = 2 + o_dtype_bytes: int = 2 + acc_dtype_bytes: int = 4 + use_bf16_output: int = 1 + use_fp8_output: int = 0 + + # Mainloop tile shape for flat query-row decode. Each loop step consumes two + # 128-token KV tiles: K for the current score update and V for the delayed + # PV update. + tile_size_q: int = 16 + tile_size_kv: int = 128 + num_insts_q: int = 1 + num_insts_kv: int = 2 + q_stages: int = 2 + kv_stages: int = 4 + # Keep page-ID stages aligned with the straight-line K/V pipeline so each + # stage has an explicit, reusable metadata window. + page_offsets_stages: int = 6 + # Each pipeline stage describes one KV tile. The 32-entry capacity covers + # the smallest supported page size: 128 tokens / 16 tokens per page = 8. + page_offsets_entries_per_stage: int = 32 + o_stages: int = 2 + + # Default CTA decomposition for SM100-class decode. A 512-thread CTA + # provides 16 warps for the two softmax groups, correction, MMA, load, + # scheduler, and page-offset work. + threads_per_cta: int = 512 + num_ctas_per_seq_q: int = 4 + num_ctas_per_seq_kv: int = 1 + num_ctas_for_all_heads: int = 1 + num_ctas_per_head_dim: int = 1 + head_dim_per_cta_v: int = 512 + head_dim_per_stage_kv: int = 128 + head_dim_per_stage_v: int = 128 + num_tokens_per_page: int = 32 + max_num_pages_per_seq_kv: int = 1 + + # Shared-memory and TMEM allocation traits. STensor allocations are aligned + # to 1 KiB, and TMEM column counts follow the tcgen05 16-column granularity + # used for S/P/O tiles. + stensor_align: int = 1024 + tmem_s_cols: int = 16 + tmem_stats_cols: int = 32 + tmem_o_cols: int = 16 + + # Warp layout for load, MMA, softmax, correction, and scheduling work. Warp + # indices are CTA-local; four-warps groups are kept contiguous for + # setmaxnreg and named-barrier participation. + softmax0_warp_idx: int = 0 + softmax1_warp_idx: int = 4 + correction_warp_idx: int = 8 + mma_warp_idx: int = 12 + load_warp_idx: int = 15 + page_offsets_warp_idx: int = 13 + scheduler_warp_idx: int = 14 + softmax_num_warps: int = 4 + correction_num_warps: int = 4 + mma_num_warps: int = 1 + load_num_warps: int = 1 + page_offsets_num_warps: int = 1 + scheduler_num_warps: int = 1 + clc_padding_warp_idx: int = 14 + clc_padding_num_warps: int = 1 + + # Register budgets used by task-local setmaxregister calls. Softmax owns + # the score registers, correction runs with a small budget after softmax, + # and load/MMA/scheduler warps share the lower budget. + softmax_regs: int = 176 + correction_regs: int = 64 + mma_load_regs: int = 96 + scheduler_regs: int = 96 + + # Feature mode flags consumed by config and schedule construction. They are + # integer flags because many branches are passed as Constexpr values into + # JIT code. + kernel_variant: str = "swaps_mma_ab" + use_paged_kv: int = 1 + supports_var_seq_lens: int = 1 + use_persistent_scheduler: int = 1 + use_clc_dynamic_persistent_scheduler: int = 0 + use_multi_ctas_kv: int = 0 + use_cluster_reduction: int = 0 + persistent_wave_sm_count: int | None = None + use_attention_sinks: int = 0 + use_sliding_window_causal: int = 0 + attention_window_size: int = 0 + + @property + def softmax0_num_warps(self) -> int: + return self.softmax_num_warps + + @property + def softmax1_num_warps(self) -> int: + return self.softmax_num_warps + + @property + def padding_warp_idx(self) -> int: + return self.page_offsets_warp_idx + + @property + def padding_num_warps(self) -> int: + return 2 + + @property + def qk_smem_tile_bytes(self) -> int: + return self.tile_size_q * self.head_dim_qk * self.qkv_dtype_bytes + + @property + def kv_smem_tile_bytes(self) -> int: + return self.tile_size_kv * self.head_dim_per_stage_kv * self.qkv_dtype_bytes + + @property + def v_smem_tile_bytes(self) -> int: + return self.tile_size_kv * self.head_dim_per_stage_v * self.qkv_dtype_bytes + + @property + def pages_per_kv_tile(self) -> int: + """Return physical pages consumed by one logical KV tile.""" + + return ceil(self.tile_size_kv / self.num_tokens_per_page) + + @property + def total_kv_tiles(self) -> int: + return ceil(self.seq_len_kv / self.tile_size_kv) + + @property + def kv_tiles_per_multi_cta_group(self) -> int: + return self.num_ctas_per_seq_kv * self.num_insts_kv + + @property + def num_steps_per_cta_kv(self) -> int: + tokens_per_multi_cta_step = ( + self.num_ctas_per_seq_kv * self.num_insts_kv * self.tile_size_kv + ) + return ceil(self.seq_len_kv / tokens_per_multi_cta_step) + + @property + def p_smem_tile_bytes(self) -> int: + return self.tile_size_kv * self.tile_size_q * self.qkv_dtype_bytes + + @property + def smem_q_tile_bytes(self) -> int: + return self.qk_smem_tile_bytes + + @property + def smem_kv_tile_bytes(self) -> int: + return self.kv_smem_tile_bytes + + @property + def smem_p_tile_bytes(self) -> int: + return self.p_smem_tile_bytes + + @property + def partial_o_dtype_bytes(self) -> int: + if self.num_ctas_per_seq_kv > 1: + return 2 + return self.o_dtype_bytes + + @property + def o_smem_tile_bytes(self) -> int: + staging_dim = max(self.head_dim_per_stage_v, 64) + return self.tile_size_q * staging_dim * self.partial_o_dtype_bytes + + @property + def o_copy_segments_per_stage(self) -> int: + bytes_per_stage = ( + self.tile_size_q * self.head_dim_per_stage_v * self.partial_o_dtype_bytes + ) + return max(1, ceil(bytes_per_stage / 2048)) + + @property + def softmax_scratch_bytes(self) -> int: + return 4 * self.tile_size_q * self.acc_dtype_bytes + + @property + def corr_scratch_bytes(self) -> int: + return 8 * self.tile_size_q * self.acc_dtype_bytes + + @property + def cluster_reduction_rows_per_slice(self) -> int: + # One cluster-reduction SMEM slice is one 16-row x 128B tile. The row count + # shrinks as the per-CTA V head dimension grows. + num_bytes_per_slice = 128 * 16 + num_bytes_per_row_o = self.head_dim_per_cta_v * self.partial_o_dtype_bytes + return max(1, num_bytes_per_slice // num_bytes_per_row_o) + + @property + def cluster_reduction_slices(self) -> int: + return ceil(self.tile_size_q / self.cluster_reduction_rows_per_slice) + + def cluster_reduction_smem_bytes_for(self, num_ctas_kv: int) -> int: + """Return multi-CTA KV cluster reduction SMEM footprint.""" + if num_ctas_kv <= 1: + return 0 + rows_per_slice = self.cluster_reduction_rows_per_slice + num_slices = self.cluster_reduction_slices + num_slices_per_cta = ceil(num_slices / num_ctas_kv) + num_rows_per_cta = num_slices_per_cta * rows_per_slice + num_bytes_per_row_o = self.head_dim_per_cta_v * self.partial_o_dtype_bytes + num_bytes_per_row_stats = 2 * self.acc_dtype_bytes + return ( + num_ctas_kv + * num_rows_per_cta + * (num_bytes_per_row_o + num_bytes_per_row_stats) + ) + + @property + def cluster_reduction_smem_bytes(self) -> int: + if self.use_multi_ctas_kv != 1 or self.use_cluster_reduction != 1: + return 0 + max_bytes = 0 + for num_ctas_kv in range(2, self.num_ctas_per_seq_kv + 1): + max_bytes = max( + max_bytes, self.cluster_reduction_smem_bytes_for(num_ctas_kv) + ) + return max_bytes + + @property + def tmem_total_cols(self) -> int: + return ( + 2 * self.tmem_s_cols + + 2 * self.tmem_stats_cols + + self.tmem_o_buffer_cols * self.o_stages * self.v_head_dim_stages + ) + + @property + def tmem_alloc_cols(self) -> int: + # Reserve the full SM100 TMEM budget for fixed column placement. + return 512 + + @property + def qk_head_dim_stages(self) -> int: + return ceil(self.head_dim_qk / self.head_dim_per_stage_kv) + + @property + def v_head_dim_stages(self) -> int: + return ceil(self.head_dim_per_cta_v / self.head_dim_per_stage_v) + + @property + def tmem_o_buffer_cols(self) -> int: + return self.tmem_o_cols + + @property + def q_smem_tile_elements(self) -> int: + return self.qk_smem_tile_bytes // self.qkv_dtype_bytes + + @property + def kv_smem_stage_elements(self) -> int: + return self.kv_smem_tile_bytes // self.qkv_dtype_bytes + + def qk_head_stage_width(self, stage_idx: int) -> int: + start = stage_idx * self.head_dim_per_stage_kv + return max(0, min(self.head_dim_per_stage_kv, self.head_dim_qk - start)) + + def v_head_stage_width(self, stage_idx: int) -> int: + start = stage_idx * self.head_dim_per_stage_v + return max(0, min(self.head_dim_per_stage_v, self.head_dim_per_cta_v - start)) + + def is_fp8_qkv(self) -> bool: + """Return whether Q/K/V tensors use E4M3 data.""" + + return self.qkv_dtype == "e4m3" + + def local_kv_tiles(self, total_kv_tiles: int) -> int: + """Return per-CTA KV tiles after multi-CTA KV splitting.""" + if self.use_multi_ctas_kv != 1: + return total_kv_tiles + tiles_per_group = self.num_ctas_per_seq_kv * self.num_insts_kv + num_groups = (total_kv_tiles + tiles_per_group - 1) // tiles_per_group + return max(self.num_insts_kv, num_groups * self.num_insts_kv) + + def loop_domain(self, local_kv_tiles: int) -> int: + """Return the decode-gen steady-state loop domain after HEAD.""" + remaining_kv_tiles = max(local_kv_tiles - self.num_insts_kv, 0) + return (remaining_kv_tiles + self.num_insts_kv - 1) // self.num_insts_kv + + +def compute_workspace_size( + *, + cfg: MlaConfig, + partial_o_dtype, + lse_dtype, +) -> int: + """Return the exact 1CTA split-KV GMEM workspace size in bytes. + + The producer layout is ``[B, SQ, H, split, D]`` for partial O and + ``[B, SQ, H, split]`` for LSE. Cluster reduction keeps these partials in + shared memory and therefore needs no GMEM workspace. + """ + + split_kv = cfg.num_ctas_per_seq_kv + if split_kv == 1 or cfg.use_cluster_reduction == 1: + return 0 + partial_rows = cfg.batch_size * cfg.seq_len_q * cfg.num_heads_q * split_kv + return partial_rows * ( + cfg.head_dim_v * partial_o_dtype.width // 8 + lse_dtype.width // 8 + ) + + +@dataclass(frozen=True) +class MlaProfile: + """Concrete tunable profile for one throughput-latency 1CTA MLA kernel variant. + + Profiles select only the tunable scheduler/decomposition knobs. The config + factory validates them against the concrete shape and expands them into the + full ``MlaConfig`` consumed by the kernel. + """ + + name: str + num_ctas_per_seq_kv: int = 1 + num_ctas_per_head_dim: int = 1 + use_persistent_scheduler: int = 1 + use_clc_dynamic_persistent_scheduler: int = 0 + use_multi_ctas_kv: int = 0 + use_cluster_reduction: int = 0 + kernel_variant: str = "swaps_mma_ab" + tile_size_q: int | None = None + + +@dataclass(frozen=True) +class FlatQueryLaunchShape: + """Logical query geometry and its normalized physical flat-row launch.""" + + logical_num_heads_q: int + logical_seq_len_q: int + num_heads_q: int + seq_len_q: int + tile_size_q: int + + @classmethod + def for_tile( + cls, + logical_num_heads_q: int, + logical_seq_len_q: int, + tile_size_q: int, + ) -> "FlatQueryLaunchShape": + """Build consecutive physical M-row tiles for one logical query.""" + + layout = FlatQueryTileLayout.for_tile( + logical_num_heads_q, + logical_seq_len_q, + tile_size_q, + ) + return cls( + logical_num_heads_q=logical_num_heads_q, + logical_seq_len_q=logical_seq_len_q, + num_heads_q=layout.tile_size_q, + seq_len_q=layout.num_tiles, + tile_size_q=tile_size_q, + ) + + +def validate_tile_size_q(tile_size_q: int) -> int: + """Validate a user-supplied Q tile size for 1CTA MLA.""" + + if tile_size_q not in SUPPORTED_TILE_SIZE_Q: + raise ValueError( + "throughput-latency 1CTA MLA tile_size_q must be one of " + f"{SUPPORTED_TILE_SIZE_Q}, got {tile_size_q}" + ) + return tile_size_q + + +def tile_size_q_from_profile_name(profile: str | None) -> int | None: + """Return the explicit Q tile encoded by profile names such as ``h32_static``.""" + + if not profile: + return None + for tile_size_q in SUPPORTED_TILE_SIZE_Q: + if profile.startswith(f"h{tile_size_q}_"): + return tile_size_q + return None + + +def validate_max_active_clusters(max_active_clusters: int) -> int: + """Return the explicit hardware active-cluster capacity.""" + + if max_active_clusters is None: + raise ValueError("max_active_clusters must be provided") + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be positive") + return max_active_clusters + + +def tile_size_q_for_heads(num_heads_q: int) -> int: + """Choose the automatic 1CTA Q tile from the physical row extent.""" + + return auto_tile_size_q_for_mla_gen( + num_heads_q=num_heads_q, + seq_len_q=1, + ) + + +def auto_tile_size_q_for_mla_gen( + *, + num_heads_q: int, + seq_len_q: int, +) -> int: + """Choose the smallest native M tile covering the flat query rows. + + Shapes wider than the largest 1CTA tile use repeated M64 tiles. The choice + depends only on query geometry, so one compiled topology is reusable across + runtime batch extents and K/V lengths. + """ + + total_q_rows = num_heads_q * seq_len_q + for tile_size_q in SUPPORTED_TILE_SIZE_Q: + if total_q_rows <= tile_size_q: + return tile_size_q + return SUPPORTED_TILE_SIZE_Q[-1] + + +def resolve_auto_flat_query_launch_shape( + *, + num_heads_q: int, + seq_len_q: int, +) -> FlatQueryLaunchShape: + """Resolve the public 1CTA profile tile and flat-row launch extent. + + Equivalent H/Q factorizations resolve to the same physical profile when + their runtime work is otherwise identical. + """ + + tile_size_q = auto_tile_size_q_for_mla_gen( + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + ) + return FlatQueryLaunchShape.for_tile( + num_heads_q, + seq_len_q, + tile_size_q, + ) + + +def profile_name( + base_name: str, num_heads_q: int, tile_size_q: int | None = None +) -> str: + tile_size_q = tile_size_q or tile_size_q_for_heads(num_heads_q) + return f"h{tile_size_q}_{base_name}" + + +def q_tile_work_count( + batch_size: int, + num_heads_q: int, + seq_len_q: int, + tile_size_q: int | None = None, +) -> int: + """Return CTA work count before KV or V-head-dim decomposition.""" + + tile_size_q = tile_size_q or tile_size_q_for_heads(num_heads_q) + return batch_size * ceil(num_heads_q * seq_len_q / tile_size_q) + + +def automatic_split_kv_step_tokens(tile_size_q: int) -> int: + """Return one steady-state K step for the selected task graph.""" + + num_kv_insts = 1 if tile_size_q >= 64 else MlaConfig.num_insts_kv + return MlaConfig.tile_size_kv * num_kv_insts + + +def select_auto_split_kv( + *, + seq_len_kv: int, + tile_size_q: int, + base_work: int, + target_work: int, +) -> int: + """Choose the smallest split count preserving the target K-step depth.""" + + split_kv_step_tokens = automatic_split_kv_step_tokens(tile_size_q) + max_split_kv = max(1, ceil(seq_len_kv / split_kv_step_tokens)) + split_kv = min( + max_split_kv, + max(1, target_work // base_work), + MAX_MLA_SPLITS_KV, + ) + local_k_steps = ceil(max_split_kv / split_kv) + return ceil(max_split_kv / local_k_steps) + + +def is_power_of_two(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +def int_flag(value: bool) -> int: + if value: + return 1 + return 0 + + +def dtype_num_bytes(dtype: str) -> int: + if dtype == "e4m3": + return 1 + return 2 + + +def dtype_config_kwargs(qkv_dtype: str, o_dtype: str) -> dict[str, int]: + """Return dtype byte widths and feature flags for the config dataclass.""" + + kwargs = { + "qkv_dtype_bytes": dtype_num_bytes(qkv_dtype), + "o_dtype_bytes": dtype_num_bytes(o_dtype), + "use_bf16_output": 0, + "use_fp8_output": 0, + } + if o_dtype == "bf16": + kwargs["use_bf16_output"] = 1 + elif o_dtype == "e4m3": + kwargs["use_fp8_output"] = 1 + + return kwargs + + +def tile_size_q_for_profile( + profile: MlaProfile, + num_heads_q: int, + seq_len_q: int, + tile_size_q: int | None = None, +) -> int: + """Resolve the effective Q tile for the selected profile.""" + + del seq_len_q + if tile_size_q is not None: + tile_size_q = validate_tile_size_q(tile_size_q) + if profile.kernel_variant == "keeps_mma_ab" and tile_size_q < 64: + raise ValueError("keeps_mma_ab profiles require tile_size_q >= 64") + if profile.kernel_variant != "keeps_mma_ab" and tile_size_q >= 64: + raise ValueError("swaps_mma_ab profiles require tile_size_q < 64") + return tile_size_q + if profile.tile_size_q is not None: + profile_tile_size_q = validate_tile_size_q(profile.tile_size_q) + if profile.kernel_variant == "keeps_mma_ab" and profile_tile_size_q < 64: + raise ValueError("keeps_mma_ab profiles require tile_size_q >= 64") + if profile.kernel_variant != "keeps_mma_ab" and profile_tile_size_q >= 64: + raise ValueError("swaps_mma_ab profiles require tile_size_q < 64") + return profile_tile_size_q + if profile.kernel_variant == "keeps_mma_ab": + return 64 + return tile_size_q_for_heads(num_heads_q) + + +def kv_stage_count(profile: MlaProfile, tile_size_q: int, qkv_dtype: str) -> int: + if profile.use_clc_dynamic_persistent_scheduler == 1 and tile_size_q == 32: + return 3 + if qkv_dtype == "e4m3": + return 9 + return MlaConfig.kv_stages + + +def softmax_register_budget(tile_size_q: int) -> int: + if tile_size_q == 32: + return 160 + return MlaConfig.softmax_regs + + +def correction_register_budget(tile_size_q: int) -> int: + if tile_size_q == 32: + return 96 + return MlaConfig.correction_regs + + +def keeps_mma_ab_config_kwargs(profile: MlaProfile, qkv_dtype: str) -> dict[str, int]: + """Return keeps-MMA-AB scheduler, pipeline, and register traits.""" + + kv_stages = MlaConfig.kv_stages + q_stages = 1 + if qkv_dtype == "e4m3": + kv_stages = 8 + q_stages = MlaConfig.q_stages + + use_persistent_scheduler = profile.use_persistent_scheduler + use_clc_dynamic_persistent_scheduler = profile.use_clc_dynamic_persistent_scheduler + if profile.use_multi_ctas_kv == 1: + use_persistent_scheduler = 0 + use_clc_dynamic_persistent_scheduler = 0 + + return { + # Keeps-MMA-AB has one QK/PV pipe. Raising this value would change the + # KV work partition without adding the second stream used by the swaps + # schedule, and would therefore skip every other KV tile. + "num_insts_kv": 1, + "kv_stages": kv_stages, + "q_stages": q_stages, + "o_stages": 1, + "threads_per_cta": 384, + "correction_warp_idx": 4, + "mma_warp_idx": 8, + "load_warp_idx": 9, + "page_offsets_warp_idx": 10, + "scheduler_warp_idx": 11, + "softmax_regs": 200, + "correction_regs": 192, + "mma_load_regs": 112, + "tmem_s_cols": 128, + "tmem_o_cols": 128, + "use_persistent_scheduler": use_persistent_scheduler, + "use_clc_dynamic_persistent_scheduler": use_clc_dynamic_persistent_scheduler, + } + + +def resolve_split_kv_reduction_policy( + *, + profile: MlaProfile, + tile_size_q: int, + head_dim: int, + head_dim_per_cta_v: int, + reduction_mode: str | None, +) -> int: + """Return whether the selected split-KV profile should use cluster reduction.""" + + if reduction_mode not in (None, "auto", "cluster", "gmem_separate"): + raise ValueError(f"unsupported reduction_mode: {reduction_mode}") + reduction_mode = reduction_mode or "auto" + + if reduction_mode == "gmem_separate": + return 0 + if reduction_mode == "cluster": + if profile.use_multi_ctas_kv != 1: + raise ValueError("explicit cluster reduction requires a split-KV profile") + if profile.kernel_variant == "keeps_mma_ab": + raise ValueError( + "explicit cluster reduction is not supported by keeps-MMA-AB profiles" + ) + if profile.num_ctas_per_seq_kv > MAX_CLUSTER_SIZE: + raise ValueError( + "explicit cluster reduction exceeds the CUDA cluster-size limit: " + f"split_kv={profile.num_ctas_per_seq_kv}, " + f"max_cluster_size={MAX_CLUSTER_SIZE}" + ) + return 1 + + cluster_capable_profile = ( + profile.use_multi_ctas_kv == 1 + and profile.use_cluster_reduction == 1 + and profile.kernel_variant != "keeps_mma_ab" + and profile.num_ctas_per_seq_kv <= MAX_CLUSTER_SIZE + ) + if not cluster_capable_profile: + return 0 + + modes = select_split_kv_modes( + family="mla_decode", + topology="1cta", + tile_size_q=tile_size_q, + head_dim=head_dim, + head_dim_per_cta_v=head_dim_per_cta_v, + split_kv=profile.num_ctas_per_seq_kv, + available_modes=("cluster", "gmem_separate"), + ) + return int_flag(modes[0] == "cluster") + + +def cluster_reduction_cluster_count(cfg: MlaConfig) -> int: + """Return the number of clusters launched by the cluster reduction grid.""" + + cluster_size = cfg.num_ctas_per_seq_kv + grid_m = cfg.num_ctas_for_all_heads * cfg.num_ctas_per_seq_q * cluster_size + grid_n = cfg.num_ctas_per_head_dim + grid_l = cfg.batch_size + return ceil(grid_m / cluster_size) * grid_n * grid_l + + +def cluster_reduction_cluster_shape(cfg: MlaConfig) -> tuple[int, int, int]: + """Return the cluster reduction cluster shape for the main kernel launch.""" + + return (cfg.num_ctas_per_seq_kv, 1, 1) + + +def resolve_auto_cluster_reduction_mode( + cfg: MlaConfig, + *, + reduction_mode: str | None, + max_active_clusters: int, +) -> str | None: + """Return the runtime reduction mode after the auto-cluster capacity check.""" + + if reduction_mode not in (None, "auto") or cfg.use_cluster_reduction != 1: + return reduction_mode + + if cluster_reduction_cluster_count(cfg) > max_active_clusters: + return "gmem_separate" + + return reduction_mode + + +def resolve_runtime_cluster_reduction_mode( + cfg: MlaConfig | None, + *, + reduction_mode: str | None, + hardware_info, + stream=None, + log=None, +) -> str | None: + """Query cluster occupancy and return the final runtime reduction mode.""" + + if ( + cfg is None + or reduction_mode not in (None, "auto") + or cfg.use_cluster_reduction != 1 + ): + return reduction_mode + + cluster_shape = cluster_reduction_cluster_shape(cfg) + cluster_size = cluster_shape[0] * cluster_shape[1] * cluster_shape[2] + active_clusters = hardware_info.get_max_active_clusters(cluster_size, stream) + cluster_count = cluster_reduction_cluster_count(cfg) + if log is not None: + log( + "cluster_reduction_occupancy", + f"clusters={cluster_count}, active={active_clusters}", + ) + + resolved_mode = resolve_auto_cluster_reduction_mode( + cfg, + reduction_mode=reduction_mode, + max_active_clusters=active_clusters, + ) + if log is not None: + if resolved_mode == "gmem_separate": + log( + "cluster_reduction", + "disabled: grid requires more than one active cluster wave", + ) + else: + log("cluster_reduction", "enabled") + return resolved_mode + + +def resolve_auto_cluster_reduction_config( + cfg: MlaConfig, + *, + reduction_mode: str | None, +) -> MlaConfig: + """Disable auto cluster reduction when the static SMEM footprint is too large.""" + + if reduction_mode not in (None, "auto") or cfg.use_cluster_reduction != 1: + return cfg + + if estimated_throughput_latency_smem_data_bytes(cfg) > MAX_CLUSTER_SMEM_DATA_BYTES: + return replace(cfg, use_cluster_reduction=0) + + return cfg + + +def validate_explicit_cluster_reduction_config( + cfg: MlaConfig, + *, + reduction_mode: str | None, +) -> None: + """Reject explicit cluster requests that exceed the SMEM launch budget.""" + + if reduction_mode != "cluster": + return + smem_data_bytes = estimated_throughput_latency_smem_data_bytes(cfg) + if smem_data_bytes > MAX_CLUSTER_SMEM_DATA_BYTES: + raise ValueError( + "explicit cluster reduction exceeds the SM100A shared-memory budget: " + f"required_data_bytes={smem_data_bytes}, " + f"max_data_bytes={MAX_CLUSTER_SMEM_DATA_BYTES}" + ) + + +def baseline_profile( + tile_size_q: int | None = None, + *, + name: str = "baseline", +) -> MlaProfile: + """Return the default nonpersistent, non-split 1CTA profile.""" + + return MlaProfile( + name=name, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + tile_size_q=tile_size_q, + ) + + +def persistent_profiles( + num_heads_q: int, tile_size_q: int | None = None +) -> tuple[MlaProfile, ...]: + """Return the persistent policy candidates in preferred order.""" + + return ( + MlaProfile( + name=profile_name("clc", num_heads_q, tile_size_q), + use_persistent_scheduler=1, + use_clc_dynamic_persistent_scheduler=1, + use_multi_ctas_kv=0, + tile_size_q=tile_size_q, + ), + MlaProfile( + name=profile_name("static", num_heads_q, tile_size_q), + use_persistent_scheduler=1, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=0, + tile_size_q=tile_size_q, + ), + MlaProfile( + name=profile_name("nonpersistent", num_heads_q, tile_size_q), + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=0, + tile_size_q=tile_size_q, + ), + ) + + +def persistent_override_profile( + profile: MlaProfile, explicit_persistent: bool | None +) -> MlaProfile: + """Apply the user persistent override to a selected non-split profile.""" + + if explicit_persistent is None: + return profile + if profile.use_multi_ctas_kv == 1: + if explicit_persistent: + raise ValueError("persistent scheduling is not supported with split-KV") + return profile + if explicit_persistent: + return replace( + profile, + use_persistent_scheduler=1, + use_clc_dynamic_persistent_scheduler=( + profile.use_clc_dynamic_persistent_scheduler + ), + ) + return replace( + profile, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + ) + + +def forced_persistent_profiles( + *, + num_heads_q: int, + tile_size_q: int | None, + explicit_persistent: bool, +) -> tuple[MlaProfile, ...]: + """Return profiles for an explicit persistent/nonpersistent user request.""" + + if explicit_persistent: + return persistent_profiles(num_heads_q, tile_size_q)[:2] + if tile_size_q is None: + return (baseline_profile(),) + return ( + MlaProfile( + name=profile_name("nonpersistent", num_heads_q, tile_size_q), + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=0, + tile_size_q=tile_size_q, + ), + ) + + +def automatic_unsplit_profiles( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + tile_size_q: int, + max_active_clusters: int, +) -> tuple[MlaProfile, ...]: + """Return automatic scheduler candidates for an unsplit 1CTA launch. + + A direct grid that fits one resident wave has no launch work for a + persistent CTA to replace. Once logical work exceeds resident capacity, + CLC becomes the preferred scheduler independent of dtype, K length, or a + shape-specific threshold. + """ + + clc, static, nonpersistent = persistent_profiles(num_heads_q, tile_size_q) + base_work = q_tile_work_count(batch_size, num_heads_q, seq_len_q, tile_size_q) + if base_work > max_active_clusters: + return (clc, static, nonpersistent) + return (nonpersistent, clc, static) + + +def keeps_mma_ab_profiles( + *, + num_heads_q: int, + batch_size: int, + seq_len_q: int, + max_active_clusters: int, +) -> tuple[MlaProfile, ...]: + """Return keeps-MMA-AB candidates for large head tiles.""" + + # Keeps-MMA-AB is only wired for the q64 head-tiled schedules. H64 and H128 + # both map to the same q64 tile; H128 is split across two head tiles. + if num_heads_q not in (64, 128): + return () + + nonpersistent = MlaProfile( + name="h64_keeps_mma_ab", + num_ctas_per_head_dim=1, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=0, + kernel_variant="keeps_mma_ab", + tile_size_q=64, + ) + split = MlaProfile( + name="h64_keeps_mma_ab_splitkv_gmem", + num_ctas_per_seq_kv=4, + num_ctas_per_head_dim=1, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=1, + kernel_variant="keeps_mma_ab", + tile_size_q=64, + ) + clc = MlaProfile( + name="h64_keeps_mma_ab_clc", + num_ctas_per_head_dim=1, + use_persistent_scheduler=1, + use_clc_dynamic_persistent_scheduler=1, + use_multi_ctas_kv=0, + kernel_variant="keeps_mma_ab", + tile_size_q=64, + ) + base_work = q_tile_work_count(batch_size, num_heads_q, seq_len_q, 64) + if base_work > max_active_clusters: + # The Keeps-MMA-AB task graph supports one persistent work tile. A + # persistent launch would therefore omit the grid suffix once work + # exceeds the resident wave, so multi-wave shapes use the direct grid. + return nonpersistent, split + return nonpersistent, clc, split + + +def keeps_mma_ab_explicit_split_profile( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + split_kv: int, + latent_dim: int, + max_active_clusters: int, +) -> MlaProfile | None: + """Return the keeps-MMA-AB profile constrained by an explicit split-KV count.""" + + # Explicit split-KV is accepted only for the q64 keeps-MMA-AB schedules; the + # swaps-MMA-AB path handles smaller tiles through ``explicit_split_profile``. + if num_heads_q not in (64, 128): + return None + if split_kv <= 1: + return MlaProfile( + name="h64_keeps_mma_ab", + num_ctas_per_head_dim=1, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=0, + kernel_variant="keeps_mma_ab", + tile_size_q=64, + ) + + base_work = q_tile_work_count(batch_size, num_heads_q, seq_len_q, 64) + head_dim_split = head_dim_split_for_work( + work_after_kv=base_work * split_kv, + target_work=validate_max_active_clusters(max_active_clusters), + latent_dim=latent_dim, + ) + + profile_suffix = "splitkv_gmem" + if split_kv != 4: + profile_suffix = f"splitkv{split_kv}_gmem" + if head_dim_split > 1: + profile_suffix = f"{profile_suffix}_hdim{latent_dim // head_dim_split}" + return MlaProfile( + name=f"h64_keeps_mma_ab_{profile_suffix}", + num_ctas_per_seq_kv=split_kv, + num_ctas_per_head_dim=head_dim_split, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=1, + kernel_variant="keeps_mma_ab", + tile_size_q=64, + ) + + +def keeps_mma_ab_forced_persistent_profile( + *, num_heads_q: int, explicit_persistent: bool +) -> MlaProfile | None: + """Return a keeps-MMA-AB profile constrained by the user persistent request.""" + + if num_heads_q not in (64, 128): + return None + return MlaProfile( + name="h64_keeps_mma_ab_clc" if explicit_persistent else "h64_keeps_mma_ab", + num_ctas_per_head_dim=1, + use_persistent_scheduler=int_flag(explicit_persistent), + use_clc_dynamic_persistent_scheduler=int_flag(explicit_persistent), + use_multi_ctas_kv=0, + kernel_variant="keeps_mma_ab", + tile_size_q=64, + ) + + +def head_dim_split_for_work( + *, work_after_kv: int, target_work: int, latent_dim: int +) -> int: + """Choose the V head-dim split used after the selected split-KV count.""" + + # Use a V head-dim split only when the current grid leaves at least a 2x SM + # gap. Smaller gaps are usually not worth the extra scheduling and epilogue + # work. + if work_after_kv * 2 > target_work: + return 1 + + if work_after_kv * 4 <= target_work and latent_dim % 4 == 0: + return 4 + if latent_dim % 2 == 0: + return 2 + return 1 + + +def explicit_split_profile( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + split_kv: int, + latent_dim: int, + max_active_clusters: int, + tile_size_q: int | None = None, +) -> MlaProfile: + """Return a swaps-MMA-AB profile constrained by a user split-KV count.""" + + max_active_clusters = validate_max_active_clusters(max_active_clusters) + if split_kv <= 1: + return baseline_profile(tile_size_q) + + base_work = q_tile_work_count(batch_size, num_heads_q, seq_len_q, tile_size_q) + target_work = max_active_clusters + head_dim_split = head_dim_split_for_work( + work_after_kv=base_work * split_kv, + target_work=target_work, + latent_dim=latent_dim, + ) + + suffix = f"splitkv{split_kv}" + if head_dim_split > 1: + suffix = f"{suffix}_hdim{latent_dim // head_dim_split}" + return MlaProfile( + name=profile_name(suffix, num_heads_q, tile_size_q), + num_ctas_per_seq_kv=split_kv, + num_ctas_per_head_dim=head_dim_split, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=1, + use_cluster_reduction=1, + tile_size_q=tile_size_q, + ) + + +def auto_split_profile( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + seq_len_kv: int, + latent_dim: int, + max_active_clusters: int, + tile_size_q: int | None = None, +) -> MlaProfile | None: + """Choose split-KV followed by V head-dimension decomposition.""" + + max_active_clusters = validate_max_active_clusters(max_active_clusters) + original_tile_size_q = tile_size_q or tile_size_q_for_heads(num_heads_q) + base_work = q_tile_work_count(batch_size, num_heads_q, seq_len_q, tile_size_q) + target_work = max_active_clusters + if base_work <= 0 or base_work >= target_work: + return None + + split_kv = select_auto_split_kv( + seq_len_kv=seq_len_kv, + tile_size_q=original_tile_size_q, + base_work=base_work, + target_work=target_work, + ) + work_after_kv = base_work * split_kv + + head_dim_split = head_dim_split_for_work( + work_after_kv=work_after_kv, + target_work=target_work, + latent_dim=latent_dim, + ) + if split_kv == 1 and head_dim_split == 1: + return None + + suffix = "baseline" + if split_kv > 1: + suffix = "splitkv" + elif head_dim_split > 1: + suffix = f"hdim{latent_dim // head_dim_split}" + if split_kv > 1 and head_dim_split > 1: + suffix = f"{suffix}_hdim{latent_dim // head_dim_split}" + return MlaProfile( + name=profile_name(suffix, num_heads_q, original_tile_size_q), + num_ctas_per_seq_kv=split_kv, + num_ctas_per_head_dim=head_dim_split, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=int_flag(split_kv > 1), + use_cluster_reduction=int_flag(split_kv > 1), + kernel_variant=( + "keeps_mma_ab" if original_tile_size_q >= 64 else "swaps_mma_ab" + ), + tile_size_q=original_tile_size_q, + ) + + +def is_throughput_latency_mla_supported_shape( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + seq_len_kv: int, + latent_dim: int = 512, + rope_dim: int = 64, + num_tokens_per_page: int = 32, +) -> bool: + """Return whether the basic dense DS MLA 1CTA path can be considered.""" + + del batch_size + return ( + latent_dim == 512 + and rope_dim == 64 + and 1 <= num_heads_q <= 128 + and is_power_of_two(num_heads_q) + and num_tokens_per_page in SUPPORTED_MLA_PAGE_SIZES + and seq_len_q >= 1 + and seq_len_kv >= 128 + ) + + +def enumerate_throughput_latency_mla_profiles( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + seq_len_kv: int, + latent_dim: int = 512, + rope_dim: int = 64, + num_tokens_per_page: int = 32, + max_active_clusters: int, + qkv_dtype: str = "bf16", + tile_size_q: int | None = None, + explicit_split_kv: int | None = None, + explicit_persistent: bool | None = None, +) -> tuple[MlaProfile, ...]: + """Return benchmarkable throughput-latency 1CTA profiles for a problem shape.""" + + max_active_clusters = validate_max_active_clusters(max_active_clusters) + if tile_size_q is not None: + tile_size_q = validate_tile_size_q(tile_size_q) + if not is_throughput_latency_mla_supported_shape( + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + latent_dim=latent_dim, + rope_dim=rope_dim, + num_tokens_per_page=num_tokens_per_page, + ): + return () + + selected_tile_size_q = tile_size_q or auto_tile_size_q_for_mla_gen( + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + ) + if explicit_split_kv is not None and explicit_split_kv > 0: + if explicit_split_kv > 1 and explicit_persistent is True: + raise ValueError("persistent scheduling is not supported with split-KV") + if explicit_split_kv == 1 and explicit_persistent is not None: + if selected_tile_size_q >= 64: + profile = keeps_mma_ab_forced_persistent_profile( + num_heads_q=num_heads_q, + explicit_persistent=explicit_persistent, + ) + if profile is None: + return () + return (profile,) + return forced_persistent_profiles( + num_heads_q=num_heads_q, + tile_size_q=selected_tile_size_q, + explicit_persistent=explicit_persistent, + ) + if selected_tile_size_q >= 64: + profile = keeps_mma_ab_explicit_split_profile( + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + split_kv=explicit_split_kv, + latent_dim=latent_dim, + max_active_clusters=max_active_clusters, + ) + if profile is None: + return () + return (profile,) + + return ( + explicit_split_profile( + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + split_kv=explicit_split_kv, + latent_dim=latent_dim, + max_active_clusters=max_active_clusters, + tile_size_q=selected_tile_size_q, + ), + ) + + if explicit_persistent is not None: + if selected_tile_size_q >= 64: + profile = keeps_mma_ab_forced_persistent_profile( + num_heads_q=num_heads_q, + explicit_persistent=explicit_persistent, + ) + if profile is None: + return () + return (profile,) + return forced_persistent_profiles( + num_heads_q=num_heads_q, + tile_size_q=selected_tile_size_q, + explicit_persistent=explicit_persistent, + ) + + profiles: list[MlaProfile] = [] + split_profile = auto_split_profile( + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + latent_dim=latent_dim, + max_active_clusters=max_active_clusters, + tile_size_q=selected_tile_size_q, + ) + if split_profile is not None: + profiles.append(split_profile) + + if selected_tile_size_q >= 64: + profiles.extend( + profile + for profile in keeps_mma_ab_profiles( + num_heads_q=num_heads_q, + batch_size=batch_size, + seq_len_q=seq_len_q, + max_active_clusters=max_active_clusters, + ) + if profile.name not in {candidate.name for candidate in profiles} + ) + else: + profiles.extend( + automatic_unsplit_profiles( + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + tile_size_q=selected_tile_size_q, + max_active_clusters=max_active_clusters, + ) + ) + + if selected_tile_size_q < 64: + baseline = baseline_profile(selected_tile_size_q) + if not any(profile.name == baseline.name for profile in profiles): + profiles.append(baseline) + + return tuple(profiles) + + +def resolve_throughput_latency_mla_profile( + *, + profile: MlaProfile | str | None, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + seq_len_kv: int, + latent_dim: int = 512, + rope_dim: int = 64, + num_tokens_per_page: int = 32, + max_active_clusters: int, + qkv_dtype: str = "bf16", + tile_size_q: int | None = None, + explicit_split_kv: int | None = None, + explicit_persistent: bool | None = None, +) -> MlaProfile: + """Resolve an explicit profile name or default to the first candidate.""" + + if isinstance(profile, MlaProfile): + return profile + + profiles = enumerate_throughput_latency_mla_profiles( + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + latent_dim=latent_dim, + rope_dim=rope_dim, + num_tokens_per_page=num_tokens_per_page, + max_active_clusters=max_active_clusters, + qkv_dtype=qkv_dtype, + tile_size_q=tile_size_q, + explicit_split_kv=explicit_split_kv, + explicit_persistent=explicit_persistent, + ) + if profile in (None, "", "default"): + if profiles: + return profiles[0] + return baseline_profile() + + for candidate in profiles: + if candidate.name == profile: + return candidate + available = ", ".join(candidate.name for candidate in profiles) or "none" + raise ValueError( + f"throughput-latency 1CTA MLA profile {profile!r} is not valid for this shape; " + f"available profiles: {available}" + ) + + +def make_throughput_latency_mla_config( + *, + batch_size: int, + num_heads_q: int, + seq_len_q: int, + seq_len_kv: int, + latent_dim: int = 512, + rope_dim: int = 64, + num_tokens_per_page: int = 32, + qkv_dtype: str = "bf16", + o_dtype: str = "bf16", + profile: MlaProfile | str | None = None, + persistent_wave_sm_count: int | None = None, + max_active_clusters: int, + reduction_mode: str | None = None, + logical_num_heads_q: int | None = None, + logical_seq_len_q: int | None = None, + tile_size_q: int | None = None, + explicit_split_kv: int | None = None, + explicit_persistent: bool | None = None, + mask_type: MaskType | str = MaskType.CAUSAL, +) -> MlaConfig: + """Return throughput-latency 1CTA MLA traits for a concrete profile.""" + + mask_type = normalize_mask_type(mask_type) + + if logical_num_heads_q is None: + logical_num_heads_q = num_heads_q + if logical_seq_len_q is None: + logical_seq_len_q = seq_len_q + + if batch_size <= 0: + raise ValueError("batch_size must be positive") + if num_heads_q <= 0: + raise ValueError("num_heads_q must be positive") + if logical_num_heads_q <= 0: + raise ValueError("logical_num_heads_q must be positive") + if num_heads_q > 128 or not is_power_of_two(num_heads_q): + raise ValueError( + "num_heads_q must be a power of two no larger than 128 for " + f"throughput-latency 1CTA MLA: num_heads_q={num_heads_q}" + ) + if seq_len_q <= 0: + raise ValueError("seq_len_q must be positive") + if logical_seq_len_q <= 0: + raise ValueError("logical_seq_len_q must be positive") + if seq_len_kv <= 0: + raise ValueError("seq_len_kv must be positive") + if latent_dim <= 0: + raise ValueError("latent_dim must be positive") + if rope_dim < 0: + raise ValueError("rope_dim must be non-negative") + if num_tokens_per_page not in SUPPORTED_MLA_PAGE_SIZES: + raise ValueError( + "num_tokens_per_page must be one of " + f"{SUPPORTED_MLA_PAGE_SIZES}, got {num_tokens_per_page}" + ) + if MlaConfig.tile_size_kv % num_tokens_per_page != 0: + raise ValueError( + "num_tokens_per_page must exactly divide the 1CTA KV tile: " + f"tile_size_kv={MlaConfig.tile_size_kv}, " + f"num_tokens_per_page={num_tokens_per_page}" + ) + pages_per_kv_tile = MlaConfig.tile_size_kv // num_tokens_per_page + if pages_per_kv_tile > MlaConfig.page_offsets_entries_per_stage: + raise ValueError( + "page-offset staging capacity is smaller than one KV tile: " + f"pages_per_kv_tile={pages_per_kv_tile}, " + "page_offsets_entries_per_stage=" + f"{MlaConfig.page_offsets_entries_per_stage}" + ) + if qkv_dtype not in ("bf16", "e4m3"): + raise ValueError(f"unsupported qkv_dtype={qkv_dtype!r}") + if o_dtype not in ("bf16", "e4m3"): + raise ValueError(f"unsupported o_dtype={o_dtype!r}") + if tile_size_q is not None: + tile_size_q = validate_tile_size_q(tile_size_q) + elif isinstance(profile, str): + tile_size_q = tile_size_q_from_profile_name(profile) + max_active_clusters = validate_max_active_clusters(max_active_clusters) + + head_dim_qk = latent_dim + rope_dim + selected_profile = resolve_throughput_latency_mla_profile( + profile=profile, + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + latent_dim=latent_dim, + rope_dim=rope_dim, + num_tokens_per_page=num_tokens_per_page, + max_active_clusters=max_active_clusters, + qkv_dtype=qkv_dtype, + tile_size_q=tile_size_q, + explicit_split_kv=explicit_split_kv, + explicit_persistent=explicit_persistent, + ) + selected_profile = persistent_override_profile( + selected_profile, + explicit_persistent, + ) + tile_size_q = tile_size_q_for_profile( + selected_profile, num_heads_q, seq_len_q, tile_size_q + ) + flat_layout = FlatQueryTileLayout.for_tile( + logical_num_heads_q, + logical_seq_len_q, + tile_size_q, + ) + if num_heads_q != flat_layout.tile_size_q or seq_len_q != flat_layout.num_tiles: + raise ValueError( + "physical 1CTA launch shape must match the flat query-row layout: " + f"got num_heads_q={num_heads_q}, seq_len_q={seq_len_q}; expected " + f"num_heads_q={flat_layout.tile_size_q}, " + f"seq_len_q={flat_layout.num_tiles} for logical " + f"H={logical_num_heads_q}, SQ={logical_seq_len_q}" + ) + if selected_profile.kernel_variant == "keeps_mma_ab" and num_heads_q < tile_size_q: + raise ValueError( + "keeps_mma_ab profiles require an effective Q tile of at least " + f"{tile_size_q}, got num_heads_q={num_heads_q}" + ) + num_ctas_for_all_heads = ceil(num_heads_q / tile_size_q) + num_ctas_per_head_dim = selected_profile.num_ctas_per_head_dim + + if num_ctas_per_head_dim <= 0: + raise ValueError("num_ctas_per_head_dim must be positive") + if selected_profile.num_ctas_per_seq_kv <= 0: + raise ValueError("num_ctas_per_seq_kv must be positive") + if selected_profile.num_ctas_per_seq_kv > seq_len_kv: + raise ValueError( + "num_ctas_per_seq_kv must not exceed seq_len_kv: " + f"num_ctas_per_seq_kv={selected_profile.num_ctas_per_seq_kv}, " + f"seq_len_kv={seq_len_kv}" + ) + if latent_dim % num_ctas_per_head_dim != 0: + raise ValueError( + "latent_dim must be divisible by num_ctas_per_head_dim: " + f"latent_dim={latent_dim}, " + f"num_ctas_per_head_dim={num_ctas_per_head_dim}" + ) + head_dim_per_cta_v = latent_dim // num_ctas_per_head_dim + use_cluster_reduction = resolve_split_kv_reduction_policy( + profile=selected_profile, + tile_size_q=tile_size_q, + head_dim=latent_dim, + head_dim_per_cta_v=head_dim_per_cta_v, + reduction_mode=reduction_mode, + ) + if use_cluster_reduction == 1 and flat_layout.tail_rows != tile_size_q: + if reduction_mode == "cluster": + raise ValueError( + "cluster reduction requires every launched Q tile to contain a full " + "tile_size_q rows: " + f"tail_rows={flat_layout.tail_rows}, tile_size_q={tile_size_q}" + ) + use_cluster_reduction = 0 + num_ctas_per_seq_q = max(1, seq_len_q) + + config_kwargs: dict[str, Any] = dict( + batch_size=batch_size, + num_heads_q=num_heads_q, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + logical_num_heads_q=logical_num_heads_q, + logical_seq_len_q=logical_seq_len_q, + mask_type=mask_type, + head_dim_qk=head_dim_qk, + head_dim_v=latent_dim, + latent_dim=latent_dim, + rope_dim=rope_dim, + qkv_dtype=qkv_dtype, + o_dtype=o_dtype, + tile_size_q=tile_size_q, + num_ctas_per_seq_q=num_ctas_per_seq_q, + num_ctas_for_all_heads=num_ctas_for_all_heads, + num_ctas_per_seq_kv=selected_profile.num_ctas_per_seq_kv, + num_ctas_per_head_dim=num_ctas_per_head_dim, + head_dim_per_cta_v=head_dim_per_cta_v, + head_dim_per_stage_v=128, + kv_stages=kv_stage_count(selected_profile, tile_size_q, qkv_dtype), + softmax0_warp_idx=0, + softmax_regs=softmax_register_budget(tile_size_q), + correction_regs=correction_register_budget(tile_size_q), + num_tokens_per_page=num_tokens_per_page, + max_num_pages_per_seq_kv=max(1, ceil(seq_len_kv / num_tokens_per_page)), + tmem_s_cols=tile_size_q, + tmem_stats_cols=MlaConfig.tmem_stats_cols, + tmem_o_cols=tile_size_q, + kernel_variant=selected_profile.kernel_variant, + use_multi_ctas_kv=selected_profile.use_multi_ctas_kv, + use_cluster_reduction=use_cluster_reduction, + use_persistent_scheduler=selected_profile.use_persistent_scheduler, + use_clc_dynamic_persistent_scheduler=( + selected_profile.use_clc_dynamic_persistent_scheduler + ), + persistent_wave_sm_count=persistent_wave_sm_count, + ) + config_kwargs.update(dtype_config_kwargs(qkv_dtype, o_dtype)) + if selected_profile.kernel_variant == "keeps_mma_ab": + config_kwargs.update(keeps_mma_ab_config_kwargs(selected_profile, qkv_dtype)) + cfg = MlaConfig(**config_kwargs) + validate_explicit_cluster_reduction_config(cfg, reduction_mode=reduction_mode) + cfg = resolve_auto_cluster_reduction_config( + cfg, + reduction_mode=reduction_mode, + ) + return cfg diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py new file mode 100644 index 000000000000..cab0923d5959 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py @@ -0,0 +1,2283 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Throughput-latency 1CTA MLA TS schedule and kernel wrapper. + +The throughput-latency 1CTA policy builds one captured task graph per concrete +``MlaConfig``. It supports the non-persistent baseline, +static-persistent, CLC-persistent, and GMEM split-KV reduction profiles selected +by ``kernel_policy.py``. Python-side launch validation catches unsupported +profile/split/workspace combinations before the JIT body constructs tensor +layouts that depend on those constants. +""" + +from dataclasses import dataclass, replace as dataclass_replace + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +from cutlass.cute.testing import assert_ as runtime_assert +from cutlass import Float32, Int32, Int64 +from ...tensor_map import ( + create_tensor_map_ragged_from_tensor, + create_tensor_map_tiled_from_view, +) +from cutlass.experimental import cuda +from cutlass.experimental import primitives as prims +from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo + +from cutlass.experimental.task_scheduling.enums import PipelineType, SignalingThreads +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + SmemAllocator, + TmemAllocator, +) +from cutlass.experimental.task_scheduling.resources import ( + PipelineConfig, + TileSchedulerConfig, + WorkQueue, +) +from cutlass.experimental.task_scheduling.task_manager import TaskManager + +from .config import ( + MlaConfig, + make_throughput_latency_mla_config, +) +from .resources import ( + SmemKvResource, + SmemPageOffsetsResource, + SmemQResource, + TmemCorrResource, + TmemOResource, + TmemPResource, + SmemPResource, + TmemSKeepsResource, + TmemSResource, + TmemSoftmaxGlobalResource, + TmemSoftmaxLocalResource, + ScheduleTokenThrottleResource, +) +from .tasks import ( + MlaDecodeTask, + create_throughput_latency_correction_task, + create_keeps_mma_ab_correction_task, + create_keeps_mma_ab_mma_task, + create_keeps_mma_ab_softmax_task, + create_load_page_offsets_task, + create_padding_task, + create_throughput_latency_scheduler_task, + create_throughput_latency_load_task, + create_throughput_latency_mma_task, + create_throughput_latency_softmax0_task, + create_throughput_latency_softmax1_task, +) +from .parallel_reduction import ( + PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_SLICE, + PARALLEL_GMEM_REDUCTION_SWAPS_ELEMENTS_PER_SLICE, + parallel_gmem_reduction_base_clusters, + parallel_gmem_reduction_launch_shape, + parallel_gmem_reduction_threads, + run_parallel_gmem_reduction_kernel, + supports_parallel_gmem_reduction, +) +from .reduction import gmem_reduction_launch_shape, run_gmem_reduction_kernel +from ..helpers.constants import TMEM_LIFECYCLE_BARRIER_ID +from ..helpers.mask import MaskType, normalize_mask_type +from ..helpers.tile import ( + runtime_query_tile_is_active, + runtime_split_pruning_is_profitable, + runtime_split_tile_is_active, + runtime_work_tile_is_active, +) +from ..parallel_reduction_topology import ( + choose_q64_parallel_reducer_cluster_size, + make_balanced_parallel_reduction_topology, + validate_parallel_reduction_workspace, +) + + +# Softmax uses exp2, so natural-scale scores are multiplied by log2(e). +LOG2_E = 1.4426950408889634 + + +@cute.jit +def _publish_neutral_standalone_partial( + cfg, + acc_o, + acc_lse, + batch_idx, + cta_idx_q, + cta_idx_kv, + cta_idx_head_dim_v, + head_base_idx, +): + """Publish ``(O=0, LSE=-inf)`` for one pruned producer tile. + + Standalone reducers use the configured split geometry. A runtime-pruned + producer must therefore initialize its workspace slot before releasing the + PDL-dependent reducer; otherwise a fixed-S reducer could consume stale + partials left by an earlier graph replay. Each thread clears aligned BF16 + vec8 fragments from this CTA's V slice. The first V-slice CTA also + publishes one neutral LSE per covered Q/head row. + """ + + if cutlass.const_expr(cfg.partial_o_dtype_bytes != cutlass.BFloat16.width // 8): + raise ValueError("neutral split-KV partials require BF16 workspace storage") + + thread_idx, _, _ = cute.arch.thread_idx() + vectors_per_row = cfg.head_dim_per_cta_v // 8 + vectors_per_tile = cfg.tile_size_q * vectors_per_row + vectors_per_thread = ( + vectors_per_tile + cfg.threads_per_cta - 1 + ) // cfg.threads_per_cta + zero_vec = cutlass.Array( + Int32, + 4, + space=cutlass.AddressSpace.rmem, + ) + for elem_idx in cutlass.range_constexpr(4): + zero_vec[elem_idx] = Int32(0) + + for iter_idx in cutlass.range_constexpr(vectors_per_thread): + vector_idx = thread_idx + Int32(iter_idx * cfg.threads_per_cta) + if vector_idx < Int32(vectors_per_tile): + local_row_idx = vector_idx // Int32(vectors_per_row) + row_vector_idx = vector_idx - local_row_idx * Int32(vectors_per_row) + head_idx = head_base_idx + local_row_idx + if head_idx < Int32(cfg.num_heads_q): + dim_idx = cta_idx_head_dim_v * Int32( + cfg.head_dim_per_cta_v + ) + row_vector_idx * Int32(8) + if dim_idx < Int32(cfg.head_dim_v): + elem_offset = ( + Int64(batch_idx) + * Int64( + cfg.seq_len_q + * cfg.num_heads_q + * cfg.num_ctas_per_seq_kv + * cfg.head_dim_v + ) + + Int64(cta_idx_q) + * Int64( + cfg.num_heads_q * cfg.num_ctas_per_seq_kv * cfg.head_dim_v + ) + + Int64(head_idx) + * Int64(cfg.num_ctas_per_seq_kv * cfg.head_dim_v) + + Int64(cta_idx_kv) * Int64(cfg.head_dim_v) + + Int64(dim_idx) + ) + dst_ptr = cutlass.inttoptr( + acc_o.iterator.raw_ptr().toint(Int64) + + elem_offset * Int64(cfg.partial_o_dtype_bytes), + mem_space=1, + dtype=Int32, + ) + dst_ptr.store( + zero_vec.data_ptr().load(count=4, alignment=16), + alignment=16, + ) + + if cta_idx_head_dim_v == Int32(0) and thread_idx < Int32(cfg.tile_size_q): + head_idx = head_base_idx + thread_idx + if head_idx < Int32(cfg.num_heads_q): + lse_offset = ( + batch_idx + * Int32(cfg.seq_len_q * cfg.num_heads_q * cfg.num_ctas_per_seq_kv) + + cta_idx_q * Int32(cfg.num_heads_q * cfg.num_ctas_per_seq_kv) + + head_idx * Int32(cfg.num_ctas_per_seq_kv) + + cta_idx_kv + ) + (acc_lse.iterator.raw_ptr() + lse_offset).store(Float32(-Float32.inf)) + + +@cute.jit +def _persistent_work_tile_is_inactive(cfg, cache_seqs, cu_seqlens_q, work_tile): + """Return whether a persistent physical Q tile has no runtime rows.""" + + cta_idx_q, _, batch_head_idx = work_tile.tile_idx + batch_idx = Int32(batch_head_idx) // Int32(cfg.num_ctas_for_all_heads) + return not runtime_work_tile_is_active( + cfg, + cache_seqs, + cu_seqlens_q, + batch_idx, + cta_idx_q, + Int32(0), + ) + + +@dataclass(kw_only=True) +class ThroughputLatencyMlaStaticWorkQueue(WorkQueue): + """Static persistent scheduler for MLA tiles shaped as (cta_q, cta_head_dim, batch_head_tile).""" + + cfg: cutlass.Constexpr[MlaConfig] = None + batch_size: object = None + cache_seqs: object = None + cu_seqlens_q: object = None + enable_runtime_skip: cutlass.Constexpr[bool] = False + + def __init__( + self, + tile_scheduler_config: TileSchedulerConfig, + cfg: cutlass.Constexpr[MlaConfig] = None, + batch_size=None, + cache_seqs=None, + cu_seqlens_q=None, + **kwargs, + ) -> None: + WorkQueue.__init__( + self, + tile_scheduler_config=tile_scheduler_config, + **kwargs, + ) + self.cfg = cfg + self.batch_size = batch_size + self.cache_seqs = cache_seqs + self.cu_seqlens_q = cu_seqlens_q + self.enable_runtime_skip = cu_seqlens_q is not None + + @cute.jit + def skip_work_tile_if(self, work_tile: WorkTileInfo): + """Skip inactive physical Q tiles while retaining queue bookkeeping.""" + + return _persistent_work_tile_is_inactive( + self.cfg, + self.cache_seqs, + self.cu_seqlens_q, + work_tile, + ) + + @cute.jit + def _work_tile_from_linear(self, linear_idx: Int32) -> WorkTileInfo: + """Decode a linear persistent index into a throughput-latency work tile.""" + ctas_q = Int32(self.cfg.num_ctas_per_seq_q) + ctas_head_dim = Int32(self.cfg.num_ctas_per_head_dim) + tiles_per_batch_head = ctas_q * ctas_head_dim + batch_head_idx = linear_idx // tiles_per_batch_head + in_batch_head_idx = linear_idx - batch_head_idx * tiles_per_batch_head + cta_idx_head_dim = in_batch_head_idx // ctas_q + cta_idx_q = in_batch_head_idx - cta_idx_head_dim * ctas_q + total_tiles = ( + tiles_per_batch_head + * Int32(self.batch_size) + * Int32(self.cfg.num_ctas_for_all_heads) + ) + return WorkTileInfo( + (cta_idx_q, cta_idx_head_dim, batch_head_idx), + linear_idx < total_tiles, + ) + + @cute.jit + def _make_initial_work_tile(self) -> WorkTileInfo: + """Return the initial work tile for the current CTA.""" + return self._work_tile_from_linear(Int32(cute.arch.block_idx()[2])) + + @cute.jit + def initial_work_tile_info(self) -> WorkTileInfo: + """Return the initial TS work-tile info.""" + return self._make_initial_work_tile() + + @cute.jit + def _get_and_advance_work_tile_impl( + self, + stage_info, + ) -> WorkTileInfo: + """Advance the persistent work tile by one grid-stride step.""" + cta_idx_q, cta_idx_head_dim, batch_head_idx = stage_info.work_tile.tile_idx + linear_idx = Int32(cta_idx_q) + Int32(self.cfg.num_ctas_per_seq_q) * ( + Int32(cta_idx_head_dim) + + Int32(self.cfg.num_ctas_per_head_dim) * Int32(batch_head_idx) + ) + next_linear_idx = linear_idx + Int32(cute.arch.grid_dim()[2]) + return self._work_tile_from_linear(next_linear_idx) + + +@dataclass(kw_only=True) +class ThroughputLatencyMlaClcWorkQueue(WorkQueue): + """CLC work queue shim for throughput-latency 1CTA captured schedules.""" + + cfg: cutlass.Constexpr[MlaConfig] = None + cache_seqs: object = None + cu_seqlens_q: object = None + enable_runtime_skip: cutlass.Constexpr[bool] = False + + def __init__( + self, + tile_scheduler_config: TileSchedulerConfig, + cfg: cutlass.Constexpr[MlaConfig] = None, + cache_seqs=None, + cu_seqlens_q=None, + **kwargs, + ) -> None: + WorkQueue.__init__( + self, + tile_scheduler_config=tile_scheduler_config, + **kwargs, + ) + self.cfg = cfg + self.cache_seqs = cache_seqs + self.cu_seqlens_q = cu_seqlens_q + self.enable_runtime_skip = cu_seqlens_q is not None + + @cute.jit + def skip_work_tile_if(self, work_tile: WorkTileInfo): + """Skip inactive physical Q tiles while retaining queue bookkeeping.""" + + return _persistent_work_tile_is_inactive( + self.cfg, + self.cache_seqs, + self.cu_seqlens_q, + work_tile, + ) + + +def _default_scales(scale_softmax_log2, output_scale): + """Return default softmax and output scales for validation-only builds.""" + if scale_softmax_log2 is None: + scale_softmax_log2 = Float32(1.0) + if output_scale is None: + output_scale = Float32(1.0) + return scale_softmax_log2, output_scale + + +def _check_persistent_scheduler_modes( + use_clc_dynamic_scheduler: bool, + use_static_persistent_scheduler: bool, +) -> None: + """Reject mutually exclusive persistent scheduler selections.""" + if use_clc_dynamic_scheduler and use_static_persistent_scheduler: + raise ValueError( + "throughput-latency 1CTA MLA cannot enable both CLC dynamic and static " + "persistent schedulers" + ) + + +def _make_static_work_queue( + cfg, + tile_sched_params, + cache_seqs, + cu_seqlens_q, + name: str, +): + """Create the static persistent work queue shared by both 1CTA variants.""" + return ThroughputLatencyMlaStaticWorkQueue( + tile_scheduler_config=TileSchedulerConfig.create_static_persistent_tile_scheduler_params( + tile_scheduler_params=tile_sched_params, + ), + cfg=cfg, + batch_size=cute.size(cache_seqs), + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + name=name, + ) + + +def _make_clc_work_queue_and_throttle( + cfg, + tile_sched_params, + clc_response_ptr, + cache_seqs, + cu_seqlens_q, + cta_layout_vmnk, + load_group, + scheduler_group, + queue_name: str, + throttle_name: str, +): + """Create the CLC work queue and the load-to-scheduler throttle edge.""" + agent = pipeline.Agent + schedule_token_throttle_pipeline_config = ( + PipelineConfig.create_async_async_pipeline_cfg( + num_stages=2, + producer_group=load_group, + consumer_group=scheduler_group, + cta_layout_vmnk=cta_layout_vmnk, + ) + ) + schedule_token_pipeline_config = PipelineConfig.create_clc_fetch_async_pipeline_cfg( + num_stages=2, + num_bytes=16, + producer_group=pipeline.CooperativeGroup(agent.Thread), + consumer_group=pipeline.CooperativeGroup( + agent.Thread, + cfg.threads_per_cta, + ), + cta_layout_vmnk=cta_layout_vmnk, + ) + return ( + ThroughputLatencyMlaClcWorkQueue( + tile_scheduler_config=TileSchedulerConfig.create_clc_dynamic_persistent_tile_scheduler_params( + tile_scheduler_params=tile_sched_params, + response_ptr=clc_response_ptr, + ), + cfg=cfg, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + pipeline_config=schedule_token_pipeline_config, + name=queue_name, + ), + ScheduleTokenThrottleResource( + pipeline_config=schedule_token_throttle_pipeline_config, + name=throttle_name, + ), + ) + + +def build_throughput_latency_mla_task_manager( + cfg: MlaConfig, + *, + total_kv_tiles: int = 4, + use_page_offsets: bool = False, + tma_desc_q_latent=None, + tma_desc_q_rope=None, + tma_desc_c_latent=None, + tma_desc_c_rope=None, + tma_desc_v=None, + c_rope_tensor=None, + page_offsets=None, + cache_seqs=None, + cu_seqlens_q=None, + head_idx=None, + batch_idx=None, + cta_idx_q=None, + cta_idx_kv=None, + cta_idx_head_dim_v=None, + scale_softmax_log2=None, + output_scale=None, + o_tensor=None, + lse_tensor=None, + acc_o_tensor=None, + acc_lse_tensor=None, + tile_sched_params=None, + clc_response_ptr=None, + use_clc_dynamic_scheduler: bool = False, + use_static_persistent_scheduler: bool = False, + verbose: bool = False, + exhaustive_deadlock_race_check: bool = False, +) -> tuple[TaskManager, object]: + """Build the throughput-latency 1CTA MLA TaskManager for validation or execution. + + The caller supplies either concrete tensors/descriptors for JIT execution or + leaves them as ``None`` for schedule-only validation. Exactly one persistent + scheduler mode may be enabled: CLC dynamic persistent or static persistent. + GMEM split-KV profiles wire an extra correction/reduction path, while + non-split profiles store O/LSE directly from the correction task. + """ + + scale_softmax_log2, output_scale = _default_scales( + scale_softmax_log2, + output_scale, + ) + _check_persistent_scheduler_modes( + use_clc_dynamic_scheduler, + use_static_persistent_scheduler, + ) + if cfg.kernel_variant == "keeps_mma_ab": + return _make_keeps_mma_ab_task_graph( + cfg, + total_kv_tiles=total_kv_tiles, + use_page_offsets=use_page_offsets, + tma_desc_q_latent=tma_desc_q_latent, + tma_desc_q_rope=tma_desc_q_rope, + tma_desc_c_latent=tma_desc_c_latent, + tma_desc_c_rope=tma_desc_c_rope, + tma_desc_v=tma_desc_v, + c_rope_tensor=c_rope_tensor, + page_offsets=page_offsets, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + cta_idx_head_dim_v=cta_idx_head_dim_v, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + o_tensor=o_tensor, + lse_tensor=lse_tensor, + acc_o_tensor=acc_o_tensor, + acc_lse_tensor=acc_lse_tensor, + tile_sched_params=tile_sched_params, + clc_response_ptr=clc_response_ptr, + use_clc_dynamic_scheduler=use_clc_dynamic_scheduler, + use_static_persistent_scheduler=use_static_persistent_scheduler, + verbose=verbose, + exhaustive_deadlock_race_check=exhaustive_deadlock_race_check, + ) + + agent = pipeline.Agent + cta_layout_vmnk = (1, 1, 1, 1) + tma_group = pipeline.CooperativeGroup(agent.Thread) + umma_group = pipeline.CooperativeGroup(agent.Thread) + load_group = pipeline.CooperativeGroup(agent.Thread, 32) + page_group = pipeline.CooperativeGroup(agent.Thread, 32) + scheduler_group = pipeline.CooperativeGroup( + agent.Thread, + cfg.scheduler_num_warps * 32, + ) + softmax0_group = pipeline.CooperativeGroup( + agent.Thread, + cfg.softmax_num_warps * 32, + ) + softmax1_group = pipeline.CooperativeGroup( + agent.Thread, + cfg.softmax_num_warps * 32, + ) + corr_group = pipeline.CooperativeGroup( + agent.Thread, + cfg.correction_num_warps * 32, + ) + + # Page offsets feed the load warp; Q/KV SMEM feed the MMA warp; TMEM + # score/output resources hand off to softmax and correction warps. + page_offsets_cfg = PipelineConfig( + num_stages=cfg.page_offsets_stages, + num_bytes=0, + producer_group=page_group, + consumer_group=load_group, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout_vmnk, + async_producer_op=pipeline.PipelineOp.AsyncLoad, + advance_on_wait=True, + ) + smem_q_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.q_stages, + num_bytes=cfg.qk_smem_tile_bytes, + producer_group=tma_group, + consumer_group=umma_group, + cta_layout_vmnk=cta_layout_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ) + smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.kv_stages, + num_bytes=cfg.kv_smem_tile_bytes, + producer_group=tma_group, + consumer_group=umma_group, + cta_layout_vmnk=cta_layout_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ) + tmem_s_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=1, + producer_group=umma_group, + consumer_group=softmax0_group, + cta_layout_vmnk=cta_layout_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + ) + tmem_s1_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=1, + producer_group=umma_group, + consumer_group=softmax1_group, + cta_layout_vmnk=cta_layout_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + ) + tmem_o_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=cfg.o_stages, + producer_group=umma_group, + consumer_group=corr_group, + cta_layout_vmnk=cta_layout_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + ) + tmem_o_cfg = dataclass_replace(tmem_o_cfg, advance_on_wait=True) + local0_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=softmax0_group, + consumer_group=corr_group, + cta_layout_vmnk=cta_layout_vmnk, + ) + local1_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=softmax1_group, + consumer_group=corr_group, + cta_layout_vmnk=cta_layout_vmnk, + ) + smem_p0_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=1, + producer_group=softmax0_group, + consumer_group=umma_group, + cta_layout_vmnk=cta_layout_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ) + smem_p1_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=1, + producer_group=softmax1_group, + consumer_group=umma_group, + cta_layout_vmnk=cta_layout_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ) + + work_queue = None + schedule_token_throttle = None + use_clc_dynamic = use_clc_dynamic_scheduler + use_static_persistent = use_static_persistent_scheduler + # Exactly one persistent scheduler flavor is selected by the profile. + # Non-persistent profiles leave work_queue unset and use block indices. + if use_clc_dynamic: + work_queue, schedule_token_throttle = _make_clc_work_queue_and_throttle( + cfg, + tile_sched_params, + clc_response_ptr, + cache_seqs, + cu_seqlens_q, + cta_layout_vmnk, + load_group, + scheduler_group, + "ll_mla_work_queue", + "ll_mla_schedule_token_throttle", + ) + if use_static_persistent: + work_queue = _make_static_work_queue( + cfg, + tile_sched_params, + cache_seqs, + cu_seqlens_q, + "ll_mla_work_queue", + ) + + smem_page_offsets = SmemPageOffsetsResource( + cfg=cfg, + pipeline_config=page_offsets_cfg, + page_offsets=page_offsets, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + name="ll_mla_page_offsets", + ) + smem_q = SmemQResource( + cfg=cfg, + pipeline_config=smem_q_cfg, + cu_seqlens_q=cu_seqlens_q, + name="ll_mla_smem_q", + ) + smem_q.tma_desc_q_latent = tma_desc_q_latent + smem_q.tma_desc_q_rope = tma_desc_q_rope + smem_q.head_idx = head_idx + smem_q.batch_idx = batch_idx + smem_q.cta_idx_q = cta_idx_q + smem_kv = SmemKvResource( + cfg=cfg, + pipeline_config=smem_kv_cfg, + tma_desc_c_latent=tma_desc_c_latent, + tma_desc_c_rope=tma_desc_c_rope, + tma_desc_v=tma_desc_v, + c_rope_tensor=c_rope_tensor, + page_offsets_kv=smem_page_offsets if use_page_offsets else None, + page_offsets=page_offsets, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + cta_idx_head_dim_v=cta_idx_head_dim_v, + name="ll_mla_smem_kv", + ) + tmem_s0 = TmemSResource( + cfg=cfg, + pipeline_config=tmem_s_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + sync_barrier_id=0, + name="ll_mla_tmem_s0", + ) + tmem_s1 = TmemSResource( + cfg=cfg, + pipeline_config=tmem_s1_cfg, + inst_id=1, + scale_softmax_log2=scale_softmax_log2, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + sync_barrier_id=1, + name="ll_mla_tmem_s1", + ) + order_p01_alloc = SmemAllocation( + name="ll_mla_order_p01", + size_bytes=16, + alignment=8, + ) + smem_p0 = SmemPResource( + cfg=cfg, + pipeline_config=smem_p0_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + order_p01_alloc=order_p01_alloc, + owns_order_p01_alloc=True, + name="ll_mla_smem_p0", + ) + smem_p1 = SmemPResource( + cfg=cfg, + pipeline_config=smem_p1_cfg, + inst_id=1, + scale_softmax_log2=scale_softmax_log2, + order_p01_alloc=order_p01_alloc, + name="ll_mla_smem_p1", + ) + tmem_s0.p_ref = smem_p0 + tmem_s1.p_ref = smem_p1 + smem_p0.tmem_s_ref = tmem_s0 + smem_p1.tmem_s_ref = tmem_s1 + tmem_o = TmemOResource( + cfg=cfg, + pipeline_config=tmem_o_cfg, + p0_ref=smem_p0, + p1_ref=smem_p1, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + name="ll_mla_tmem_o", + ) + local0 = TmemSoftmaxLocalResource( + cfg=cfg, + pipeline_config=local0_cfg, + inst_id=0, + name="ll_mla_local0", + ) + local1 = TmemSoftmaxLocalResource( + cfg=cfg, + pipeline_config=local1_cfg, + inst_id=1, + name="ll_mla_local1", + ) + global0 = TmemSoftmaxGlobalResource(cfg=cfg, inst_id=0, name="ll_mla_global0") + global1 = TmemSoftmaxGlobalResource(cfg=cfg, inst_id=1, name="ll_mla_global1") + corr0 = TmemCorrResource( + cfg=cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + o_tensor=o_tensor, + lse_tensor=lse_tensor, + acc_o_tensor=acc_o_tensor, + acc_lse_tensor=acc_lse_tensor, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + cta_idx_head_dim_v=cta_idx_head_dim_v, + name="ll_mla_corr0", + ) + corr1 = TmemCorrResource( + cfg=cfg, + inst_id=1, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + o_tensor=o_tensor, + lse_tensor=lse_tensor, + acc_o_tensor=acc_o_tensor, + acc_lse_tensor=acc_lse_tensor, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + cta_idx_head_dim_v=cta_idx_head_dim_v, + name="ll_mla_corr1", + ) + local_kv_tiles = cfg.local_kv_tiles(total_kv_tiles) + loop_domain = cfg.loop_domain(local_kv_tiles) + load_domain = loop_domain + mma_domain = loop_domain + softmax_domain = loop_domain + 1 + corr_domain = loop_domain + task_domain_kwargs = { + "seqlens_kv": cache_seqs, + "cu_seqlens_q": cu_seqlens_q, + } + + tasks = [] + if use_page_offsets: + tasks.append( + create_load_page_offsets_task( + smem_page_offsets, + work_queue, + cfg, + domain=load_domain, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ) + ) + tasks.extend( + [ + create_throughput_latency_load_task( + smem_q, + smem_kv, + work_queue, + schedule_token_throttle, + cfg, + domain=load_domain, + smem_page_offsets=smem_page_offsets if use_page_offsets else None, + use_page_offsets=use_page_offsets, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ), + create_throughput_latency_softmax0_task( + tmem_s0, + local0, + smem_p0, + global0, + work_queue, + cfg, + domain=softmax_domain, + task_class=MlaDecodeTask, + domain_bias=1, + **task_domain_kwargs, + ), + create_throughput_latency_softmax1_task( + tmem_s1, + local1, + smem_p1, + global1, + work_queue, + cfg, + domain=softmax_domain, + task_class=MlaDecodeTask, + domain_bias=1, + **task_domain_kwargs, + ), + create_throughput_latency_correction_task( + local0, + local1, + tmem_o, + corr0, + corr1, + work_queue, + cfg, + domain=corr_domain, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ), + create_throughput_latency_mma_task( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue, + cfg, + domain=mma_domain, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ), + ] + ) + if not (use_page_offsets and use_clc_dynamic): + tasks.append( + create_padding_task( + cfg, + work_queue, + warp_idx=( + cfg.page_offsets_warp_idx + cfg.page_offsets_num_warps + if use_page_offsets + else (cfg.clc_padding_warp_idx if use_clc_dynamic else None) + ), + num_warps=( + 1 + if use_page_offsets + else (cfg.clc_padding_num_warps if use_clc_dynamic else None) + ), + task_class=MlaDecodeTask, + ) + ) + if use_clc_dynamic: + tasks.append( + create_throughput_latency_scheduler_task( + work_queue, + schedule_token_throttle, + cfg, + task_class=MlaDecodeTask, + ) + ) + + deps = { + smem_q: [], + smem_kv: [], + tmem_s0: [smem_q, smem_kv], + tmem_s1: [smem_q, smem_kv], + smem_p0: [tmem_s0], + smem_p1: [tmem_s1], + global0: [tmem_s0], + global1: [tmem_s1], + local0: [tmem_s0], + local1: [tmem_s1], + tmem_o: [smem_p0, smem_p1, smem_kv], + corr0: [local0, tmem_o], + corr1: [local0, local1, tmem_o], + } + if use_page_offsets: + deps[smem_page_offsets] = [] + deps[smem_kv].append(smem_page_offsets) + if work_queue is not None: + work_queue_deps = [work_queue] if use_clc_dynamic else [] + deps = { + smem_q: [work_queue], + smem_kv: [work_queue], + tmem_s0: [smem_q, smem_kv, work_queue], + tmem_s1: [smem_q, smem_kv, work_queue], + smem_p0: [tmem_s0, work_queue], + smem_p1: [tmem_s1, work_queue], + global0: [tmem_s0, work_queue], + global1: [tmem_s1, work_queue], + local0: [tmem_s0, work_queue], + local1: [tmem_s1, work_queue], + tmem_o: [smem_p0, smem_p1, smem_kv, work_queue], + corr0: [local0, tmem_o, work_queue], + corr1: [local0, local1, tmem_o, work_queue], + work_queue: ( + work_queue_deps + [schedule_token_throttle] + if schedule_token_throttle is not None + else work_queue_deps + ), + } + if use_page_offsets: + deps[smem_page_offsets] = [work_queue] + deps[smem_kv].append(smem_page_offsets) + if schedule_token_throttle is not None: + deps[schedule_token_throttle] = [work_queue] + + dma_release_labels = { + (smem_kv, tmem_s0): {"k_desc_0"}, + (smem_kv, tmem_s1): {"k_desc_1"}, + (smem_kv, tmem_o): {"v_desc_0", "v_desc_1"}, + } + + smem_allocator = SmemAllocator() + smem_allocator.add_resource(smem_q) + if use_page_offsets: + smem_allocator.add_resource(smem_page_offsets) + smem_allocator.add_resource(smem_kv) + smem_allocator.add_resource(smem_p0) + smem_allocator.add_resource(smem_p1) + smem_allocator.add_resource(tmem_s0) + smem_allocator.add_resource(tmem_s1) + smem_allocator.add_resource(tmem_o) + smem_allocator.add_resource(local0) + smem_allocator.add_resource(local1) + smem_allocator.add_resource(global0) + smem_allocator.add_resource(global1) + smem_allocator.add_resource(corr0) + smem_allocator.add_resource(corr1) + smem_allocator.add_tmem_ptr( + SmemAllocation("ll_mla_tmem_ptr_i32", dtype=cutlass.Int32, alignment=4) + ) + smem_allocator.compute_layout() + + tmem_allocator = TmemAllocator() + tmem_allocator.add_resource(tmem_s0) + tmem_allocator.add_resource(tmem_s1) + tmem_allocator.add_resource(local0) + tmem_allocator.add_resource(local1) + tmem_allocator.add_resource(tmem_o) + tmem_allocator.compute_layout() + + task_manager = TaskManager( + tasks=tasks, + resource_dependency_graph=deps, + dma_consumer_release_labels=dma_release_labels, + smem_allocator=smem_allocator, + tmem_allocator=tmem_allocator, + verbose=verbose, + exhaustive_deadlock_race_check=exhaustive_deadlock_race_check, + ) + return task_manager, corr1 + + +def _make_keeps_mma_ab_task_graph( + cfg: MlaConfig, + *, + total_kv_tiles: int, + use_page_offsets: bool = True, + tma_desc_q_latent=None, + tma_desc_q_rope=None, + tma_desc_c_latent=None, + tma_desc_c_rope=None, + tma_desc_v=None, + c_rope_tensor=None, + page_offsets=None, + cache_seqs=None, + cu_seqlens_q=None, + head_idx=None, + batch_idx=None, + cta_idx_q=None, + cta_idx_kv=None, + cta_idx_head_dim_v=None, + scale_softmax_log2=None, + output_scale=None, + o_tensor=None, + lse_tensor=None, + acc_o_tensor=None, + acc_lse_tensor=None, + tile_sched_params=None, + clc_response_ptr=None, + use_clc_dynamic_scheduler: bool = False, + use_static_persistent_scheduler: bool = False, + verbose: bool = False, + exhaustive_deadlock_race_check: bool = False, +) -> tuple[TaskManager, object]: + """Build the keeps-MMA-AB 1CTA task graph for the generic builder. + + This variant uses one score/P pipe with P stored in TMEM. The swaps path + keeps its two SMEM-P pipes and is intentionally left untouched. + """ + + scale_softmax_log2, output_scale = _default_scales( + scale_softmax_log2, + output_scale, + ) + _check_persistent_scheduler_modes( + use_clc_dynamic_scheduler, + use_static_persistent_scheduler, + ) + agent = pipeline.Agent + cta_layout_vmnk = (1, 1, 1, 1) + tma_group = pipeline.CooperativeGroup(agent.Thread) + umma_group = pipeline.CooperativeGroup(agent.Thread) + load_group = pipeline.CooperativeGroup(agent.Thread, 32) + page_group = pipeline.CooperativeGroup(agent.Thread, 32) + softmax_group = pipeline.CooperativeGroup( + agent.Thread, + cfg.softmax_num_warps * 32, + ) + corr_group = pipeline.CooperativeGroup( + agent.Thread, + cfg.correction_num_warps * 32, + ) + + page_offsets_cfg = PipelineConfig( + num_stages=cfg.page_offsets_stages, + num_bytes=0, + producer_group=page_group, + consumer_group=load_group, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout_vmnk, + async_producer_op=pipeline.PipelineOp.AsyncLoad, + advance_on_wait=True, + ) + smem_q_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.q_stages, + num_bytes=cfg.qk_smem_tile_bytes, + producer_group=tma_group, + consumer_group=umma_group, + cta_layout_vmnk=cta_layout_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ) + smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.kv_stages, + num_bytes=cfg.kv_smem_tile_bytes, + producer_group=tma_group, + consumer_group=umma_group, + cta_layout_vmnk=cta_layout_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ) + tmem_s_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=2, + producer_group=umma_group, + consumer_group=softmax_group, + cta_layout_vmnk=cta_layout_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + ) + local_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=2, + producer_group=softmax_group, + consumer_group=corr_group, + cta_layout_vmnk=cta_layout_vmnk, + ) + tmem_p_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=2, + producer_group=softmax_group, + consumer_group=umma_group, + cta_layout_vmnk=cta_layout_vmnk, + consumer_signaling_threads=SignalingThreads.CtaLeader, + ) + tmem_o_cfg = PipelineConfig.create_umma_async_pipeline_cfg( + num_stages=1, + producer_group=umma_group, + consumer_group=corr_group, + cta_layout_vmnk=cta_layout_vmnk, + producer_signaling_threads=SignalingThreads.CtaLeader, + ) + + work_queue = None + schedule_token_throttle = None + use_clc_dynamic = use_clc_dynamic_scheduler + if use_static_persistent_scheduler: + work_queue = _make_static_work_queue( + cfg, + tile_sched_params, + cache_seqs, + cu_seqlens_q, + "ll_mla_q64_work_queue", + ) + if use_clc_dynamic: + scheduler_group = pipeline.CooperativeGroup(agent.Thread, 32) + work_queue, schedule_token_throttle = _make_clc_work_queue_and_throttle( + cfg, + tile_sched_params, + clc_response_ptr, + cache_seqs, + cu_seqlens_q, + cta_layout_vmnk, + load_group, + scheduler_group, + "ll_mla_q64_work_queue", + "ll_mla_q64_schedule_token_throttle", + ) + + smem_page_offsets = SmemPageOffsetsResource( + cfg=cfg, + pipeline_config=page_offsets_cfg, + page_offsets=page_offsets, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + name="ll_mla_q64_page_offsets", + ) + smem_q = SmemQResource( + cfg=cfg, + pipeline_config=smem_q_cfg, + tma_desc_q_latent=tma_desc_q_latent, + tma_desc_q_rope=tma_desc_q_rope, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + name="ll_mla_q64_smem_q", + ) + smem_kv = SmemKvResource( + cfg=cfg, + pipeline_config=smem_kv_cfg, + tma_desc_c_latent=tma_desc_c_latent, + tma_desc_c_rope=tma_desc_c_rope, + tma_desc_v=tma_desc_v, + c_rope_tensor=c_rope_tensor, + page_offsets_kv=smem_page_offsets if use_page_offsets else None, + page_offsets=page_offsets, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + cta_idx_head_dim_v=cta_idx_head_dim_v, + name="ll_mla_q64_smem_kv", + ) + tmem_s = TmemSKeepsResource( + cfg=cfg, + pipeline_config=tmem_s_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + sync_barrier_id=0, + name="ll_mla_q64_tmem_s", + ) + tmem_p = TmemPResource( + cfg=cfg, + pipeline_config=tmem_p_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + tmem_alias_ref=tmem_s, + name="ll_mla_q64_tmem_p", + ) + tmem_s.p_ref = tmem_p + tmem_o = TmemOResource( + cfg=cfg, + pipeline_config=tmem_o_cfg, + p_tmem_ref=tmem_p, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + name="ll_mla_q64_tmem_o", + ) + local = TmemSoftmaxLocalResource( + cfg=cfg, + pipeline_config=local_cfg, + inst_id=0, + tmem_alias_ref=tmem_s, + name="ll_mla_q64_local", + ) + global_softmax = TmemSoftmaxGlobalResource( + cfg=cfg, inst_id=0, name="ll_mla_q64_global" + ) + corr = TmemCorrResource( + cfg=cfg, + inst_id=1, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + o_tensor=o_tensor, + lse_tensor=lse_tensor, + acc_o_tensor=acc_o_tensor, + acc_lse_tensor=acc_lse_tensor, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + cta_idx_head_dim_v=cta_idx_head_dim_v, + name="ll_mla_q64_corr", + ) + + local_kv_tiles = cfg.local_kv_tiles(total_kv_tiles) + loop_domain = cfg.loop_domain(local_kv_tiles) + task_domain_kwargs = { + "seqlens_kv": cache_seqs, + "cu_seqlens_q": cu_seqlens_q, + } + + tasks = [] + if use_page_offsets: + tasks.append( + create_load_page_offsets_task( + smem_page_offsets, + work_queue, + cfg, + domain=loop_domain, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ) + ) + tasks.extend( + [ + create_throughput_latency_load_task( + smem_q, + smem_kv, + work_queue, + schedule_token_throttle, + cfg, + domain=loop_domain, + smem_page_offsets=smem_page_offsets if use_page_offsets else None, + use_page_offsets=use_page_offsets, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ), + create_keeps_mma_ab_softmax_task( + tmem_s, + local, + tmem_p, + global_softmax, + work_queue, + cfg, + domain=loop_domain + 1, + task_class=MlaDecodeTask, + domain_bias=1, + **task_domain_kwargs, + ), + create_keeps_mma_ab_correction_task( + local, + tmem_o, + corr, + work_queue, + cfg, + domain=loop_domain, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ), + create_keeps_mma_ab_mma_task( + smem_q, + smem_kv, + tmem_s, + tmem_p, + tmem_o, + work_queue, + cfg, + domain=loop_domain, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ), + ] + ) + if not use_clc_dynamic: + tasks.append( + create_padding_task( + cfg, + work_queue, + warp_idx=11, + num_warps=1, + task_class=MlaDecodeTask, + ) + ) + elif not use_page_offsets: + tasks.append( + create_padding_task( + cfg, + work_queue, + warp_idx=cfg.page_offsets_warp_idx, + num_warps=cfg.page_offsets_num_warps, + task_class=MlaDecodeTask, + ) + ) + if use_clc_dynamic: + tasks.append( + create_throughput_latency_scheduler_task( + work_queue, + schedule_token_throttle, + cfg, + task_class=MlaDecodeTask, + ) + ) + + deps = { + smem_q: [], + smem_kv: [], + tmem_s: [smem_q, smem_kv], + local: [tmem_s], + global_softmax: [tmem_s], + tmem_p: [tmem_s], + tmem_o: [tmem_p, smem_kv], + corr: [local, tmem_o], + } + if use_page_offsets: + deps[smem_page_offsets] = [] + deps[smem_kv].append(smem_page_offsets) + if work_queue is not None: + deps = { + smem_q: [work_queue], + smem_kv: [work_queue], + tmem_s: [smem_q, smem_kv, work_queue], + local: [tmem_s, work_queue], + global_softmax: [tmem_s, work_queue], + tmem_p: [tmem_s, work_queue], + tmem_o: [tmem_p, smem_kv, work_queue], + corr: [local, tmem_o, work_queue], + work_queue: ( + [work_queue, schedule_token_throttle] + if schedule_token_throttle is not None + else [] + ), + } + if use_page_offsets: + deps[smem_page_offsets] = [work_queue] + deps[smem_kv].append(smem_page_offsets) + if schedule_token_throttle is not None: + deps[schedule_token_throttle] = [work_queue] + + dma_release_labels = { + (smem_kv, tmem_s): {"k_desc_0"}, + (smem_kv, tmem_o): {"v_desc_0"}, + } + + smem_allocator = SmemAllocator() + smem_allocator.add_resource(smem_q) + if use_page_offsets: + smem_allocator.add_resource(smem_page_offsets) + for resource in (smem_kv, tmem_s, local, global_softmax, tmem_p, tmem_o, corr): + smem_allocator.add_resource(resource) + smem_allocator.add_tmem_ptr( + SmemAllocation("ll_mla_q64_tmem_ptr_i32", dtype=Int32, alignment=4) + ) + smem_allocator.compute_layout() + + tmem_allocator = TmemAllocator() + tmem_allocator.add_resource(tmem_s) + tmem_allocator.add_resource(local) + tmem_allocator.add_resource(tmem_o) + tmem_allocator.compute_layout() + + task_manager = TaskManager( + tasks=tasks, + resource_dependency_graph=deps, + dma_consumer_release_labels=dma_release_labels, + smem_allocator=smem_allocator, + tmem_allocator=tmem_allocator, + verbose=verbose, + exhaustive_deadlock_race_check=exhaustive_deadlock_race_check, + ) + return task_manager, corr + + +class ThroughputLatencyMlaDecodeTs: + """Dense throughput-latency 1CTA MLA TS wrapper.""" + + def __init__( + self, + *, + batch_size: int, + num_heads: int, + seq_len_q: int, + seq_len_k: int, + latent_dim: int = 512, + rope_dim: int = 64, + page_size: int = 32, + max_active_clusters: int, + acc_dtype=None, + lse_dtype=None, + qkv_dtype: str = "bf16", + out_dtype: str = "bf16", + profile: str | None = None, + persistent_wave_sm_count: int | None = None, + reduction_mode: str | None = None, + logical_num_heads: int, + logical_seq_len_q: int, + tile_size_q: int, + explicit_split_kv: int | None = None, + explicit_persistent: bool | None = None, + mask_type: MaskType | str = MaskType.CAUSAL, + ): + """Initialize one selected physical tile profile over logical flat Q rows.""" + import cutlass as _cutlass + + if acc_dtype is None: + acc_dtype = _cutlass.Float32 + if lse_dtype is None: + lse_dtype = _cutlass.Float32 + + self.batch_size = batch_size + self.num_heads = num_heads + self.seq_len_q = seq_len_q + self.seq_len_k = seq_len_k + self.latent_dim = latent_dim + self.rope_dim = rope_dim + self.page_size = page_size + self.max_active_clusters = max_active_clusters + self.qkv_dtype = qkv_dtype + self.out_dtype = out_dtype + self.profile = profile + self.persistent_wave_sm_count = persistent_wave_sm_count + self.reduction_mode = reduction_mode + self.logical_num_heads = logical_num_heads + self.logical_seq_len_q = logical_seq_len_q + self.tile_size_q = tile_size_q + self.explicit_split_kv = explicit_split_kv + self.explicit_persistent = explicit_persistent + self.mask_type = normalize_mask_type(mask_type) + + cfg = self._make_config() + # Parallel standalone reduction requires a fixed split profile and a + # static producer schedule so its topology and workspace contract are + # compile-time invariant. Other profiles use the general reducer. + self.use_parallel_reduction = ( + supports_parallel_gmem_reduction(cfg) + and lse_dtype == _cutlass.Float32 + and cfg.use_persistent_scheduler == 0 + and cfg.use_clc_dynamic_persistent_scheduler == 0 + ) + self.parallel_reduction_topology = None + self.parallel_reduction_elements_per_slice = ( + PARALLEL_GMEM_REDUCTION_SWAPS_ELEMENTS_PER_SLICE + if cfg.tile_size_q in (8, 16, 32) + else PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_SLICE + ) + if cfg.use_multi_ctas_kv == 1 and cfg.use_cluster_reduction != 1: + # Both reducer implementations use the same normalized partial-O + # workspace. Validate every separate-GMEM launch against it. + validate_parallel_reduction_workspace( + batch_size=cfg.batch_size, + num_heads_q=cfg.num_heads_q, + seq_len_q=cfg.seq_len_q, + splits_kv=cfg.num_ctas_per_seq_kv, + head_dim=cfg.head_dim_v, + ) + if self.use_parallel_reduction: + base_clusters = parallel_gmem_reduction_base_clusters( + cfg, + self.parallel_reduction_elements_per_slice, + ) + cluster_size = ( + choose_q64_parallel_reducer_cluster_size( + cfg.num_ctas_per_seq_kv, + base_clusters=base_clusters, + sm_count=self.max_active_clusters, + ) + if cfg.tile_size_q == 64 + else 1 + ) + self.parallel_reduction_topology = ( + make_balanced_parallel_reduction_topology( + cfg.num_ctas_per_seq_kv, + cluster_size=cluster_size, + ) + ) + self.acc_dtype = acc_dtype + self.lse_dtype = lse_dtype + + def compile_topology_signature(self) -> tuple[object, ...]: + """Describe batch-derived reducer choices without retaining batch.""" + + cfg = self._make_config() + if cfg.use_multi_ctas_kv != 1 or cfg.use_cluster_reduction == 1: + return ("no_separate_reducer",) + if self.use_parallel_reduction: + return ( + "parallel", + self.parallel_reduction_topology, + self.parallel_reduction_elements_per_slice, + ) + grid, smem, threads, reduction_ctas = gmem_reduction_launch_shape( + cfg, + cfg.seq_len_q, + cfg.batch_size, + self.lse_dtype.width, + self.max_active_clusters, + ) + return ( + "reference", + grid[0], + grid[1], + smem, + threads, + reduction_ctas, + ) + + def compile_signature(self) -> tuple[object, ...]: + """Return the complete batch-independent JIT identity.""" + + cfg = dataclass_replace(self._make_config(), batch_size=1) + return ( + cfg, + self.acc_dtype, + self.lse_dtype, + self.use_parallel_reduction, + self.parallel_reduction_topology, + self.parallel_reduction_elements_per_slice, + self.compile_topology_signature(), + ) + + def _make_config(self): + """Create the static throughput-latency MLA config for this wrapper.""" + cfg = make_throughput_latency_mla_config( + batch_size=self.batch_size, + num_heads_q=self.num_heads, + seq_len_q=self.seq_len_q, + seq_len_kv=self.seq_len_k, + latent_dim=self.latent_dim, + rope_dim=self.rope_dim, + num_tokens_per_page=self.page_size, + qkv_dtype=self.qkv_dtype, + o_dtype=self.out_dtype, + profile=self.profile, + persistent_wave_sm_count=self.persistent_wave_sm_count, + max_active_clusters=self.max_active_clusters, + reduction_mode=self.reduction_mode, + logical_num_heads_q=self.logical_num_heads, + logical_seq_len_q=self.logical_seq_len_q, + tile_size_q=self.tile_size_q, + explicit_split_kv=self.explicit_split_kv, + explicit_persistent=self.explicit_persistent, + mask_type=self.mask_type, + ) + return cfg + + def validate_split_kv_launch(self, split_kv: int, workspace) -> None: + """Validate host-side split-KV arguments before compiling the JIT body.""" + + cfg = self._make_config() + if ( + cfg.use_multi_ctas_kv == 1 + and cfg.use_cluster_reduction != 1 + and workspace is None + ): + raise ValueError( + "multi-CTA-KV throughput-latency 1CTA MLA requires workspace" + ) + if cfg.use_multi_ctas_kv == 1 and split_kv < cfg.num_ctas_per_seq_kv: + raise ValueError( + "split_kv is smaller than the configured multi-CTA-KV split count" + ) + + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + workspace: cute.Tensor, + ): + """Construct throughput-latency 1CTA split-KV GMEM reduction tensors.""" + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + align = 256 // cutlass.Float16.width + acc_o_layout = cute.make_layout( + (H, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, + cute.assume(split_kv * H * D, align), + cute.assume(H * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=cutlass.BFloat16) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( + (H, split_kv, S, B), + stride=(split_kv, 1, H * split_kv, H * split_kv * S), + ) + acc_lse_iter = cute.recast_ptr( + workspace.iterator + + Int64(cute.cosize(acc_o_layout)) * Int64(cutlass.Float16.width // 8), + dtype=self.lse_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_offsets: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor | None, + block_split_kvs: cute.Tensor, + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: object, + ): + """Execute throughput-latency 1CTA MLA with fixed or compact ragged Q.""" + cfg = self._make_config() + + # Public fixed tensors use [H,D,SQ,B]/[H,SQ,B]; compact ragged tensors + # use [H,D,totalQ]/[H,totalQ]. Both flatten H x Q for TMA because the + # scheduler operates on consecutive physical row tiles. Resource-level + # predicates own the final partial tile and map outputs back to storage. + tma_box0 = min(128 // cfg.qkv_dtype_bytes, cfg.head_dim_per_stage_kv) + tma_page_tokens = cfg.num_tokens_per_page + if cutlass.const_expr(q_latent.stride[1] != 1 or q_rope.stride[1] != 1): + raise ValueError("q_latent and q_rope must have leading dimension 1") + runtime_assert( + q_latent.stride[2] == q_latent.shape[0] * q_latent.stride[0], + "q_latent must be compact across the head and query dimensions", + ) + runtime_assert( + q_rope.stride[2] == q_rope.shape[0] * q_rope.stride[0], + "q_rope must be compact across the head and query dimensions", + ) + if cutlass.const_expr(c_latent.stride[1] != 1 or c_rope.stride[1] != 1): + raise ValueError("c_latent and c_rope must have leading dimension 1") + if cutlass.const_expr(o.stride[1] != 1): + raise ValueError("o must have leading dimension 1") + runtime_assert( + o.stride[0] == o.shape[1] * o.stride[1], + "o must be compact from the dimension axis into the head axis", + ) + runtime_assert( + o.stride[2] == o.shape[0] * o.stride[0], + "o must be compact from the head axis into the query axis", + ) + if cutlass.const_expr(cu_seqlens_q is None): + runtime_assert( + o.stride[3] == o.shape[2] * o.stride[2], + "o must be compact from the query axis into the batch axis", + ) + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") + runtime_assert( + lse.stride[1] == lse.shape[0] * lse.stride[0], + "lse must be compact from the head axis into the query axis", + ) + if cutlass.const_expr(cu_seqlens_q is None): + runtime_assert( + lse.stride[2] == lse.shape[1] * lse.stride[1], + "lse must be compact from the query axis into the batch axis", + ) + else: + runtime_assert( + cute.size(cu_seqlens_q) == cute.size(cache_seqs) + Int32(1), + "cu_seqlens_q must contain batch_size + 1 offsets", + ) + if cutlass.const_expr( + cfg.use_multi_ctas_kv == 1 + and cfg.use_cluster_reduction != 1 + and workspace is None + ): + raise ValueError( + "multi-CTA-KV throughput-latency 1CTA MLA requires workspace" + ) + workspace_split_kv = split_kv + if cutlass.const_expr(cfg.use_multi_ctas_kv == 1): + workspace_split_kv = cutlass.Int32(cfg.num_ctas_per_seq_kv) + runtime_assert( + split_kv >= workspace_split_kv, + "split_kv is smaller than the configured multi-CTA-KV split count", + ) + + if cutlass.const_expr(cu_seqlens_q is not None): + q_latent_tma = cute.make_tensor( + q_latent.iterator, + cute.make_layout( + (q_latent.shape[1], q_latent.shape[0] * q_latent.shape[2]), + stride=(q_latent.stride[1], q_latent.stride[0]), + ), + ) + tma_desc_q_latent = create_tensor_map_ragged_from_tensor( + q_latent_tma, + box_dims=(tma_box0, cfg.tile_size_q), + ragged_dim=1, + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + q_latent_tma = cute.make_tensor( + q_latent.iterator, + cute.make_layout( + ( + q_latent.shape[1], + q_latent.shape[0] * q_latent.shape[2], + q_latent.shape[3], + ), + stride=( + q_latent.stride[1], + q_latent.stride[0], + q_latent.stride[3], + ), + ), + ) + tma_desc_q_latent = create_tensor_map_tiled_from_view( + q_latent_tma, + box_dims=(tma_box0, cfg.tile_size_q, 1), + stride_order=(0, 1, 2), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + q_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + q_rope_swizzle = cuda.TensorMapSwizzle.s64b + if cutlass.const_expr(cu_seqlens_q is not None): + q_rope_tma = cute.make_tensor( + q_rope.iterator, + cute.make_layout( + (q_rope.shape[1], q_rope.shape[0] * q_rope.shape[2]), + stride=(q_rope.stride[1], q_rope.stride[0]), + ), + ) + tma_desc_q_rope = create_tensor_map_ragged_from_tensor( + q_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q), + ragged_dim=1, + stride_order=(0, 1), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + q_rope_tma = cute.make_tensor( + q_rope.iterator, + cute.make_layout( + ( + q_rope.shape[1], + q_rope.shape[0] * q_rope.shape[2], + q_rope.shape[3], + ), + stride=( + q_rope.stride[1], + q_rope.stride[0], + q_rope.stride[3], + ), + ), + ) + tma_desc_q_rope = create_tensor_map_tiled_from_view( + q_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q, 1), + stride_order=(0, 1, 2), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + c_latent_tma = cute.make_tensor( + c_latent.iterator, + cute.select(c_latent.layout, mode=[1, 0, 2]), + ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_latent_tma, + box_dims=(tma_box0, tma_page_tokens, 1), + stride_order=(0, 1, 2), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + c_rope_tma = cute.make_tensor( + c_rope.iterator, + cute.select(c_rope.layout, mode=[1, 0, 2]), + ) + c_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + c_rope_swizzle = cuda.TensorMapSwizzle.s64b + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), tma_page_tokens, 1), + stride_order=(0, 1, 2), + swizzle=c_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + softmax_scale_log2 = softmax_scale * LOG2_E + use_gmem_reduction = cutlass.const_expr( + cfg.use_multi_ctas_kv == 1 and cfg.use_cluster_reduction != 1 + ) + batch_size = Int32(cute.size(cache_seqs)) + if cutlass.const_expr(cu_seqlens_q is None): + runtime_assert( + cute.size(o.shape[3]) == batch_size, + "fixed output batch size must match cache_seqs", + ) + acc_o, acc_lse = self.initialize_workspace( + cutlass.Int32(cfg.num_heads_q), + cfg.head_dim_v, + cutlass.Int32(cfg.seq_len_q), + batch_size, + workspace_split_kv, + workspace if use_gmem_reduction else None, + ) + + use_clc_dynamic = cutlass.const_expr( + cfg.use_clc_dynamic_persistent_scheduler == 1 + ) + use_static_persistent = cutlass.const_expr( + cfg.use_persistent_scheduler == 1 + and cfg.use_clc_dynamic_persistent_scheduler != 1 + ) + tile_sched_params = None + if cutlass.const_expr(use_clc_dynamic): + tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + ( + cfg.num_ctas_per_seq_q, + cfg.num_ctas_per_head_dim, + batch_size * Int32(cfg.num_ctas_for_all_heads), + ), + (1, 1, 1), + ) + grid = tile_sched_params.get_grid_shape() + elif cutlass.const_expr(use_static_persistent): + tile_sched_params = utils.PersistentTileSchedulerParams( + ( + cfg.num_ctas_per_seq_q, + cfg.num_ctas_per_head_dim, + batch_size * Int32(cfg.num_ctas_for_all_heads), + ), + (1, 1, 1), + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, + self.max_active_clusters, + ) + else: + grid = ( + cfg.num_ctas_for_all_heads + * cfg.num_ctas_per_seq_q + * cfg.num_ctas_per_seq_kv, + cfg.num_ctas_per_head_dim, + batch_size, + ) + cluster_shape = None + if cutlass.const_expr(cfg.use_cluster_reduction == 1): + cluster_shape = (cfg.num_ctas_per_seq_kv, 1, 1) + self.dense_kernel( + tma_desc_q_latent, + tma_desc_q_rope, + tma_desc_c_latent, + tma_desc_c_rope, + tma_desc_c_latent, + c_rope, + page_offsets, + o, + lse, + acc_o, + acc_lse, + cache_seqs, + cu_seqlens_q, + softmax_scale_log2, + output_scale, + tile_sched_params, + ).launch( + grid=grid, + block=[cfg.threads_per_cta, 1, 1], + cluster=cluster_shape, + stream=stream, + min_blocks_per_mp=1, + use_pdl=acc_o is not None, + ) + if cutlass.const_expr(acc_o is not None): + if cutlass.const_expr(self.use_parallel_reduction): + topology = self.parallel_reduction_topology + reduction_grid, reduction_cluster = ( + parallel_gmem_reduction_launch_shape( + cfg, + topology, + self.parallel_reduction_elements_per_slice, + ) + ) + reduction_grid = ( + reduction_grid[0], + reduction_grid[1], + batch_size, + ) + reduction_threads = parallel_gmem_reduction_threads( + self.parallel_reduction_elements_per_slice + ) + parallel_reducer = self.parallel_gmem_reduction_kernel( + o, + lse, + acc_o, + acc_lse, + cache_seqs, + cu_seqlens_q, + ) + if cutlass.const_expr(topology.cluster_size == 1): + parallel_reducer.launch( + grid=reduction_grid, + block=[reduction_threads, 1, 1], + stream=stream, + min_blocks_per_mp=1, + use_pdl=True, + ) + else: + parallel_reducer.launch( + grid=reduction_grid, + block=[reduction_threads, 1, 1], + cluster=reduction_cluster, + stream=stream, + min_blocks_per_mp=1, + use_pdl=True, + ) + return + ( + reduction_grid, + reduction_smem, + reduction_threads, + reduction_ctas, + ) = gmem_reduction_launch_shape( + cfg, + cfg.seq_len_q, + cfg.batch_size, + self.lse_dtype.width, + self.max_active_clusters, + ) + reduction_grid = ( + reduction_grid[0], + reduction_grid[1], + batch_size, + ) + self.gmem_reduction_kernel( + o, + lse, + acc_o, + acc_lse, + cache_seqs, + cu_seqlens_q, + cutlass.Int32(reduction_ctas), + ).launch( + grid=reduction_grid, + block=[reduction_threads, 1, 1], + smem=reduction_smem, + stream=stream, + min_blocks_per_mp=1, + use_pdl=True, + ) + + @cute.kernel + def dense_kernel( + self, + tma_desc_q_latent: cutlass.GridConstant[cuda.TensorMap], + tma_desc_q_rope: cutlass.GridConstant[cuda.TensorMap], + tma_desc_c_latent: cutlass.GridConstant[cuda.TensorMap], + tma_desc_c_rope: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v: cutlass.GridConstant[cuda.TensorMap], + c_rope: cute.Tensor, + page_offsets: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + acc_o: cute.Tensor, + acc_lse: cute.Tensor, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + tile_sched_params: object, + ): + """Execute one flat-Q, batch, KV-split, and V head-dimension tile.""" + cfg = self._make_config() + + # The grid is expressed in physical flat-Q tile coordinates. Decode it + # once here; Q/O resources retain responsibility for logical-row + # mapping, compact-ragged offsets, and padded-tail publication. + cta_idx_x, cta_idx_head_dim_v, batch_idx = cute.arch.block_idx() + if cutlass.const_expr(cfg.use_cluster_reduction == 1): + cta_idx_x = cta_idx_x // Int32(cfg.num_ctas_per_seq_kv) + ctas_per_head_tile = Int32(cfg.num_ctas_per_seq_q) + cta_idx_head_q = cta_idx_x // ctas_per_head_tile + cta_idx_q = cta_idx_x - cta_idx_head_q * ctas_per_head_tile + cta_idx_kv = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + else: + ctas_per_head_tile = Int32(cfg.num_ctas_per_seq_q * cfg.num_ctas_per_seq_kv) + cta_idx_head_q = cta_idx_x // ctas_per_head_tile + cta_idx_x_in_head = cta_idx_x - cta_idx_head_q * ctas_per_head_tile + cta_idx_q = cta_idx_x_in_head // Int32(cfg.num_ctas_per_seq_kv) + cta_idx_kv = cta_idx_x_in_head - cta_idx_q * Int32(cfg.num_ctas_per_seq_kv) + head_idx = cta_idx_head_q * Int32(cfg.tile_size_q) + run_task_graph = cutlass.Boolean(True) + query_tile_is_active = cutlass.Boolean(True) + split_tile_is_active = cutlass.Boolean(True) + enable_runtime_split_pruning = ( + cfg.use_multi_ctas_kv == 1 + and runtime_split_pruning_is_profitable(cfg.num_ctas_per_seq_kv) + ) + has_runtime_activity_guard = tile_sched_params is None and ( + cu_seqlens_q is not None or enable_runtime_split_pruning + ) + if cutlass.const_expr(has_runtime_activity_guard): + # Decode launch coordinates against the configured maximum grid, + # then independently drop compact-Q padding and profitable split + # suffixes. S2/S3 retain configured split work: too few mainloop + # CTAs can retire to amortize the activity branch plus mandatory + # neutral publication/reduction. Q padding remains prunable. + if cutlass.const_expr(cu_seqlens_q is not None): + query_tile_is_active = runtime_query_tile_is_active( + cfg, + cu_seqlens_q, + batch_idx, + cta_idx_q, + ) + query_tile_is_active = cute.arch.make_warp_uniform(query_tile_is_active) + if cutlass.const_expr(enable_runtime_split_pruning): + split_tile_is_active = runtime_split_tile_is_active( + cfg, + cache_seqs, + cu_seqlens_q, + batch_idx, + cta_idx_q, + cta_idx_kv, + ) + split_tile_is_active = cute.arch.make_warp_uniform(split_tile_is_active) + run_task_graph = query_tile_is_active and split_tile_is_active + if cutlass.const_expr(tile_sched_params is not None): + cta_idx_q = None + batch_idx = None + cta_idx_head_dim_v = None + cta_idx_kv = Int32(0) + head_idx = None + use_clc_dynamic = cutlass.const_expr( + cfg.use_clc_dynamic_persistent_scheduler == 1 + ) + use_static_persistent = cutlass.const_expr( + cfg.use_persistent_scheduler == 1 + and cfg.use_clc_dynamic_persistent_scheduler != 1 + ) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + if warp_idx == Int32(cfg.load_warp_idx): + prims.prefetch_tensormap(tma_desc_q_latent.get_ptr()) + prims.prefetch_tensormap(tma_desc_q_rope.get_ptr()) + prims.prefetch_tensormap(tma_desc_c_latent.get_ptr()) + prims.prefetch_tensormap(tma_desc_c_rope.get_ptr()) + prims.prefetch_tensormap(tma_desc_v.get_ptr()) + + clc_response_ptr = None + if cutlass.const_expr(use_clc_dynamic): + clc_response_ptr = cute.arch.alloc_smem(cutlass.Int128, 2) + + task_manager, cluster_corr_resource = build_throughput_latency_mla_task_manager( + cfg, + total_kv_tiles=cfg.total_kv_tiles, + use_page_offsets=True, + tma_desc_q_latent=tma_desc_q_latent.get_ptr(), + tma_desc_q_rope=tma_desc_q_rope.get_ptr(), + tma_desc_c_latent=tma_desc_c_latent.get_ptr(), + tma_desc_c_rope=tma_desc_c_rope.get_ptr(), + tma_desc_v=tma_desc_v.get_ptr(), + c_rope_tensor=c_rope, + page_offsets=page_offsets, + cache_seqs=cache_seqs, + cu_seqlens_q=cu_seqlens_q, + head_idx=head_idx, + batch_idx=batch_idx, + cta_idx_q=cta_idx_q, + cta_idx_kv=cta_idx_kv, + cta_idx_head_dim_v=cta_idx_head_dim_v, + scale_softmax_log2=softmax_scale_log2, + output_scale=output_scale, + o_tensor=o, + lse_tensor=lse, + acc_o_tensor=acc_o, + acc_lse_tensor=acc_lse, + tile_sched_params=tile_sched_params, + clc_response_ptr=clc_response_ptr, + use_clc_dynamic_scheduler=use_clc_dynamic, + use_static_persistent_scheduler=use_static_persistent, + ) + task_manager.setup_resources_and_tasks() + smem_allocator = task_manager.smem_allocator + assert smem_allocator is not None + tmem_ptr_alloc = smem_allocator.tmem_ptr_alloc + assert tmem_ptr_alloc is not None + tmem_ptr_i32 = smem_allocator.get(tmem_ptr_alloc) + + if cutlass.const_expr(cfg.use_cluster_reduction == 1): + resource_context = ResourceContext( + smem_base=smem_allocator.smem_base, + tmem_ptr_i32=tmem_ptr_i32, + ) + cluster_corr_resource.create_cluster_function_variables(resource_context) + + prims.fence_mbarrier_init() + if cutlass.const_expr(cfg.use_cluster_reduction == 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync( + barrier_id=TMEM_LIFECYCLE_BARRIER_ID, + thread_count=cfg.threads_per_cta, + ) + + if warp_idx == Int32(cfg.mma_warp_idx): + prims.tcgen05_alloc( + tmem_ptr_i32, + cfg.tmem_alloc_cols, + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_relinquish_alloc_permit(group=prims.CTAGroup.CTA_1) + + prims.barrier_cta_sync( + barrier_id=TMEM_LIFECYCLE_BARRIER_ID, + thread_count=cfg.threads_per_cta, + ) + if cutlass.const_expr(has_runtime_activity_guard): + if run_task_graph: + task_manager.run() + elif query_tile_is_active: + if cutlass.const_expr(cfg.use_cluster_reduction == 1): + # Publish an online-softmax identity and retain this rank's + # configured cluster row-owner reduction duties. + cluster_corr_resource.publish_neutral_cluster_partial_and_reduce( + batch_idx, + head_idx, + cta_idx_q, + cta_idx_kv, + cta_idx_head_dim_v, + ) + elif cutlass.const_expr(acc_o is not None): + # Standalone reducers retain configured split loops, so a + # pruned producer must overwrite its workspace slot before + # releasing the PDL-dependent reducer. + _publish_neutral_standalone_partial( + cfg, + acc_o, + acc_lse, + batch_idx, + cta_idx_q, + cta_idx_kv, + cta_idx_head_dim_v, + head_idx, + ) + else: + # In particular, SQ1/S2 without packed-Q metadata lowers to the + # original straight-line task graph with no runtime activity test. + task_manager.run() + prims.barrier_cta_sync( + barrier_id=TMEM_LIFECYCLE_BARRIER_ID, + thread_count=cfg.threads_per_cta, + ) + + if warp_idx == Int32(cfg.mma_warp_idx): + tmem_arr_for_dealloc = prims.make_tmem_ptr( + tmem_ptr_i32.load(), + Float32, + ) + prims.tcgen05_dealloc( + tmem_arr_for_dealloc, + cfg.tmem_alloc_cols, + group=prims.CTAGroup.CTA_1, + ) + if cutlass.const_expr(acc_o is not None): + # Every CTA in the producer grid must release the dependent + # reducer, including runtime-padded Q/split CTAs that skipped the + # task graph. One elected thread emits the convergent CTA signal. + thread_idx, _, _ = cute.arch.thread_idx() + if thread_idx == Int32(0): + prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) + + @cute.kernel + def gmem_reduction_kernel( + self, + output: cute.Tensor, + lse: cute.Tensor, + acc_output: cute.Tensor, + acc_lse: cute.Tensor, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor, + num_reduction_ctas: cutlass.Int32, + ): + """Dispatch the throughput-latency split-KV reduction body.""" + cfg = self._make_config() + run_gmem_reduction_kernel( + self, + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + num_reduction_ctas, + ) + + @cute.kernel + def parallel_gmem_reduction_kernel( + self, + output: cute.Tensor, + lse: cute.Tensor, + acc_output: cute.Tensor, + acc_lse: cute.Tensor, + cache_seqs: cute.Tensor, + cu_seqlens_q: cute.Tensor, + ): + """Dispatch the automatically selected parallel standalone reducer.""" + + cfg = self._make_config() + topology = self.parallel_reduction_topology + run_parallel_gmem_reduction_kernel( + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + topology.cluster_size, + topology.slots_per_rank, + topology.actual_splits, + self.parallel_reduction_elements_per_slice, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py new file mode 100644 index 000000000000..bee9c7bfa7a7 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py @@ -0,0 +1,808 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parallel standalone split-KV reduction for throughput-latency 1CTA MLA.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 +from cutlass.experimental import primitives as prims + +from ...separate_reduction import finalize_log2_sum_exp, normalized_lse_weight +from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ..helpers.math import ceil_div +from ..helpers.ops import ( + fmax_f32, + warp_reduce_max_f32, + warp_reduce_sum_f32, +) +from ..helpers.query import flat_query_row_state, public_query_flat_row +from ..parallel_reduction_topology import ParallelReductionTopology +from .config import MlaConfig + + +GMEM_REDUCTION_WARP_LANES = 32 +GMEM_REDUCTION_WARPS_PER_CTA = 4 + +# The parallel reducer always uses 128 threads. Q8/Q16/Q32 assign one scalar +# from a 128-element output band to each thread. Q64 uses up to one BF16 vec8 +# per thread, but caps a slice at four D-per-CTA rows so the four warps can +# still compute one row's shared LSE state each. +PARALLEL_GMEM_REDUCTION_THREADS = 128 +PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_THREAD = 8 +PARALLEL_GMEM_REDUCTION_SWAPS_ELEMENTS_PER_SLICE = 128 +PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_SLICE = ( + PARALLEL_GMEM_REDUCTION_THREADS * PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_THREAD +) +_PARALLEL_GMEM_REDUCTION_SUPPORTED_SLICE_ELEMENTS = ( + PARALLEL_GMEM_REDUCTION_SWAPS_ELEMENTS_PER_SLICE, + 512, + PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_SLICE, +) + + +@cute.jit +def _split_o_row_base_offset( + cfg: MlaConfig, + batch_idx: Int32, + q_idx: Int32, + head_idx: Int32, + dim_idx: Int32, +) -> Int64: + """Return the S=0 split-O offset, widening before the first product.""" + + return ( + ( + (Int64(batch_idx) * Int64(cfg.seq_len_q) + Int64(q_idx)) + * Int64(cfg.num_heads_q) + + Int64(head_idx) + ) + * Int64(cfg.num_ctas_per_seq_kv) + ) * Int64(cfg.head_dim_v) + Int64(dim_idx) + + +@cute.jit +def _output_element_offset( + cfg: MlaConfig, + output_query_row: Int32, + dim_idx: Int32, +) -> Int64: + """Linearize one final-output element without a 32-bit row product.""" + + return Int64(output_query_row) * Int64(cfg.head_dim_v) + Int64(dim_idx) + + +def _validate_parallel_reduction_slice(elements_per_slice: int) -> None: + if elements_per_slice not in _PARALLEL_GMEM_REDUCTION_SUPPORTED_SLICE_ELEMENTS: + raise ValueError( + "parallel reducer supports only 128-, 512-, or 1,024-element slices" + ) + + +def supports_parallel_gmem_reduction(cfg: MlaConfig) -> bool: + """Return whether ``cfg`` matches the parallel-reducer envelope.""" + + return ( + cfg.use_multi_ctas_kv == 1 + and cfg.use_cluster_reduction != 1 + and cfg.head_dim_per_cta_v in (128, 256, 512) + and cfg.tile_size_q in (8, 16, 32, 64) + and 2 <= cfg.num_ctas_per_seq_kv <= 128 + and ( + (cfg.o_dtype == "bf16" and cfg.use_bf16_output == 1) + or (cfg.o_dtype == "e4m3" and cfg.use_fp8_output == 1) + ) + # Split partial O is always BF16 in initialize_workspace. Keep the + # explicit byte check so a future workspace dtype change fails closed. + and cfg.partial_o_dtype_bytes == 2 + ) + + +def _parallel_reduction_effective_slice_elements( + cfg: MlaConfig, + elements_per_slice: int, +) -> int: + """Cap one slice to the four row-statistics warps in a reducer CTA.""" + + _validate_parallel_reduction_slice(elements_per_slice) + return min( + elements_per_slice, + GMEM_REDUCTION_WARPS_PER_CTA * cfg.head_dim_per_cta_v, + ) + + +def parallel_gmem_reduction_launch_shape( + cfg: MlaConfig, + topology: ParallelReductionTopology, + elements_per_slice: int = PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_SLICE, +) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + """Return the parallel reducer grid and cluster shape. + + Grid X packs the cluster rank inside each output slice. Grid Y + packs effective Q and the V head-dimension CTA so the helper remains exact + about D-per-CTA rather than assuming total D is always 512. + """ + + elements_per_slice = _parallel_reduction_effective_slice_elements( + cfg, elements_per_slice + ) + output_slices = ceil_div( + cfg.num_heads_q * cfg.head_dim_per_cta_v, + elements_per_slice, + ) + return ( + ( + output_slices * topology.cluster_size, + cfg.seq_len_q * cfg.num_ctas_per_head_dim, + cfg.batch_size, + ), + (topology.cluster_size, 1, 1), + ) + + +def parallel_gmem_reduction_base_clusters( + cfg: MlaConfig, + elements_per_slice: int = PARALLEL_GMEM_REDUCTION_ELEMENTS_PER_SLICE, +) -> int: + """Return logical reducer clusters for the selected output slice.""" + + elements_per_slice = _parallel_reduction_effective_slice_elements( + cfg, elements_per_slice + ) + output_slices = ceil_div( + cfg.num_heads_q * cfg.head_dim_per_cta_v, + elements_per_slice, + ) + return output_slices * cfg.seq_len_q * cfg.num_ctas_per_head_dim * cfg.batch_size + + +def parallel_gmem_reduction_threads(elements_per_slice: int) -> int: + """Return reducer threads for an output slice.""" + + _validate_parallel_reduction_slice(elements_per_slice) + return PARALLEL_GMEM_REDUCTION_THREADS + + +def parallel_gmem_reduction_elements_per_thread(elements_per_slice: int) -> int: + """Return scalar output elements owned by one reducer thread. + + A 128-element band uses one scalar per thread. Wider slices assign BF16 + vec4 or vec8 fragments while retaining one shared-statistics warp per row. + """ + + return ceil_div( + elements_per_slice, + parallel_gmem_reduction_threads(elements_per_slice), + ) + + +@cute.jit +def _run_parallel_gmem_reduction_g1_shared_stats( + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + elements_per_slice: cutlass.Constexpr[int], +): + """Reduce G1 partials with the compact row-shared schedule.""" + + slice_idx, block_idx_y, batch_idx = cute.arch.block_idx() + thread_idx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_idx = thread_idx % Int32(GMEM_REDUCTION_WARP_LANES) + head_dim_cta_idx = block_idx_y % Int32(cfg.num_ctas_per_head_dim) + q_idx = block_idx_y // Int32(cfg.num_ctas_per_head_dim) + elements_per_slice = _parallel_reduction_effective_slice_elements( + cfg, elements_per_slice + ) + rows_per_slice = ceil_div(elements_per_slice, cfg.head_dim_per_cta_v) + reducer_threads = parallel_gmem_reduction_threads(elements_per_slice) + output_elements_per_thread = ceil_div(elements_per_slice, reducer_threads) + slice_element_base = slice_idx * Int32(elements_per_slice) + + smem_scale = cutlass.Array( + Float32, + rows_per_slice * cfg.num_ctas_per_seq_kv, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + # One statistics warp per covered row stripes over split slots, then + # publishes normalized scales for all output threads to reuse. + if warp_idx < Int32(rows_per_slice): + stats_row = warp_idx + stats_element = slice_element_base + stats_row * Int32(cfg.head_dim_per_cta_v) + head_idx = stats_element // Int32(cfg.head_dim_per_cta_v) + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + if head_idx < Int32(cfg.num_heads_q) and valid_output_row: + row_lse = acc_lse[head_idx, None, q_idx, batch_idx] + lse_per_lane = ceil_div( + cfg.num_ctas_per_seq_kv, + GMEM_REDUCTION_WARP_LANES, + ) + lane_lse = cutlass.Array( + Float32, + lse_per_lane, + space=cutlass.AddressSpace.rmem, + ) + lse_max = Float32(-Float32.inf) + for lane_slot_i in cutlass.range_constexpr(lse_per_lane): + split_idx = lane_idx + Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) + lane_lse[lane_slot_i] = ( + Float32(row_lse[split_idx]) + if split_idx < Int32(cfg.num_ctas_per_seq_kv) + else Float32(-Float32.inf) + ) + lse_max = fmax_f32(lse_max, lane_lse[lane_slot_i]) + + lse_max = warp_reduce_max_f32(lse_max) + lse_max = lse_max if lse_max != Float32(-Float32.inf) else Float32(0.0) + lse_sum = Float32(0.0) + for lane_slot_i in cutlass.range_constexpr(lse_per_lane): + lse_sum += cute.math.exp2( + lane_lse[lane_slot_i] - lse_max, + fastmath=True, + ) + lse_sum = warp_reduce_sum_f32(lse_sum) + global_lse = finalize_log2_sum_exp(lse_max, lse_sum) + if ( + lane_idx == Int32(0) + and head_dim_cta_idx == Int32(0) + and stats_element % Int32(cfg.head_dim_per_cta_v) == Int32(0) + ): + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + + for lane_slot_i in cutlass.range_constexpr(lse_per_lane): + split_idx = lane_idx + Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) + if split_idx < Int32(cfg.num_ctas_per_seq_kv): + smem_scale[ + stats_row * Int32(cfg.num_ctas_per_seq_kv) + split_idx + ] = normalized_lse_weight(lane_lse[lane_slot_i], global_lse) + + prims.barrier_cta_sync( + barrier_id=SPLIT_REDUCTION_SCALE_BARRIER_ID, + thread_count=reducer_threads, + ) + + thread_elem_offset = thread_idx * Int32(output_elements_per_thread) + flat_element = slice_element_base + thread_elem_offset + head_idx = flat_element // Int32(cfg.head_dim_per_cta_v) + dim_in_cta = flat_element - head_idx * Int32(cfg.head_dim_per_cta_v) + row_in_slice = (flat_element - slice_element_base) // Int32(cfg.head_dim_per_cta_v) + dim_idx = head_dim_cta_idx * Int32(cfg.head_dim_per_cta_v) + dim_in_cta + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + valid_output = ( + thread_elem_offset < Int32(elements_per_slice) + and head_idx < Int32(cfg.num_heads_q) + and dim_idx < Int32(cfg.head_dim_v) + and valid_output_row + ) + if valid_output: + output_acc = cutlass.Array( + Float32, + output_elements_per_thread, + space=cutlass.AddressSpace.rmem, + ) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + output_acc[elem_i] = Float32(0.0) + # Form the 64-bit row base once. Split strides are compile-time + # constants in this fixed-S kernel, avoiding repeated wide integer + # linearization in the hot accumulation loop. + acc_row_ptr = acc_output.iterator.raw_ptr() + _split_o_row_base_offset( + cfg, + batch_idx, + q_idx, + head_idx, + dim_idx, + ) + for split_i in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): + split_idx = Int32(split_i) + scale = Float32( + smem_scale[row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx] + ) + split_ptr = acc_row_ptr + Int64(split_i * cfg.head_dim_v) + if cutlass.const_expr(output_elements_per_thread == 1): + output_acc[0] += Float32(split_ptr.load()) * scale + else: + partial_o = split_ptr.load( + count=output_elements_per_thread, + alignment=output_elements_per_thread * cfg.partial_o_dtype_bytes, + ).to(Float32) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + output_acc[elem_i] += partial_o[elem_i] * scale + + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + out_elem_offset = _output_element_offset(cfg, output_query_row, dim_idx) + if cutlass.const_expr(output_elements_per_thread == 1): + (output.iterator.raw_ptr() + out_elem_offset).store( + output.element_type(output_acc[0]) + ) + else: + output_regs = cutlass.Array( + output.element_type, + output_elements_per_thread, + space=cutlass.AddressSpace.rmem, + ) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + output_regs[elem_i] = output.element_type(output_acc[elem_i]) + (output.iterator.raw_ptr() + out_elem_offset).store( + output_regs.data_ptr().load( + count=output_elements_per_thread, + alignment=output_elements_per_thread * cfg.o_dtype_bytes, + ), + alignment=output_elements_per_thread * cfg.o_dtype_bytes, + ) + + +@cute.jit +def _run_parallel_gmem_reduction_shared_stats( + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + cluster_size: cutlass.Constexpr[int], + slots_per_rank: cutlass.Constexpr[int], + actual_splits: cutlass.Constexpr[int], + elements_per_slice: cutlass.Constexpr[int], +): + """Reduce split-KV with one cooperative statistics warp per output row. + + Every cluster rank owns at most four D-per-CTA rows and a cyclic subset of + split slots. One warp calculates each row's statistics, then all output + threads reuse the normalized scales while reducing BF16 fragments. Rank + zero repeats the same row-shared scheme for the final DSMEM merge. + """ + + block_idx_x, block_idx_y, batch_idx = cute.arch.block_idx() + thread_idx, _, _ = cute.arch.thread_idx() + cluster_rank = cute.arch.block_idx_in_cluster() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_idx = thread_idx % Int32(GMEM_REDUCTION_WARP_LANES) + + slice_idx = block_idx_x // Int32(cluster_size) + head_dim_cta_idx = block_idx_y % Int32(cfg.num_ctas_per_head_dim) + q_idx = block_idx_y // Int32(cfg.num_ctas_per_head_dim) + local_slots = slots_per_rank + elements_per_slice = _parallel_reduction_effective_slice_elements( + cfg, elements_per_slice + ) + rows_per_slice = ceil_div(elements_per_slice, cfg.head_dim_per_cta_v) + reducer_threads = parallel_gmem_reduction_threads(elements_per_slice) + output_elements_per_thread = ceil_div(elements_per_slice, reducer_threads) + slice_element_base = slice_idx * Int32(elements_per_slice) + + neg_inf = Float32(-Float32.inf) + smem_local_scale = cutlass.Array( + Float32, + rows_per_slice * local_slots, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + smem_local_lse = cutlass.Array( + Float32, + rows_per_slice, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + # Statistics warp r owns row r. This ownership is independent of vector + # ownership; a vector warp can consume a different row fragment. + if warp_idx < Int32(rows_per_slice): + stats_row = warp_idx + stats_element = slice_element_base + stats_row * Int32(cfg.head_dim_per_cta_v) + stats_head_idx = stats_element // Int32(cfg.head_dim_per_cta_v) + ( + _, + _, + _, + _, + valid_stats_row, + ) = flat_query_row_state( + stats_head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + valid_stats = stats_head_idx < Int32(cfg.num_heads_q) and valid_stats_row + lse_per_lane = ceil_div(local_slots, GMEM_REDUCTION_WARP_LANES) + lane_lse = cutlass.Array( + Float32, + lse_per_lane, + space=cutlass.AddressSpace.rmem, + ) + lane_exp = cutlass.Array( + Float32, + lse_per_lane, + space=cutlass.AddressSpace.rmem, + ) + for lane_slot_i in cutlass.range_constexpr(lse_per_lane): + lane_lse[lane_slot_i] = neg_inf + local_lse = neg_inf + if valid_stats: + # Form the row view only after validating the flat output row. + # Invalid/padded rows publish a neutral state without an acc_lse + # address or load. + row_lse = acc_lse[stats_head_idx, None, q_idx, batch_idx] + local_max = neg_inf + for lane_slot_i in cutlass.range_constexpr(lse_per_lane): + local_slot_idx = lane_idx + Int32( + lane_slot_i * GMEM_REDUCTION_WARP_LANES + ) + split_idx = local_slot_idx * Int32(cluster_size) + cluster_rank + if local_slot_idx < Int32(local_slots) and split_idx < Int32( + actual_splits + ): + lane_lse[lane_slot_i] = Float32(row_lse[split_idx]) + local_max = fmax_f32(local_max, lane_lse[lane_slot_i]) + + local_max = warp_reduce_max_f32(local_max) + safe_local_max = local_max if local_max != neg_inf else Float32(0.0) + local_sum = Float32(0.0) + for lane_slot_i in cutlass.range_constexpr(lse_per_lane): + lane_exp[lane_slot_i] = Float32( + cute.math.exp2( + lane_lse[lane_slot_i] - safe_local_max, + fastmath=True, + ) + ) + local_sum += lane_exp[lane_slot_i] + local_sum = warp_reduce_sum_f32(local_sum) + local_lse = finalize_log2_sum_exp(safe_local_max, local_sum) + + if lane_idx == Int32(0): + smem_local_lse[stats_row] = local_lse + + for lane_slot_i in cutlass.range_constexpr(lse_per_lane): + local_slot_idx = lane_idx + Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) + if local_slot_idx < Int32(local_slots): + smem_local_scale[stats_row * Int32(local_slots) + local_slot_idx] = ( + normalized_lse_weight(lane_lse[lane_slot_i], local_lse) + ) + + prims.barrier_cta_sync( + barrier_id=SPLIT_REDUCTION_SCALE_BARRIER_ID, + thread_count=reducer_threads, + ) + + thread_elem_offset = thread_idx * Int32(output_elements_per_thread) + flat_element = slice_element_base + thread_elem_offset + head_idx = flat_element // Int32(cfg.head_dim_per_cta_v) + dim_in_cta = flat_element - head_idx * Int32(cfg.head_dim_per_cta_v) + row_in_slice = (flat_element - slice_element_base) // Int32(cfg.head_dim_per_cta_v) + dim_idx = head_dim_cta_idx * Int32(cfg.head_dim_per_cta_v) + dim_in_cta + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + valid_output = ( + thread_elem_offset < Int32(elements_per_slice) + and head_idx < Int32(cfg.num_heads_q) + and dim_idx < Int32(cfg.head_dim_v) + and valid_output_row + ) + local_output = cutlass.Array( + Float32, + output_elements_per_thread, + space=cutlass.AddressSpace.rmem, + ) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + local_output[elem_i] = Float32(0.0) + if valid_output: + # Offset to this rank once; each unrolled local slot then advances by + # the compile-time cluster stride. This retains 64-bit safety without + # paying a full wide linearization for every partial load. + acc_row_ptr = ( + acc_output.iterator.raw_ptr() + + _split_o_row_base_offset( + cfg, + batch_idx, + q_idx, + head_idx, + dim_idx, + ) + + Int64(cluster_rank) * Int64(cfg.head_dim_v) + ) + for local_slot_i in cutlass.range_constexpr(local_slots): + split_idx = Int32(local_slot_i * cluster_size) + cluster_rank + if split_idx < Int32(actual_splits): + scale = Float32( + smem_local_scale[ + row_in_slice * Int32(local_slots) + Int32(local_slot_i) + ] + ) + split_ptr = acc_row_ptr + Int64( + local_slot_i * cluster_size * cfg.head_dim_v + ) + if cutlass.const_expr(output_elements_per_thread == 1): + local_output[0] += Float32(split_ptr.load()) * scale + else: + partial_o = split_ptr.load( + count=output_elements_per_thread, + alignment=output_elements_per_thread + * cfg.partial_o_dtype_bytes, + ).to(Float32) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + local_output[elem_i] += partial_o[elem_i] * scale + + # Publish normalized O in BF16 while retaining FP32 accumulation. + smem_o = cutlass.Array( + acc_output.element_type, + reducer_threads * output_elements_per_thread, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + smem_o_offset = thread_idx * Int32(output_elements_per_thread) + if cutlass.const_expr(output_elements_per_thread == 1): + smem_o[smem_o_offset] = acc_output.element_type(local_output[0]) + else: + smem_output_regs = cutlass.Array( + acc_output.element_type, + output_elements_per_thread, + space=cutlass.AddressSpace.rmem, + ) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + smem_output_regs[elem_i] = acc_output.element_type(local_output[elem_i]) + (smem_o.data_ptr() + smem_o_offset).store( + smem_output_regs.data_ptr().load( + count=output_elements_per_thread, + alignment=output_elements_per_thread * cfg.partial_o_dtype_bytes, + ), + alignment=output_elements_per_thread * cfg.partial_o_dtype_bytes, + ) + + prims.barrier_cta_sync(0) + prims.barrier_cluster_arrive() + prims.barrier_cluster_wait() + + smem_global_lse = cutlass.Array( + Float32, + rows_per_slice, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + smem_peer_scale = cutlass.Array( + Float32, + rows_per_slice * cluster_size, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + if cluster_rank == Int32(0): + if warp_idx < Int32(rows_per_slice): + stats_row = warp_idx + stats_element = slice_element_base + stats_row * Int32( + cfg.head_dim_per_cta_v + ) + stats_head_idx = stats_element // Int32(cfg.head_dim_per_cta_v) + ( + _, + _, + _, + _, + valid_stats_row, + ) = flat_query_row_state( + stats_head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + valid_stats = stats_head_idx < Int32(cfg.num_heads_q) and valid_stats_row + peer_lse = neg_inf + for peer_rank_i in cutlass.range_constexpr(cluster_size): + if valid_stats and lane_idx == Int32(peer_rank_i): + peer_lse = ( + prims.mapa( + smem_local_lse.data_ptr(), + Int32(peer_rank_i), + ) + + stats_row + ).load() + + global_max = warp_reduce_max_f32(peer_lse) + safe_global_max = global_max if global_max != neg_inf else Float32(0.0) + peer_exp = Float32( + cute.math.exp2( + peer_lse - safe_global_max, + fastmath=True, + ) + ) + global_sum = warp_reduce_sum_f32(peer_exp) + global_lse = finalize_log2_sum_exp(safe_global_max, global_sum) + if lane_idx == Int32(0): + smem_global_lse[stats_row] = global_lse + if lane_idx < Int32(cluster_size): + smem_peer_scale[stats_row * Int32(cluster_size) + lane_idx] = ( + normalized_lse_weight(peer_lse, global_lse) + ) + + # Only rank zero executes this CTA-uniform branch and barrier. + prims.barrier_cta_sync(0) + + global_output = cutlass.Array( + Float32, + output_elements_per_thread, + space=cutlass.AddressSpace.rmem, + ) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + global_output[elem_i] = Float32(0.0) + if valid_output: + for peer_rank_i in cutlass.range_constexpr(cluster_size): + peer_scale = Float32( + smem_peer_scale[ + row_in_slice * Int32(cluster_size) + Int32(peer_rank_i) + ] + ) + peer_o = prims.mapa( + smem_o.data_ptr(), + Int32(peer_rank_i), + ) + if cutlass.const_expr(output_elements_per_thread == 1): + global_output[0] += ( + Float32((peer_o + smem_o_offset).load()) * peer_scale + ) + else: + peer_output = ( + (peer_o + smem_o_offset) + .load( + count=output_elements_per_thread, + alignment=( + output_elements_per_thread * cfg.partial_o_dtype_bytes + ), + ) + .to(Float32) + ) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + global_output[elem_i] += peer_output[elem_i] * peer_scale + + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + if head_dim_cta_idx == Int32(0) and dim_in_cta == Int32(0): + (lse.iterator.raw_ptr() + output_query_row).store( + smem_global_lse[row_in_slice] + ) + out_elem_offset = _output_element_offset(cfg, output_query_row, dim_idx) + if cutlass.const_expr(output_elements_per_thread == 1): + (output.iterator.raw_ptr() + out_elem_offset).store( + output.element_type(global_output[0]) + ) + else: + output_regs = cutlass.Array( + output.element_type, + output_elements_per_thread, + space=cutlass.AddressSpace.rmem, + ) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + output_regs[elem_i] = output.element_type(global_output[elem_i]) + (output.iterator.raw_ptr() + out_elem_offset).store( + output_regs.data_ptr().load( + count=output_elements_per_thread, + alignment=output_elements_per_thread * cfg.o_dtype_bytes, + ), + alignment=output_elements_per_thread * cfg.o_dtype_bytes, + ) + + # Every peer remains alive until rank zero finishes all DSMEM O loads. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + + +@cute.jit +def run_parallel_gmem_reduction_kernel( + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + cluster_size: cutlass.Constexpr[int], + slots_per_rank: cutlass.Constexpr[int], + actual_splits: cutlass.Constexpr[int], + elements_per_slice: cutlass.Constexpr[int], +): + """Reduce split-KV partials with row-shared local and peer statistics.""" + + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + + if cutlass.const_expr(cluster_size == 1): + _run_parallel_gmem_reduction_g1_shared_stats( + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + elements_per_slice, + ) + return + + _run_parallel_gmem_reduction_shared_stats( + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + cluster_size, + slots_per_rank, + actual_splits, + elements_per_slice, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py new file mode 100644 index 000000000000..c5a694dbdba1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py @@ -0,0 +1,851 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Split-KV GMEM reduction helpers for throughput-latency 1CTA MLA.""" + +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims +from cutlass import Float32, Int32, Int64 + +from .config import MlaConfig +from ..helpers.math import ceil_div +from ..helpers.ops import ( + fmax_f32, + warp_reduce_max_f32, + warp_reduce_sum_f32, + vector_from_scalars, +) +from ..helpers.query import ( + flat_query_row_state, + public_query_flat_row, + split_o_element_offset, +) +from ..helpers.tile import runtime_seq_len_kv_for_q + + +# Baseline split-KV reduction uses one 128-thread CTA. Each thread owns eight +# BF16 elements, which is one 16B vector and covers a full 512-wide output head +# per CTA. +GMEM_REDUCTION_THREADS = 128 +GMEM_REDUCTION_WARP_LANES = 32 +GMEM_REDUCTION_WARPS_PER_CTA = GMEM_REDUCTION_THREADS // GMEM_REDUCTION_WARP_LANES +GMEM_REDUCTION_ELEMENTS_PER_THREAD = 8 +GMEM_REDUCTION_VECTOR_BYTES = 16 +GMEM_REDUCTION_VECTOR_HEAD_DIM = 512 + +# Keeps-MMA-AB split reduction uses a wider CTA so one slice can cover 512 16B +# vectors and split a large Q/head tile across multiple reducer CTAs. +GMEM_REDUCTION_SLICE_THREADS = 512 + +# Scalar stats are reduced in 128-wide bands to match the warp-reduction and +# SMEM exchange layout used by the O reduction. +GMEM_REDUCTION_SCALAR_DIM_TILE = 128 + +# The small-B fast path lets one reducer CTA cover two 512-wide BF16 heads. +GMEM_REDUCTION_HEADS_PER_CTA = 2 + +# CTA-local barrier used after writer warps publish split-KV softmax scales in +# SMEM and before all reducer threads consume those scales for O accumulation. +GMEM_REDUCTION_SCALE_BARRIER_ID = 4 +GMEM_REDUCTION_HEAD_SEGMENTS_PER_CTA = ( + GMEM_REDUCTION_HEADS_PER_CTA + * GMEM_REDUCTION_VECTOR_HEAD_DIM + // (GMEM_REDUCTION_THREADS * GMEM_REDUCTION_ELEMENTS_PER_THREAD) +) +GMEM_REDUCTION_Q_ROWS_PER_CTA = ( + GMEM_REDUCTION_THREADS + * GMEM_REDUCTION_ELEMENTS_PER_THREAD + // GMEM_REDUCTION_VECTOR_HEAD_DIM +) + + +def uses_two_head_gmem_reduction(cfg: MlaConfig) -> bool: + """Return whether one reducer CTA covers two full BF16 output heads.""" + return ( + cfg.head_dim_v == GMEM_REDUCTION_VECTOR_HEAD_DIM + and cfg.o_dtype_bytes == 2 + and cfg.batch_size <= 2 + ) + + +def uses_q_row_gmem_reduction(cfg: MlaConfig) -> bool: + """Return whether one reducer CTA covers multiple Q rows for one head.""" + return ( + cfg.seq_len_q >= 2 * GMEM_REDUCTION_Q_ROWS_PER_CTA + and cfg.seq_len_q % GMEM_REDUCTION_Q_ROWS_PER_CTA == 0 + and cfg.head_dim_v == GMEM_REDUCTION_VECTOR_HEAD_DIM + and cfg.o_dtype_bytes == 2 + ) + + +def uses_slice_split_gmem_reduction(cfg: MlaConfig) -> bool: + """Return whether the slice-split reducer applies. + + The keeps-MMA-AB split-KV path writes one Q/head tile of BF16 partial O per + split. A 512-thread reducer CTA covers a contiguous slice of rows from that + tile, with each thread loading/storing one 16B vector. Multiple reducer CTAs + can split the row slices for one Q/head tile when there are spare SMs. + """ + + return ( + cfg.kernel_variant == "keeps_mma_ab" + and cfg.tile_size_q in (64, 128) + and cfg.head_dim_per_cta_v in (64, 128, 256, 512) + and cfg.partial_o_dtype_bytes == 2 + ) + + +def slice_split_rows_per_slice(cfg: MlaConfig) -> int: + """Return rows covered by one 512-thread reducer slice.""" + + return ( + GMEM_REDUCTION_SLICE_THREADS + * GMEM_REDUCTION_ELEMENTS_PER_THREAD + // cfg.head_dim_per_cta_v + ) + + +def slice_split_num_slices(cfg: MlaConfig) -> int: + """Return row-slice count in one keeps-MMA-AB Q/head tile.""" + + return ceil_div(cfg.tile_size_q, slice_split_rows_per_slice(cfg)) + + +def slice_split_num_reduction_ctas( + cfg: MlaConfig, + seq_len_q, + batch_size, + max_active_clusters, +) -> int: + """Return reducer CTAs per Q/head tile, capped to roughly two SM waves.""" + + base_ctas = ( + int(seq_len_q) + * cfg.num_ctas_for_all_heads + * cfg.num_ctas_per_head_dim + * int(batch_size) + ) + if base_ctas <= 0: + return 1 + max_ctas_for_reduction = max(1, (int(max_active_clusters) * 2) // base_ctas) + return min(max_ctas_for_reduction, slice_split_num_slices(cfg)) + + +def gmem_reduction_launch_shape( + cfg: MlaConfig, + seq_len_q, + batch_size, + lse_width, + max_active_clusters, +): + """Return grid, dynamic SMEM bytes, block threads, and CTAs per tile.""" + + if uses_slice_split_gmem_reduction(cfg): + num_reduction_ctas = slice_split_num_reduction_ctas( + cfg, + seq_len_q, + batch_size, + max_active_clusters, + ) + return ( + ( + int(seq_len_q) * num_reduction_ctas, + cfg.num_ctas_for_all_heads * cfg.num_ctas_per_head_dim, + int(batch_size), + ), + 0, + GMEM_REDUCTION_SLICE_THREADS, + num_reduction_ctas, + ) + + lse_bytes = lse_width // 8 + if uses_two_head_gmem_reduction(cfg): + return ( + ( + seq_len_q, + ceil_div(cfg.num_heads_q, GMEM_REDUCTION_HEADS_PER_CTA) + * GMEM_REDUCTION_HEAD_SEGMENTS_PER_CTA, + batch_size, + ), + GMEM_REDUCTION_HEADS_PER_CTA * cfg.num_ctas_per_seq_kv * lse_bytes, + GMEM_REDUCTION_THREADS, + 1, + ) + if uses_q_row_gmem_reduction(cfg): + return ( + ( + ceil_div(seq_len_q, GMEM_REDUCTION_Q_ROWS_PER_CTA), + cfg.num_heads_q, + batch_size, + ), + GMEM_REDUCTION_Q_ROWS_PER_CTA * cfg.num_ctas_per_seq_kv * lse_bytes, + GMEM_REDUCTION_THREADS, + 1, + ) + return ( + ( + seq_len_q, + cfg.num_heads_q * ceil_div(cfg.head_dim_v, GMEM_REDUCTION_SCALAR_DIM_TILE), + batch_size, + ), + cfg.num_ctas_per_seq_kv * lse_bytes, + GMEM_REDUCTION_THREADS, + 1, + ) + + +@cute.jit +def runtime_seq_len_kv_for_reduction( + cfg: MlaConfig, + cache_seqs, + batch_idx, + cta_idx_q, + cu_seqlens_q=None, +): + """Return the runtime KV length visible to split-KV reduction.""" + return runtime_seq_len_kv_for_q( + cfg, + cache_seqs, + batch_idx, + cta_idx_q, + cu_seqlens_q, + ) + + +@cute.jit +def run_gmem_reduction_kernel( + kernel, + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + num_reduction_ctas, +): + """Reduce split-KV partial O/LSE rows written by the 1CTA main kernel. + + Each reducer CTA owns either a full head row, a small group of Q rows, or a + slice of a keeps-MMA-AB head tile. For its owned rows it loads all split-KV + partial LSE values, computes the row max, rescales each partial O by + ``exp2(partial_lse - max_lse)``, sums the weighted O vectors and + denominators, then writes final O and LSE. Runtime-pruned producers publish + neutral partials, so the reducer retains configured, compile-time split + geometry without consuming stale workspace from an earlier graph replay. + + The grid and workspace remain in physical flat-query tile coordinates. + Every reducer variant maps those coordinates to fixed or compact-ragged + public storage; inactive tail rows may synchronize and consume workspace + slots but never publish O or LSE. + """ + # Pair with the dense kernel PDL signal before reading partials. + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + q_idx, head_dim_tile_idx, batch_idx = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_idx = tidx % Int32(GMEM_REDUCTION_WARP_LANES) + + use_two_head_reduction = uses_two_head_gmem_reduction(cfg) + use_q_row_reduction = uses_q_row_gmem_reduction(cfg) + use_slice_split_reduction = uses_slice_split_gmem_reduction(cfg) + if cutlass.const_expr(use_slice_split_reduction): + cta_idx_q = q_idx % Int32(cfg.seq_len_q) + cta_idx_for_reduction = q_idx // Int32(cfg.seq_len_q) + head_dim_cta_idx = head_dim_tile_idx % Int32(cfg.num_ctas_per_head_dim) + head_group_idx = head_dim_tile_idx // Int32(cfg.num_ctas_per_head_dim) + head_base_idx = head_group_idx * Int32(cfg.tile_size_q) + head_dim_offset = head_dim_cta_idx * Int32(cfg.head_dim_per_cta_v) + rows_per_slice = slice_split_rows_per_slice(cfg) + num_slices = slice_split_num_slices(cfg) + num_slices_per_cta = ceil_div(Int32(num_slices), num_reduction_ctas) + start_slice_idx = cta_idx_for_reduction * num_slices_per_cta + end_slice_idx = cute.math.min( + start_slice_idx + num_slices_per_cta, + Int32(num_slices), + ) + acc_ptr = acc_output.iterator.raw_ptr() + out_ptr = output.iterator.raw_ptr() + for slice_idx in range(start_slice_idx, end_slice_idx): + base_vec_offset = tidx * Int32(GMEM_REDUCTION_ELEMENTS_PER_THREAD) + row_in_slice = base_vec_offset // Int32(cfg.head_dim_per_cta_v) + dim_in_cta = base_vec_offset - row_in_slice * Int32(cfg.head_dim_per_cta_v) + row_in_tile = slice_idx * Int32(rows_per_slice) + row_in_slice + head_idx = head_base_idx + row_in_tile + dim_idx = head_dim_offset + dim_in_cta + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + head_idx, + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + valid_row = head_idx < Int32(cfg.num_heads_q) and valid_output_row + valid_dim = dim_idx < Int32(cfg.head_dim_v) + + if valid_row and valid_dim: + row_lse = acc_lse[head_idx, None, cta_idx_q, batch_idx] + local_lse = cutlass.Array( + kernel.lse_dtype, + cfg.num_ctas_per_seq_kv, + space=cutlass.AddressSpace.rmem, + ) + lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) + for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): + local_lse[split_idx] = row_lse[split_idx] + lse_max = fmax_f32(lse_max, local_lse[split_idx]) + + lse_max = ( + lse_max + if lse_max != -kernel.lse_dtype.inf + else kernel.lse_dtype(0.0) + ) + sum_lse = kernel.lse_dtype(0.0) + for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): + sum_lse += cute.math.exp2( + local_lse[split_idx] - lse_max, + fastmath=True, + ) + has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype( + 0.0 + ) + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if has_finite_mass + else -kernel.lse_dtype.inf + ) + if dim_in_cta == Int32(0): + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + + acc_vec = vector_from_scalars( + ( + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + ), + dtype=Float32, + ) + for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): + scale = Float32( + cute.math.exp2( + local_lse[split_idx] - global_lse, + fastmath=True, + ) + if has_finite_mass + else kernel.acc_dtype(0.0) + ) + acc_elem_offset = split_o_element_offset( + cfg, + batch_idx, + cta_idx_q, + head_idx, + Int32(split_idx), + dim_idx, + ) + partial_vec = ( + (acc_ptr + acc_elem_offset) + .load( + count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, + alignment=GMEM_REDUCTION_VECTOR_BYTES, + ) + .to(Float32) + ) + acc_vec = acc_vec + partial_vec * scale + + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + out_elem_offset = Int64(output_query_row) * Int64( + cfg.head_dim_v + ) + Int64(dim_idx) + (out_ptr + out_elem_offset).store( + acc_vec.to(output.element_type), + alignment=GMEM_REDUCTION_VECTOR_BYTES, + ) + return + + if cutlass.const_expr(use_two_head_reduction): + # Small-batch BF16 path: one CTA reduces two full output heads and each + # thread writes 8 elements, matching a 16B GMEM vector store. + head_group_idx = head_dim_tile_idx // Int32( + GMEM_REDUCTION_HEAD_SEGMENTS_PER_CTA + ) + segment_idx = head_dim_tile_idx - head_group_idx * Int32( + GMEM_REDUCTION_HEAD_SEGMENTS_PER_CTA + ) + smem_lse_scale = cutlass.Array( + kernel.lse_dtype, + GMEM_REDUCTION_HEADS_PER_CTA * cfg.num_ctas_per_seq_kv, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + for row_group in cutlass.range_constexpr( + ceil_div(GMEM_REDUCTION_HEADS_PER_CTA, GMEM_REDUCTION_WARPS_PER_CTA) + ): + local_head_idx = Int32(row_group * GMEM_REDUCTION_WARPS_PER_CTA) + warp_idx + global_head_idx = ( + head_group_idx * Int32(GMEM_REDUCTION_HEADS_PER_CTA) + local_head_idx + ) + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + global_head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + if ( + local_head_idx < Int32(GMEM_REDUCTION_HEADS_PER_CTA) + and global_head_idx < Int32(cfg.num_heads_q) + and valid_output_row + ): + row_lse = acc_lse[global_head_idx, None, q_idx, batch_idx] + # Warp lanes stripe over split-KV LSE slots; each lane owns + # split_idx = lane + n * warp_size. + lse_per_thread = ceil_div( + cfg.num_ctas_per_seq_kv, GMEM_REDUCTION_WARP_LANES + ) + local_lse = cutlass.Array( + kernel.lse_dtype, + lse_per_thread, + space=cutlass.AddressSpace.rmem, + ) + lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) + for i in cutlass.range_constexpr(lse_per_thread): + split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) + local_lse[i] = ( + row_lse[split_idx] + if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) + else -kernel.lse_dtype.inf + ) + lse_max = fmax_f32(lse_max, local_lse[i]) + lse_max = warp_reduce_max_f32(lse_max) + lse_max = ( + lse_max + if lse_max != -kernel.lse_dtype.inf + else kernel.lse_dtype(0.0) + ) + sum_lse = kernel.lse_dtype(0.0) + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = warp_reduce_sum_f32(sum_lse) + has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype( + 0.0 + ) + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if has_finite_mass + else -kernel.lse_dtype.inf + ) + if lane_idx == Int32(0): + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + for i in cutlass.range_constexpr(lse_per_thread): + split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) + if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): + smem_lse_scale[ + local_head_idx * Int32(cfg.num_ctas_per_seq_kv) + split_idx + ] = ( + cute.math.exp2(local_lse[i] - global_lse, fastmath=True) + if has_finite_mass + else kernel.acc_dtype(0.0) + ) + + prims.barrier_cta_sync( + barrier_id=GMEM_REDUCTION_SCALE_BARRIER_ID, + thread_count=GMEM_REDUCTION_THREADS, + ) + + acc_ptr = acc_output.iterator.raw_ptr() + vecs_per_head = cfg.head_dim_v // GMEM_REDUCTION_ELEMENTS_PER_THREAD + base_vec_idx = segment_idx * Int32(GMEM_REDUCTION_THREADS) + tidx + local_head_idx = base_vec_idx // Int32(vecs_per_head) + dim_vec_idx = base_vec_idx - local_head_idx * Int32(vecs_per_head) + dim_idx = dim_vec_idx * Int32(GMEM_REDUCTION_ELEMENTS_PER_THREAD) + global_head_idx = ( + head_group_idx * Int32(GMEM_REDUCTION_HEADS_PER_CTA) + local_head_idx + ) + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + global_head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + if global_head_idx < Int32(cfg.num_heads_q) and valid_output_row: + acc_vec = vector_from_scalars( + ( + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + ), + dtype=Float32, + ) + for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): + scale = Float32( + smem_lse_scale[ + local_head_idx * Int32(cfg.num_ctas_per_seq_kv) + + Int32(split_idx) + ] + ) + acc_elem_offset = split_o_element_offset( + cfg, + batch_idx, + q_idx, + global_head_idx, + Int32(split_idx), + dim_idx, + ) + partial_vec = ( + (acc_ptr + acc_elem_offset) + .load( + count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, + alignment=16, + ) + .to(Float32) + ) + acc_vec = acc_vec + partial_vec * scale + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + out_elem_offset = Int64(output_query_row) * Int64(cfg.head_dim_v) + Int64( + dim_idx + ) + (output.iterator.raw_ptr() + out_elem_offset).store( + acc_vec.to(output.element_type), + alignment=16, + ) + return + + if cutlass.const_expr(use_q_row_reduction): + # Multi-token-Q BF16 path: one CTA reduces several Q rows for one head, + # again using 8 output elements per thread. + slice_idx = q_idx + head_idx = head_dim_tile_idx + rows_per_slice = GMEM_REDUCTION_Q_ROWS_PER_CTA + # Warp lanes stripe over split-KV LSE slots for each Q row in the slice. + lse_per_thread = ceil_div(cfg.num_ctas_per_seq_kv, GMEM_REDUCTION_WARP_LANES) + smem_lse_scale = cutlass.Array( + kernel.lse_dtype, + rows_per_slice * cfg.num_ctas_per_seq_kv, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + if warp_idx < Int32(rows_per_slice): + row_q_idx = slice_idx * Int32(rows_per_slice) + warp_idx + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + head_idx, + row_q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + if valid_output_row: + row_lse = acc_lse[head_idx, None, row_q_idx, batch_idx] + local_lse = cutlass.Array( + kernel.lse_dtype, + lse_per_thread, + space=cutlass.AddressSpace.rmem, + ) + lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) + for i in cutlass.range_constexpr(lse_per_thread): + split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) + local_lse[i] = ( + row_lse[split_idx] + if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) + else -kernel.lse_dtype.inf + ) + lse_max = fmax_f32(lse_max, local_lse[i]) + lse_max = warp_reduce_max_f32(lse_max) + lse_max = ( + lse_max + if lse_max != -kernel.lse_dtype.inf + else kernel.lse_dtype(0.0) + ) + sum_lse = kernel.lse_dtype(0.0) + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = warp_reduce_sum_f32(sum_lse) + has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype( + 0.0 + ) + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if has_finite_mass + else -kernel.lse_dtype.inf + ) + if lane_idx == Int32(0): + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + for i in cutlass.range_constexpr(lse_per_thread): + split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) + if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): + smem_lse_scale[ + warp_idx * Int32(cfg.num_ctas_per_seq_kv) + split_idx + ] = ( + cute.math.exp2(local_lse[i] - global_lse, fastmath=True) + if has_finite_mass + else kernel.acc_dtype(0.0) + ) + + prims.barrier_cta_sync( + barrier_id=GMEM_REDUCTION_SCALE_BARRIER_ID, + thread_count=GMEM_REDUCTION_THREADS, + ) + + base_vec_offset = tidx * Int32(GMEM_REDUCTION_ELEMENTS_PER_THREAD) + row_in_slice = base_vec_offset // Int32(cfg.head_dim_v) + dim_idx = base_vec_offset - row_in_slice * Int32(cfg.head_dim_v) + q_idx = slice_idx * Int32(rows_per_slice) + row_in_slice + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + if valid_output_row: + acc_vec = vector_from_scalars( + ( + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + ), + dtype=Float32, + ) + acc_ptr = acc_output.iterator.raw_ptr() + for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): + scale = Float32( + smem_lse_scale[ + row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx + ] + ) + acc_elem_offset = split_o_element_offset( + cfg, + batch_idx, + q_idx, + head_idx, + Int32(split_idx), + dim_idx, + ) + partial_vec = ( + (acc_ptr + acc_elem_offset) + .load( + count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, + alignment=16, + ) + .to(Float32) + ) + acc_vec = acc_vec + partial_vec * scale + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + out_elem_offset = Int64(output_query_row) * Int64(cfg.head_dim_v) + Int64( + dim_idx + ) + (output.iterator.raw_ptr() + out_elem_offset).store( + acc_vec.to(output.element_type), + alignment=16, + ) + return + + # Scalar fallback: one reducer coordinate owns one flat query row and one + # head-dimension band, then publishes only after logical-row validation. + reduction_dim_tiles = ceil_div(cfg.head_dim_v, GMEM_REDUCTION_SCALAR_DIM_TILE) + head_idx = head_dim_tile_idx // Int32(reduction_dim_tiles) + dim_tile_idx = head_dim_tile_idx - head_idx * Int32(reduction_dim_tiles) + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + head_idx, + q_idx, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=cu_seqlens_q, + batch_idx=batch_idx, + ) + + smem_lse_scale = cutlass.Array( + kernel.lse_dtype, + cfg.num_ctas_per_seq_kv, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + g_lse = acc_lse[head_idx, None, q_idx, batch_idx] + if warp_idx == Int32(0): + # One warp reduces split-KV statistics; lanes stripe split slots by + # warp size before publishing per-split O rescale factors. + lse_per_thread = ceil_div(cfg.num_ctas_per_seq_kv, GMEM_REDUCTION_WARP_LANES) + local_lse = cutlass.Array( + kernel.lse_dtype, + lse_per_thread, + space=cutlass.AddressSpace.rmem, + ) + lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) + for i in cutlass.range_constexpr(lse_per_thread): + split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) + local_lse[i] = ( + g_lse[split_idx] + if valid_output_row + and cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) + else -kernel.lse_dtype.inf + ) + lse_max = fmax_f32(lse_max, local_lse[i]) + lse_max = warp_reduce_max_f32(lse_max) + lse_max = lse_max if lse_max != -kernel.lse_dtype.inf else kernel.lse_dtype(0.0) + sum_lse = kernel.lse_dtype(0.0) + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = warp_reduce_sum_f32(sum_lse) + has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype(0.0) + global_lse = ( + lse_max + cute.math.log2(sum_lse, fastmath=True) + if has_finite_mass + else -kernel.lse_dtype.inf + ) + if lane_idx == Int32(0) and valid_output_row: + if dim_tile_idx == Int32(0): + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + for i in cutlass.range_constexpr(lse_per_thread): + split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) + if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): + smem_lse_scale[split_idx] = ( + cute.math.exp2(local_lse[i] - global_lse, fastmath=True) + if has_finite_mass + else kernel.acc_dtype(0.0) + ) + + prims.barrier_cta_sync( + barrier_id=GMEM_REDUCTION_SCALE_BARRIER_ID, + thread_count=GMEM_REDUCTION_THREADS, + ) + + dim_idx = dim_tile_idx * Int32(GMEM_REDUCTION_SCALAR_DIM_TILE) + tidx + g_acc_o = acc_output[head_idx, None, None, q_idx, batch_idx] + r_acc_o = cutlass.Array(kernel.acc_dtype, 1, space=cutlass.AddressSpace.rmem) + out_element_dtype = output.element_type + r_o = cutlass.Array(out_element_dtype, 1, space=cutlass.AddressSpace.rmem) + r_acc_o[0] = kernel.acc_dtype(0.0) + if valid_output_row: + for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): + scale = Float32(smem_lse_scale[split_idx]) + if dim_idx < Int32(cfg.head_dim_v): + r_acc_o[0] = r_acc_o[0] + Float32(g_acc_o[split_idx, dim_idx]) * scale + r_o.store(r_acc_o.load(0, 1).to(out_element_dtype), 0) + if dim_idx < Int32(cfg.head_dim_v) and valid_output_row: + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + cu_seqlens_q, + ) + ( + output.iterator.raw_ptr() + + Int64(output_query_row) * Int64(cfg.head_dim_v) + + Int64(dim_idx) + ).store(r_o[0]) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/__init__.py new file mode 100644 index 000000000000..13a688d9c25a --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/__init__.py @@ -0,0 +1,51 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resources for the throughput-latency 1CTA MLA task schedule.""" + +from .common import MlaResource, ScheduleTokenThrottleResource +from .smem_p import SmemPResource +from .smem_resources import ( + SmemKResource, + SmemKvResource, + SmemPageOffsetsResource, + SmemQResource, + SmemVResource, +) +from .tmem_corr import TmemCorrResource +from .tmem_o import TmemOResource +from .tmem_p import TmemPResource +from .tmem_s import TmemSKeepsResource, TmemSResource +from .tmem_softmax_stats import ( + TmemSoftmaxGlobalResource, + TmemSoftmaxLocalResource, +) + +__all__ = [ + "MlaResource", + "ScheduleTokenThrottleResource", + "SmemKResource", + "SmemKvResource", + "SmemPResource", + "SmemPageOffsetsResource", + "SmemQResource", + "SmemVResource", + "TmemCorrResource", + "TmemOResource", + "TmemPResource", + "TmemSKeepsResource", + "TmemSResource", + "TmemSoftmaxGlobalResource", + "TmemSoftmaxLocalResource", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/common.py new file mode 100644 index 000000000000..b8a7aa203e3a --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/common.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Common resource state and persistent-work throttling for MLA decode.""" + +from dataclasses import dataclass +from typing import ClassVar, Optional + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.experimental.task_scheduling.memory import SmemAllocation, TmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ..config import MlaConfig + +# BF16 tcgen05 SMEM descriptors in this schedule use 128B-row swizzled tiles. +# The stride is one swizzle group, and the 16 KiB leading offset selects the +# second 64-wide K block inside the staged 128-token K/V tile. +TCGEN05_BF16_SWIZZLE_STRIDE_BYTES = 1024 +TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES = 16384 +TCGEN05_BF16_K_BLOCK_WIDTH = 64 + + +def _install_task_local_specs(resource: object, specs: tuple[tuple, ...]) -> None: + """Install TaskLocalVariable fields declared by resource classes.""" + for spec in specs: + field_name, dtype, default, docs = spec[:4] + runtime_slot_name = spec[4] if len(spec) > 4 else None + object.__setattr__( + resource, + field_name, + TaskLocalVariable( + dtype=dtype, + default=default, + docs=docs, + runtime_slot_name=runtime_slot_name, + ), + ) + + +# ===================================================================== +# MlaResource — Shared SMEM/TMEM allocation helpers +# ===================================================================== + + +@dataclass(kw_only=True) +class MlaResource(MemoryResource): + """Base resource that owns common MLA config and task-local declarations.""" + + cfg: cutlass.Constexpr[MlaConfig] = None + cu_seqlens_q: object = None + _alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + _tmem_alloc: cutlass.Constexpr[Optional[TmemAllocation]] = None + _tmem_base_addr: object = None + _task_local_specs: ClassVar[tuple[tuple, ...]] = () + + def __post_init__(self) -> None: + _install_task_local_specs(self, self._task_local_specs) + + def get_smem_requirements(self): + """Return SMEM allocations required by this resource.""" + return [] + + def get_tmem_requirements(self): + """Return TMEM allocations required by this resource.""" + return [] + + @cute.jit + def _init_tmem_state(self, stage_info: StageInfo) -> None: + """Initialize shared resource state from the TS allocation context.""" + context = stage_info.context + if cutlass.const_expr( + context is not None + and context.tmem_ptr_i32 is not None + and self._tmem_alloc is not None + ): + self._tmem_base_addr = Int32(context.tmem_ptr_i32.load()) + Int32( + self._tmem_alloc.offset + ) + + +# ===================================================================== +# ScheduleTokenThrottleResource — Dynamic scheduler pacing marker +# ===================================================================== + + +@dataclass(kw_only=True) +class ScheduleTokenThrottleResource(MemoryResource): + """No-op named throttle resource for dynamic persistent schedule token pacing.""" + + # Producer side: the load task marks that its schedule token slot can be reused. + # There is no data payload; the resource exists to make the ordering edge + # visible to TS. + @producer_work + @cute.jit + def publish_schedule_token(self, stage_info: StageInfo): + """Signal that the load task has yielded its schedule token slot.""" + del stage_info + + # Consumer side: the scheduler task consumes the marker before it advances + # to the next CLC work tile. + @consumer_work + @cute.jit + def consume_schedule_token(self, stage_info: StageInfo): + """Wait-side marker before the scheduler reuses a schedule token slot.""" + del stage_info diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py new file mode 100644 index 000000000000..deae8aa5fed2 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py @@ -0,0 +1,380 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SMEM probability resource for the swaps-MMA-AB PV operand.""" + +from dataclasses import dataclass +from typing import Optional + +from cutlass.experimental import primitives as prims + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import SmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + consumer_work, + producer_work, +) + +from ...helpers.constants import ( + SMEM_WORD_BYTE_SHIFT, + SMEM_WORD_BYTES, + TCGEN05_16X256B_REGS_PER_LOAD, + TMEM_LIFECYCLE_BARRIER_ID, +) + + +from ...helpers.layout import ( + _TASK_CACHE_LANE_IDX, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _TASK_CACHE_WARP_IDX, + decode_gen_task_cache, + num_o_stsm_row_blocks, + num_packed_p_regs, + num_softmax_scale_groups, + p_stsm_smem_offset_bytes, + smem_array, +) +from ...helpers.math import ( + neg_max_f32, + pack_float2_to_bf16, + qkv_dtype, +) +from ...helpers.ops import ( + fp8_log2_quant_scale, + pack_float4_to_fp8_e4m3, + store_transposed_smem8b_x2, + store_transposed_smem8b_x4, +) + +from .common import ( + MlaResource, +) + +# ===================================================================== +# SmemPResource — P in SMEM, AsyncUmma pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemPResource(MlaResource): + """SMEM probability tile exchanged from softmax to PV MMA. + + AsyncUmma synchronizes the softmax producer with the UMMA PV consumer. + Softmax writes P with stmatrix and fences the async-shared view before + commit; MmaTask waits before PV and releases after the tensor-core read. + """ + + inst_id: cutlass.Constexpr[int] = 0 + scale_softmax_log2: Float32 = None + tmem_s_ref: Optional[MemoryResource] = None + order_p01_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + owns_order_p01_alloc: cutlass.Constexpr[bool] = False + _smem_p: object = None + _smem_p_i32: object = None + _order_p01_barrier_ptr: object = None + _order_p01_phase: object = None + + def get_smem_requirements(self): + """Return P SMEM plus optional ordering-barrier allocation.""" + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.p_smem_tile_bytes, + alignment=self.cfg.stensor_align, + ) + allocs = [self._alloc] + if self.owns_order_p01_alloc and self.order_p01_alloc is not None: + allocs.append(self.order_p01_alloc) + return allocs + + @cute.jit + def _init_smem_state_from_context(self, context) -> None: + """Create P SMEM views and optional producer-order barrier pointer.""" + self._smem_p = smem_array( + context, + self._alloc, + qkv_dtype(self.cfg), + self.cfg.p_smem_tile_bytes // self.cfg.qkv_dtype_bytes, + ) + self._smem_p_i32 = smem_array( + context, + self._alloc, + Int32, + self.cfg.p_smem_tile_bytes // 4, + ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self.order_p01_alloc is not None + ): + self._order_p01_barrier_ptr = cute.make_ptr( + Int64, + context.smem_base.data_ptr() + self.order_p01_alloc.offset, + mem_space=cute.AddressSpace.smem, + ) + self._order_p01_phase = Int32(1 if self.inst_id == 0 else 0) + + @cute.jit + def initialize_runtime_state_internal( + self, + context=None, + captured_schedule: cutlass.Constexpr[bool] = False, + ) -> None: + """Initialize pipeline state and CTA-wide P-order barriers.""" + super().initialize_runtime_state_internal(context, captured_schedule) + self._init_smem_state_from_context(context) + if cutlass.const_expr( + self.owns_order_p01_alloc + and self.cfg.use_clc_dynamic_persistent_scheduler == 1 + and self._order_p01_barrier_ptr is not None + ): + self._init_order_p01_barriers() + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + """Create P SMEM views for captured schedule aux work.""" + self._init_smem_state_from_context(stage_info.context) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_materialize_state(self, stage_info: StageInfo) -> None: + """Initialize P SMEM state before softmax materializes probabilities.""" + + # Producer aux work prepares P SMEM for the softmax task before it + # materializes probabilities. + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize P SMEM state before PV MMA consumes descriptors.""" + + # Consumer aux work prepares the same P SMEM view for the MMA task. + self._init_smem_state(stage_info) + + @cute.jit + def _init_order_p01_barriers(self): + """Initialize the ordered P0/P1 barriers used by CLC scheduling.""" + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bar_init_warp = Int32(2 if self.inst_id == 0 else 6) + if warp_idx == bar_init_warp: + with cute.arch.elect_one(): + cute.arch.mbarrier_init( + self._order_p01_barrier_ptr, + Int32(128), + ) + cute.arch.mbarrier_init( + self._order_p01_barrier_ptr + 1, + Int32(128), + ) + cute.arch.mbarrier_init_fence() + prims.barrier_cta_sync( + barrier_id=TMEM_LIFECYCLE_BARRIER_ID, + thread_count=self.cfg.threads_per_cta, + ) + + @cute.jit + def _ordered_sequence_wait(self): + """Wait for this P instance's ordered materialization turn.""" + + if cutlass.const_expr(self.cfg.use_clc_dynamic_persistent_scheduler == 1): + cute.arch.mbarrier_wait( + self._order_p01_barrier_ptr + self.inst_id, + self._order_p01_phase, + ) + + @cute.jit + def _ordered_sequence_arrive(self): + """Signal the peer P instance after materialization completes.""" + + if cutlass.const_expr(self.cfg.use_clc_dynamic_persistent_scheduler == 1): + signaling_id = 1 if self.inst_id == 0 else 0 + cute.arch.mbarrier_arrive( + self._order_p01_barrier_ptr + signaling_id, + ) + self._order_p01_phase = self._order_p01_phase ^ Int32(1) + + @producer_work + @cute.jit + def materialize_p( + self, + stage_info: StageInfo, + *, + new_max_arr, + s_arr, + local_sum_arr, + ): + """Materialize grouped-head softmax probabilities in P SMEM for BMM2.""" + # Producer work for SmemPResource: softmax writes P into the acquired + # SMEM stage. The AsyncUmma commit protects the payload until PV MMA + # waits on p_desc() and releases it. + cfg = self.cfg + task_cache = decode_gen_task_cache(stage_info) + + num_scale_groups = num_softmax_scale_groups(cfg) + neg_scaled_max = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + local_sums = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + for idx in cutlass.range_constexpr(num_scale_groups): + new_max = new_max_arr[idx] + safe_new_max = new_max + if safe_new_max == neg_max_f32(): + safe_new_max = Float32(0.0) + neg_scaled_max[idx] = -self.scale_softmax_log2 * safe_new_max + if cutlass.const_expr(cfg.is_fp8_qkv()): + neg_scaled_max[idx] += fp8_log2_quant_scale() + local_sums[idx] = Float32(0.0) + + # Convert S registers to P and accumulate local softmax sums. FP8 packs + # four probabilities per register; BF16 packs two. + packed_p_reg_count = num_packed_p_regs(cfg) + regs_p = cutlass.Array( + Int32, packed_p_reg_count, space=cutlass.AddressSpace.rmem + ) + self._ordered_sequence_wait() + if cutlass.const_expr(cfg.is_fp8_qkv()): + for packed_idx in cutlass.range_constexpr(packed_p_reg_count): + p_vals = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for elem_idx in cutlass.range_constexpr(4): + s_idx = packed_idx * 4 + elem_idx + pair_idx = s_idx // 2 + scale_base = ( + (pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2 + ) * 2 + scale_idx = scale_base + (s_idx % 2) + p_val = Float32(0.0) + if new_max_arr[scale_idx] != neg_max_f32(): + p_val = cute.math.exp2( + s_arr[s_idx] * self.scale_softmax_log2 + + neg_scaled_max[scale_idx], + fastmath=True, + ) + p_vals[elem_idx] = p_val + local_sums[scale_idx] += p_val + regs_p[packed_idx] = pack_float4_to_fp8_e4m3( + p_vals[0], p_vals[1], p_vals[2], p_vals[3] + ) + if cutlass.const_expr(packed_idx == packed_p_reg_count // 2): + self._ordered_sequence_arrive() + else: + for pair_idx in cutlass.range_constexpr(packed_p_reg_count): + scale0 = ((pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2) * 2 + scale1 = scale0 + 1 + s0 = pair_idx * 2 + s1 = s0 + 1 + p0 = Float32(0.0) + p1 = Float32(0.0) + new_max0 = new_max_arr[scale0] + new_max1 = new_max_arr[scale1] + if new_max0 != neg_max_f32(): + p0 = cute.math.exp2( + s_arr[s0] * self.scale_softmax_log2 + neg_scaled_max[scale0], + fastmath=True, + ) + if new_max1 != neg_max_f32(): + p1 = cute.math.exp2( + s_arr[s1] * self.scale_softmax_log2 + neg_scaled_max[scale1], + fastmath=True, + ) + local_sums[scale0] += p0 + local_sums[scale1] += p1 + regs_p[pair_idx] = pack_float2_to_bf16(p0, p1) + if cutlass.const_expr(pair_idx == packed_p_reg_count // 2): + self._ordered_sequence_arrive() + + # Publish local sums before committing P so the softmax sum update can + # consume them while PV MMA waits on the AsyncUmma P stage. + for scale_idx in cutlass.range_constexpr(num_scale_groups): + self.tmem_s_ref._p_local_sum_arr[scale_idx] = local_sums[scale_idx] + local_sum_arr[scale_idx] = local_sums[scale_idx] + + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + if cutlass.const_expr(cfg.is_fp8_qkv()): + # FP8 P uses byte-transposed STSM stores so PV MMA can consume the + # same logical P tile layout as the BF16 path. + if cutlass.const_expr(packed_p_reg_count == 2): + store_transposed_smem8b_x2( + self._smem_p_i32, + regs_p[0], + regs_p[1], + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.tile_size_kv, + ) + else: + for stsm_chunk_idx in cutlass.range_constexpr(packed_p_reg_count // 4): + reg_base = stsm_chunk_idx * TCGEN05_16X256B_REGS_PER_LOAD + store_transposed_smem8b_x4( + self._smem_p_i32, + regs_p[reg_base], + regs_p[reg_base + 1], + regs_p[reg_base + 2], + regs_p[reg_base + 3], + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.tile_size_kv, + stsm_idx=stsm_chunk_idx, + ) + # Every softmax producer thread reaches the AsyncUmma commit after + # this proxy fence. Its 128-thread full-mbarrier arrival count is + # already the cross-warp rendezvous consumed by the MMA wait, so a + # second named CTA barrier would only serialize the producer. + cute.arch.fence_view_async_shared() + return + + for stsm_chunk_idx in cutlass.range_constexpr(packed_p_reg_count // 4): + stsm_group_idx = stsm_chunk_idx // num_o_stsm_row_blocks(cfg) + stsm_row_block_idx = stsm_chunk_idx % num_o_stsm_row_blocks(cfg) + smem_offset_bytes = p_stsm_smem_offset_bytes( + warp_idx, + lane_idx, + stsm_group_idx, + stsm_row_block_idx, + cfg.tile_size_q, + ) + smem_dst = self._smem_p_i32.data_ptr( + smem_offset_bytes >> SMEM_WORD_BYTE_SHIFT + ) + prims.stmatrix( + smem_dst, + ( + regs_p.data_ptr() + stsm_chunk_idx * TCGEN05_16X256B_REGS_PER_LOAD + ).load( + count=TCGEN05_16X256B_REGS_PER_LOAD, + alignment=SMEM_WORD_BYTES, + ), + prims.MMALayout.COL, + shape=prims.StoreShape.M8N8, + ) + cute.arch.fence_view_async_shared() + + @consumer_work + @cute.jit + def p_desc(self, stage_info: StageInfo): + """Publish the P SMEM descriptor consumed by PV MMA.""" + # The descriptor itself is reconstructed by TmemOResource from this + # resource's SMEM pointer. This consumer work is still needed as the TS + # wait/release edge that keeps P live until PV MMA is done. + del stage_info diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py new file mode 100644 index 000000000000..8e2544d61276 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py @@ -0,0 +1,1174 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SMEM staging resources for Q, page offsets, and K/V tiles.""" + +from dataclasses import dataclass +from typing import Any, ClassVar, Optional + +from cutlass.experimental import primitives as prims +from ....tensor_map import transform_ragged_coords + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 +from cutlass.cutlass_dsl import Boolean, dsl_user_op, if_generate +from cutlass.pipeline import PipelineAsync, PipelineState +from cutlass.experimental import primitives as cprims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import SmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...helpers.constants import ( + CP_ASYNC_CACHE_CA, + PAGE_OFFSET_BYTES, + WARP_LANES, +) + +from ...helpers.layout import ( + decode_gen_task_cache, + head_dim_cta_offset_v, + q_stage_smem_element_offset, + smem_array, + tma_inner_dim_elems, +) +from ...helpers.math import ( + qkv_dtype, + qkv_major_k_stride_bytes_for, + qkv_smem_swizzle, + qkv_smem_swizzle_for_head_dim, +) +from ...helpers.ops import ( + lane_idx_from_thread, +) +from ...helpers.query import query_batch_bounds +from ...helpers.stage import MlaStage +from ...helpers.tile import ( + batch_idx_for_stage_cfg, + cta_idx_head_dim_v_for_stage, + cta_idx_kv_for_stage, + cta_idx_q_for_stage, + global_kv_tile_idx, + head_idx_for_stage, + local_kv_tile_idx, + runtime_seq_len_kv_from_task_cache, + staged_kv_head_dim_call_idx, +) + +from .common import ( + TCGEN05_BF16_K_BLOCK_WIDTH, + TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES, + TCGEN05_BF16_SWIZZLE_STRIDE_BYTES, + MlaResource, +) + +# ===================================================================== +# SmemQResource — Q SMEM buffer with TmaUmmaAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemQResource(MlaResource): + """Stage Q latent/RoPE tiles in SMEM and publish Q MMA descriptors.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "q_desc_var", + Int64, + Int64(0), + "SMEM descriptor for the staged Q latent tile.", + "q_desc", + ), + ( + "q_desc_rope_var", + Int64, + Int64(0), + "SMEM descriptor for the staged Q rope tile.", + "q_desc_rope", + ), + ) + tma_desc_q_latent: object = None + tma_desc_q_rope: object = None + head_idx: object = None + batch_idx: object = None + cta_idx_q: object = None + _smem_q: object = None + _q_desc_base: object = None + q_desc_var: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + q_desc_rope_var: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def get_smem_requirements(self): + """Return the SMEM allocation used for staged Q tiles.""" + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.qk_smem_tile_bytes * self.cfg.q_stages, + alignment=self.cfg.stensor_align, + ) + return [self._alloc] + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + """Create the Q SMEM view and base descriptor state.""" + context = stage_info.context + elem_count = ( + self.cfg.qk_smem_tile_bytes * self.cfg.q_stages + ) // self.cfg.qkv_dtype_bytes + self._smem_q = smem_array(context, self._alloc, qkv_dtype(self.cfg), elem_count) + if self._smem_q is not None: + # The BF16 leading offset advances by at most one 64-wide K block. + # FP8 uses a dtype-specific major-K stride and therefore rebuilds + # both descriptor offsets below. + q_leading_byte_offset = Int32( + self.cfg.tile_size_q + * min(self.cfg.head_dim_per_stage_kv, TCGEN05_BF16_K_BLOCK_WIDTH) + * self.cfg.qkv_dtype_bytes + ) + q_stride_byte_offset = Int32(TCGEN05_BF16_SWIZZLE_STRIDE_BYTES) + if cutlass.const_expr(self.cfg.is_fp8_qkv()): + q_leading_byte_offset = Int32( + self.cfg.tile_size_q + * self.cfg.head_dim_per_stage_kv + * self.cfg.qkv_dtype_bytes + ) + q_stride_byte_offset = Int32( + qkv_major_k_stride_bytes_for( + self.cfg, self.cfg.head_dim_per_stage_kv + ) + ) + self._q_desc_base = cprims.Tcgen05SmemDesc.build( + self._smem_q, + leading_byte_offset=q_leading_byte_offset, + stride_byte_offset=q_stride_byte_offset, + layout=qkv_smem_swizzle(self.cfg), + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize the Q SMEM view for producer-side TMA loads.""" + + # Producer aux work runs before the load task starts issuing Q TMA + # copies. It creates the SMEM view used by producer_work(load_q). + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize the Q SMEM view for consumer-side MMA descriptors.""" + + # Consumer aux work creates the same SMEM view in the MMA task so the + # consumer can build descriptors after it waits on the Q stage. + self._init_smem_state(stage_info) + + @cute.jit + def _query_tma_coords(self, dim_offset, local_flat_query_row, batch_idx): + """Return fixed or ragged TMA coordinates for one Q slice.""" + + if cutlass.const_expr(self.cu_seqlens_q is None): + return (dim_offset, local_flat_query_row, batch_idx) + + query_start, query_length = query_batch_bounds( + self.cu_seqlens_q, + batch_idx, + self.cfg.logical_seq_len_q, + ) + storage_flat_query_row = ( + query_start * Int32(self.cfg.logical_num_heads_q) + local_flat_query_row + ) + ragged_extent = ( + query_length * Int32(self.cfg.logical_num_heads_q) - local_flat_query_row + ) + return transform_ragged_coords( + (dim_offset, storage_flat_query_row), + ragged_dim_idx=1, + ragged_box_size=self.cfg.tile_size_q, + ragged_extent=ragged_extent, + ) + + @cute.jit + def _load_q_stage( + self, + stage_info: StageInfo, + stage_base, + qk_stage_idx: int, + dim_offset: Int32, + tma_desc, + ): + """Issue one staged TMA load for a Q head-dimension slice.""" + if cutlass.const_expr(tma_desc is None): + return + active_width = self.cfg.qk_head_stage_width(qk_stage_idx) + inner_width = min(tma_inner_dim_elems(self.cfg), active_width) + head_idx = head_idx_for_stage(self.head_idx, self.cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, self.cfg, stage_info) + # Ragged storage adds the cumulative batch row offset while preserving + # the unclamped logical row so a fully padded CTA becomes an OOB + # zero-fill instead of reloading the final real query row. + local_flat_query_row = cta_idx_q * Int32(self.cfg.num_heads_q) + head_idx + query_coords = self._query_tma_coords( + dim_offset, + local_flat_query_row, + batch_idx, + ) + if prims.elect_sync(): + smem_offset = Int32(q_stage_smem_element_offset(self.cfg, qk_stage_idx)) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr(smem_offset), + tma_desc, + query_coords, + stage_info.barrier, + ) + if cutlass.const_expr(active_width > inner_width): + second_query_coords = self._query_tma_coords( + dim_offset + Int32(inner_width), + local_flat_query_row, + batch_idx, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr( + smem_offset + Int32(inner_width * self.cfg.tile_size_q) + ), + tma_desc, + second_query_coords, + stage_info.barrier, + ) + + @producer_work + @cute.jit + def load_q(self, stage_info: StageInfo): + """Load the grouped-head Q tile as 8 latent chunks plus 1 rope chunk.""" + # This is the actual Q payload producer: the load task fills the + # acquired Q SMEM stage and the pipeline commit makes it visible to MMA. + if cutlass.const_expr( + self.tma_desc_q_latent is None or self.tma_desc_q_rope is None + ): + return + cfg = self.cfg + stage_elems = cfg.q_smem_tile_elements + stage_base = self._smem_q.subview(stage_info.stage_idx * Int32(stage_elems)) + latent_stages = cfg.latent_dim // cfg.head_dim_per_stage_kv + for qk_stage_idx in cutlass.range_constexpr(latent_stages): + dim_offset = Int32(qk_stage_idx * cfg.head_dim_per_stage_kv) + self._load_q_stage( + stage_info, + stage_base, + qk_stage_idx, + dim_offset, + self.tma_desc_q_latent, + ) + self._load_q_stage( + stage_info, + stage_base, + latent_stages, + Int32(0), + self.tma_desc_q_rope, + ) + + @consumer_work(returns=("q_desc", "q_desc_rope")) + @cute.jit + def q_desc(self, stage_info: StageInfo): + """Publish the Q SMEM descriptor for staged QK MMA.""" + # The MMA task has waited for the Q stage. It converts the live SMEM + # stage into descriptors and returns them as task-local values for the + # downstream TmemSResource producer_work(qk_mma). + stage_base = self._smem_q.subview( + stage_info.stage_idx * Int32(self.cfg.q_smem_tile_elements) + ) + # Build the same descriptor as producer init, but relative to the live + # pipeline stage that the MMA task has already waited on. + q_leading_byte_offset = Int32( + self.cfg.tile_size_q + * min(self.cfg.head_dim_per_stage_kv, TCGEN05_BF16_K_BLOCK_WIDTH) + * self.cfg.qkv_dtype_bytes + ) + q_stride_byte_offset = Int32(TCGEN05_BF16_SWIZZLE_STRIDE_BYTES) + if cutlass.const_expr(self.cfg.is_fp8_qkv()): + q_leading_byte_offset = Int32( + self.cfg.tile_size_q + * self.cfg.head_dim_per_stage_kv + * self.cfg.qkv_dtype_bytes + ) + q_stride_byte_offset = Int32( + qkv_major_k_stride_bytes_for(self.cfg, self.cfg.head_dim_per_stage_kv) + ) + desc_q = cprims.Tcgen05SmemDesc.build( + stage_base, + leading_byte_offset=q_leading_byte_offset, + stride_byte_offset=q_stride_byte_offset, + layout=qkv_smem_swizzle(self.cfg), + ) + desc_q_rope = desc_q + if cutlass.const_expr(self.cfg.is_fp8_qkv() and self.cfg.rope_dim == 64): + rope_stage_idx = self.cfg.latent_dim // self.cfg.head_dim_per_stage_kv + rope_stage_offset = Int32( + q_stage_smem_element_offset(self.cfg, rope_stage_idx) + ) + desc_q_rope = cprims.Tcgen05SmemDesc.build( + stage_base.subview(rope_stage_offset), + leading_byte_offset=Int32( + self.cfg.tile_size_q * self.cfg.rope_dim * self.cfg.qkv_dtype_bytes + ), + stride_byte_offset=Int32( + qkv_major_k_stride_bytes_for(self.cfg, self.cfg.rope_dim) + ), + layout=qkv_smem_swizzle_for_head_dim(self.cfg, self.cfg.rope_dim), + ) + return desc_q, desc_q_rope + + +# ===================================================================== +# SmemPageOffsetsResource — Page-offset SMEM buffer with Async pipeline +# ===================================================================== + + +@dataclass(frozen=True) +class _StructuredWaitPipelineAsync(PipelineAsync): + """Page-offset pipeline with a public structured mbarrier retry loop.""" + + @cute.jit + def _retry_wait( + self, + sync_object: object, + state: PipelineState, + *, + loc: Any = None, + ip: Any = None, + ) -> None: + while not sync_object.try_wait( + state.index, + state.phase, + loc=loc, + ip=ip, + ): + pass + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc: Any = None, + ip: Any = None, + ) -> None: + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self._retry_wait(self.sync_object_empty, state, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + @dsl_user_op + def consumer_wait( + self, + state: PipelineState, + try_wait_token: Optional[Boolean] = None, + *, + loc: Any = None, + ip: Any = None, + ) -> None: + if_generate( + try_wait_token is None or try_wait_token == 0, + lambda: self._retry_wait(self.sync_object_full, state, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + + +@dataclass(kw_only=True) +class SmemPageOffsetsResource(MlaResource): + """Prefetch page ids for the K/V tile currently owned by the load task.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "cached_page_ids", + cutlass.Array, + None, + "Cached logical page ids for the staged K/V tile.", + ), + ) + page_offsets: object = None + cache_seqs: object = None + batch_idx: object = None + cta_idx_q: object = None + cta_idx_kv: object = None + _smem_page_offsets: object = None + cached_page_ids: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def create_pipeline(self, pipeline_config): + """Preserve stock barrier objects and specialize only their wait path.""" + base = super().create_pipeline(pipeline_config) + if not isinstance(base, PipelineAsync): + raise TypeError( + "page-offset staging requires cutlass.pipeline.PipelineAsync, " + f"got {type(base).__name__}" + ) + return _StructuredWaitPipelineAsync( + sync_object_full=base.sync_object_full, + sync_object_empty=base.sync_object_empty, + num_stages=base.num_stages, + producer_mask=base.producer_mask, + consumer_mask=base.consumer_mask, + ) + + def get_smem_requirements(self): + """Return the SMEM allocation for cached page offsets.""" + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=( + self.cfg.page_offsets_stages + * self.cfg.page_offsets_entries_per_stage + * PAGE_OFFSET_BYTES + ), + alignment=128, + ) + return [self._alloc] + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + """Create the SMEM page-offset array view.""" + context = stage_info.context + count = self.cfg.page_offsets_stages * self.cfg.page_offsets_entries_per_stage + self._smem_page_offsets = smem_array( + context, + self._alloc, + Int32, + count, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize page-offset SMEM before producer copies begin.""" + + # Producer aux work creates the page-offset SMEM view for the page load + # warp before it issues cp.async copies. + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=cached_page_ids) + @cute.jit + def init_read_state(self, stage_info: StageInfo): + """Create register caches for the active K/V page-ID streams.""" + + # Consumer aux work allocates the register cache that read_offsets() + # returns after the page-offset stage is ready. + self._init_smem_state(stage_info) + cache_slots = 2 if self.cfg.kernel_variant != "keeps_mma_ab" else 1 + return cutlass.Array( + Int32, + max(1, cache_slots * self.cfg.pages_per_kv_tile), + space=cutlass.AddressSpace.rmem, + ) + + @cute.jit + def page_id(self, page_fragment_idx: int): + """Return one cached logical page id from the consumer stage.""" + offset = self.consumer_work_stage * Int32( + self.cfg.page_offsets_entries_per_stage + ) + Int32(page_fragment_idx) + return Int32(self._smem_page_offsets[offset]) + + @cute.jit + def _producer_load_page_offsets( + self, + stage_info: StageInfo, + inst_id: int, + is_v: int, + *, + section: cutlass.Constexpr[MlaStage], + ): + """Load page ids for one K or V producer stage into SMEM.""" + if cutlass.const_expr(self.page_offsets is None): + return + cfg = self.cfg + lane_idx = lane_idx_from_thread(cute.arch.thread_idx()[0]) + pages_per_tile = Int32(cfg.pages_per_kv_tile) + if lane_idx < pages_per_tile: + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + cfg, + decode_gen_task_cache(stage_info), + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + local_tile_idx = local_kv_tile_idx( + cfg, stage_info, inst_id, is_v, section=section + ) + tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) + last_valid_page = cute.math.max( + (seq_len_kv + Int32(cfg.num_tokens_per_page - 1)) + // Int32(cfg.num_tokens_per_page) + - Int32(1), + Int32(0), + ) + logical_page_idx = cute.math.min( + tile_idx * pages_per_tile + lane_idx, + last_valid_page, + ) + smem_offset = ( + stage_info.stage_idx * Int32(cfg.page_offsets_entries_per_stage) + + lane_idx + ) + page_offsets_batch = self.page_offsets[None, batch_idx] + page_offsets_flat = cute.flat_divide(page_offsets_batch, (1,)) + gmem_ptr = page_offsets_flat[None, logical_page_idx].iterator.llvm_ptr + smem_ptr = cutlass.inttoptr( + self._smem_page_offsets.data_ptr(smem_offset).toint(cutlass.Int32), + 3, + cutlass.Int32, + ) + prims.cp_async_shared_global( + smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA + ) + + @producer_work + @cute.jit + def load_k0(self, stage_info: StageInfo, *, section: cutlass.Constexpr[MlaStage]): + """Prefetch page offsets for K instance 0.""" + + # Producer work for the page-offset stage that feeds K instance 0. + self._producer_load_page_offsets(stage_info, 0, 0, section=section) + + @producer_work + @cute.jit + def load_k1(self, stage_info: StageInfo, *, section: cutlass.Constexpr[MlaStage]): + """Prefetch page offsets for K instance 1.""" + + # Producer work for the page-offset stage that feeds K instance 1. + self._producer_load_page_offsets(stage_info, 1, 0, section=section) + + @producer_work + @cute.jit + def load_v0(self, stage_info: StageInfo, *, section: cutlass.Constexpr[MlaStage]): + """Prefetch deferred page offsets for V instance 0.""" + + # Producer work for the deferred page-offset stage that feeds V instance 0. + self._producer_load_page_offsets(stage_info, 0, 1, section=section) + + @producer_work + @cute.jit + def load_v1(self, stage_info: StageInfo, *, section: cutlass.Constexpr[MlaStage]): + """Prefetch deferred page offsets for V instance 1.""" + + # Producer work for the deferred page-offset stage that feeds V instance 1. + self._producer_load_page_offsets(stage_info, 1, 1, section=section) + + @consumer_work(returns=cached_page_ids) + @cute.jit + def read_offsets( + self, + stage_info: StageInfo, + *, + cached_page_ids, + cache_slot: cutlass.Constexpr[int] = 0, + ): + """Cache staged page IDs in the selected K stream's register slot.""" + # The load task waits on the page-offset stage, snapshots the page ids + # into registers, and keeps swaps K0/K1 live until their delayed V use. + del stage_info + cache_base = cache_slot * self.cfg.pages_per_kv_tile + for page_frag in cutlass.range_constexpr(self.cfg.pages_per_kv_tile): + cached_page_ids[cache_base + page_frag] = self.page_id(page_frag) + return cached_page_ids + + +# ===================================================================== +# SmemKvResource — K/V SMEM buffer with TmaUmmaAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemKvResource(MlaResource): + """Shared K/V staging resource for latent and rope tiles.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "kv_desc_var", + Int64, + Int64(0), + "SMEM descriptor for staged K tiles.", + "kv_desc", + ), + ( + "v_desc_0_var", + Int64, + Int64(0), + "First V descriptor consumed by PV MMA.", + "v_desc_0", + ), + ( + "v_desc_1_var", + Int64, + Int64(0), + "Second V descriptor consumed by PV MMA.", + "v_desc_1", + ), + ) + tma_desc_c_latent: object = None + tma_desc_c_rope: object = None + tma_desc_v: object = None + c_rope_tensor: object = None + page_offsets_kv: object = None + page_offsets: object = None + cache_seqs: object = None + head_idx: object = None + batch_idx: object = None + cta_idx_q: object = None + cta_idx_kv: object = None + cta_idx_head_dim_v: object = None + _smem_kv: object = None + _k_desc_base: object = None + _v_desc_base: object = None + kv_desc_var: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + v_desc_0_var: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + v_desc_1_var: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def get_smem_requirements(self): + """Return the SMEM allocation used for staged K/V tiles.""" + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else self.cfg.kv_stages + ) + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.kv_smem_tile_bytes * num_stages, + alignment=self.cfg.stensor_align, + ) + return [self._alloc] + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + """Create K/V SMEM views and base descriptors.""" + context = stage_info.context + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else self.cfg.kv_stages + ) + elem_count = ( + self.cfg.kv_smem_tile_bytes * num_stages + ) // self.cfg.qkv_dtype_bytes + self._smem_kv = smem_array( + context, self._alloc, qkv_dtype(self.cfg), elem_count + ) + if self._smem_kv is not None: + # FP8 descriptors use dtype-specific major-K strides rather than + # the BF16 128B-row swizzle group. + k_leading_byte_offset = Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + v_leading_byte_offset = ( + Int32(0) + if self.cfg.head_dim_per_stage_v == TCGEN05_BF16_K_BLOCK_WIDTH + else Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + ) + stride_byte_offset = Int32(TCGEN05_BF16_SWIZZLE_STRIDE_BYTES) + if cutlass.const_expr(self.cfg.is_fp8_qkv()): + k_leading_byte_offset = Int32(self.cfg.kv_smem_tile_bytes) + v_leading_byte_offset = Int32(0) + stride_byte_offset = Int32( + qkv_major_k_stride_bytes_for( + self.cfg, self.cfg.head_dim_per_stage_kv + ) + ) + self._k_desc_base = cprims.Tcgen05SmemDesc.build( + self._smem_kv, + leading_byte_offset=k_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=qkv_smem_swizzle(self.cfg), + ) + self._v_desc_base = cprims.Tcgen05SmemDesc.build( + self._smem_kv, + leading_byte_offset=v_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=qkv_smem_swizzle(self.cfg), + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize K/V SMEM state before producer TMA loads.""" + + # Producer aux work initializes K/V SMEM views for the load task. + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize K/V SMEM state before descriptor consumption.""" + + # Consumer aux work initializes K/V SMEM views for the MMA task. + self._init_smem_state(stage_info) + + @cute.jit + def _stage_base(self, stage_info: StageInfo): + """Return the base of the current K/V SMEM pipeline stage.""" + + stage_elems = self.cfg.kv_smem_stage_elements + return self._smem_kv.subview(stage_info.stage_idx * Int32(stage_elems)) + + @cute.jit + def _producer_copy_rope_stage( + self, + stage_info: StageInfo, + page_ids, + stage_base, + ): + """Copy paged rope values into the staged K/V SMEM tile.""" + cfg = self.cfg + lane_idx = lane_idx_from_thread(cute.arch.thread_idx()[0]) + total_elems = Int32(cfg.tile_size_kv * cfg.head_dim_per_stage_kv) + for elem_idx in cutlass.range( + lane_idx, total_elems, Int32(WARP_LANES), unroll=1 + ): + token_idx = elem_idx // Int32(cfg.head_dim_per_stage_kv) + dim_idx = elem_idx - token_idx * Int32(cfg.head_dim_per_stage_kv) + src_dim_idx = dim_idx + if dim_idx >= Int32(cfg.rope_dim): + src_dim_idx = dim_idx - Int32(cfg.rope_dim) + page_frag = token_idx // Int32(cfg.num_tokens_per_page) + token_in_page = token_idx - page_frag * Int32(cfg.num_tokens_per_page) + page_id = Int32(page_ids[page_frag]) + stage_base[elem_idx] = self.c_rope_tensor[ + token_in_page, src_dim_idx, page_id + ] + cute.arch.fence_view_async_shared() + if prims.elect_sync(): + prims.mbarrier_complete_tx( + stage_info.barrier, Int32(cfg.kv_smem_tile_bytes) + ) + + @cute.jit + def _producer_load_kv_stage( + self, + stage_info: StageInfo, + inst_id: int, + is_v: int, + *, + stage_idx: cutlass.Constexpr[int], + section: cutlass.Constexpr[MlaStage], + cached_page_ids, + page_id_slot: cutlass.Constexpr[int] = 0, + ): + """Load one staged K or V tile using the consumed page-id payload.""" + cfg = self.cfg + qk_stage_idx = staged_kv_head_dim_call_idx( + cfg, + stage_info, + inst_id, + is_v, + stage_idx=stage_idx, + section=section, + ) + if cutlass.const_expr(is_v): + active_width = cfg.v_head_stage_width(qk_stage_idx) + cta_idx_head_dim_v = cta_idx_head_dim_v_for_stage( + self.cta_idx_head_dim_v, stage_info + ) + dim_offset = head_dim_cta_offset_v(cfg, cta_idx_head_dim_v) + Int32( + qk_stage_idx * cfg.head_dim_per_stage_v + ) + tma_desc = self.tma_desc_v + else: + active_width = cfg.qk_head_stage_width(qk_stage_idx) + if cutlass.const_expr( + qk_stage_idx < cfg.latent_dim // cfg.head_dim_per_stage_kv + ): + dim_offset = Int32(qk_stage_idx * cfg.head_dim_per_stage_kv) + tma_desc = self.tma_desc_c_latent + else: + dim_offset = Int32(0) + tma_desc = self.tma_desc_c_rope + + if cutlass.const_expr(tma_desc is None): + return + local_tile_idx = local_kv_tile_idx( + cfg, stage_info, inst_id, is_v, section=section + ) + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + cfg, + decode_gen_task_cache(stage_info), + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) + stage_base = self._stage_base(stage_info) + inner_width = min(tma_inner_dim_elems(cfg), active_width) + uses_compact_fp8_rope_stage = ( + cfg.load_num_warps == 1 + and cfg.is_fp8_qkv() + and not is_v + and cfg.rope_dim == 64 + and active_width == cfg.rope_dim + and cfg.head_dim_per_stage_kv == 128 + ) + + if cutlass.const_expr(cfg.use_paged_kv == 1): + if cutlass.const_expr( + self.page_offsets_kv is None and self.page_offsets is None + ): + return + page_ids = None + if cutlass.const_expr(self.page_offsets_kv is not None): + page_ids = cached_page_ids + pages_per_tile = cfg.pages_per_kv_tile + first_page_elems = inner_width * cfg.num_tokens_per_page + first_tile_elems = inner_width * cfg.tile_size_kv + active_tile_elems = active_width * cfg.tile_size_kv + last_valid_page = cute.math.max( + (seq_len_kv + Int32(cfg.num_tokens_per_page - 1)) + // Int32(cfg.num_tokens_per_page) + - Int32(1), + Int32(0), + ) + # Keep the small set of independent page TMA issues straight-line. + for page_frag in cutlass.range(pages_per_tile, unroll_full=True): + if cutlass.const_expr(self.page_offsets_kv is not None): + page_id = Int32( + page_ids[page_id_slot * cfg.pages_per_kv_tile + page_frag] + ) + else: + logical_page_idx = cute.math.min( + tile_idx * Int32(pages_per_tile) + Int32(page_frag), + last_valid_page, + ) + page_id = Int32(self.page_offsets[logical_page_idx, batch_idx]) + page_base = Int32(page_frag * first_page_elems) + smem_page_offset = page_base + if prims.elect_sync(): + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr(smem_page_offset), + tma_desc, + (dim_offset, Int32(0), page_id), + stage_info.barrier, + ) + if cutlass.const_expr(active_width > inner_width): + second_half_offset = Int32(first_tile_elems) + smem_page_offset + if prims.elect_sync(): + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr(second_half_offset), + tma_desc, + (dim_offset + Int32(inner_width), Int32(0), page_id), + stage_info.barrier, + ) + if cutlass.const_expr( + active_width < cfg.head_dim_per_stage_kv + and not uses_compact_fp8_rope_stage + ): + if prims.elect_sync(): + # The shared K/V TMA pipeline expects one full QK-sized + # stage. Duplicate short K/V slices into the unused half + # so the expected transaction byte count is satisfied. + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr( + Int32(active_tile_elems) + smem_page_offset + ), + tma_desc, + (dim_offset, Int32(0), page_id), + stage_info.barrier, + ) + if cutlass.const_expr(uses_compact_fp8_rope_stage): + if prims.elect_sync(): + # The FP8 RoPE MMA consumes only the 64-wide s64b payload. + # Retire the intentionally unwritten half of the fixed + # 128-wide pipeline stage instead of rereading every page. + missing_stage_bytes = ( + cfg.tile_size_kv + * (cfg.head_dim_per_stage_kv - active_width) + * cfg.qkv_dtype_bytes + ) + prims.mbarrier_complete_tx( + stage_info.barrier, Int32(missing_stage_bytes) + ) + else: + tile_offset = tile_idx * Int32(cfg.tile_size_kv) + if prims.elect_sync(): + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base, + tma_desc, + ( + dim_offset, + tile_offset, + head_idx_for_stage(self.head_idx, cfg, stage_info), + batch_idx, + ), + stage_info.barrier, + ) + if cutlass.const_expr(active_width > inner_width): + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr(Int32(inner_width * cfg.tile_size_kv)), + tma_desc, + ( + dim_offset + Int32(inner_width), + tile_offset, + head_idx_for_stage(self.head_idx, cfg, stage_info), + batch_idx, + ), + stage_info.barrier, + ) + if cutlass.const_expr(active_width < cfg.head_dim_per_stage_kv): + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr(Int32(active_width * cfg.tile_size_kv)), + tma_desc, + ( + dim_offset, + tile_offset, + head_idx_for_stage(self.head_idx, cfg, stage_info), + batch_idx, + ), + stage_info.barrier, + ) + + @producer_work + @cute.jit + def load_k0( + self, + stage_info: StageInfo, + *, + stage_idx: cutlass.Constexpr[int], + section: cutlass.Constexpr[MlaStage], + cached_page_ids, + page_id_slot: cutlass.Constexpr[int] = 0, + ): + """Load K instance 0 into the acquired K/V SMEM stage.""" + + # ProdWork: fill K instance 0 from the page IDs consumed by the load task. + self._producer_load_kv_stage( + stage_info, + 0, + 0, + stage_idx=stage_idx, + section=section, + cached_page_ids=cached_page_ids, + page_id_slot=page_id_slot, + ) + + @producer_work + @cute.jit + def load_k1( + self, + stage_info: StageInfo, + *, + stage_idx: cutlass.Constexpr[int], + section: cutlass.Constexpr[MlaStage], + cached_page_ids, + page_id_slot: cutlass.Constexpr[int] = 0, + ): + """Load K instance 1 into the acquired K/V SMEM stage.""" + + # ProdWork: fill K instance 1 from the page IDs consumed by the load task. + self._producer_load_kv_stage( + stage_info, + 1, + 0, + stage_idx=stage_idx, + section=section, + cached_page_ids=cached_page_ids, + page_id_slot=page_id_slot, + ) + + @producer_work + @cute.jit + def load_v0( + self, + stage_info: StageInfo, + *, + stage_idx: cutlass.Constexpr[int], + section: cutlass.Constexpr[MlaStage], + cached_page_ids, + page_id_slot: cutlass.Constexpr[int] = 0, + ): + """Load deferred V instance 0 into the acquired K/V SMEM stage.""" + + # ProdWork: fill V instance 0 from the page IDs consumed by the load task. + self._producer_load_kv_stage( + stage_info, + 0, + 1, + stage_idx=stage_idx, + section=section, + cached_page_ids=cached_page_ids, + page_id_slot=page_id_slot, + ) + + @producer_work + @cute.jit + def load_v1( + self, + stage_info: StageInfo, + *, + stage_idx: cutlass.Constexpr[int], + section: cutlass.Constexpr[MlaStage], + cached_page_ids, + page_id_slot: cutlass.Constexpr[int] = 0, + ): + """Load deferred V instance 1 into the acquired K/V SMEM stage.""" + + # ProdWork: fill V instance 1 from the page IDs consumed by the load task. + self._producer_load_kv_stage( + stage_info, + 1, + 1, + stage_idx=stage_idx, + section=section, + cached_page_ids=cached_page_ids, + page_id_slot=page_id_slot, + ) + + @cute.jit + def _set_desc( + self, + stage_info: StageInfo, + is_v: int, + inst_id: int = 0, + *, + k_subtile_idx: cutlass.Constexpr[int] = 0, + ): + """Build a K or V SMEM descriptor for the current producer stage.""" + # Descriptor offsets mirror _init_smem_state, but are rebuilt against the + # current pipeline stage because K and delayed-V can be consumed from + # different stages. + stride_byte_offset = Int32(TCGEN05_BF16_SWIZZLE_STRIDE_BYTES) + if cutlass.const_expr(self.cfg.is_fp8_qkv()): + stride_byte_offset = Int32( + qkv_major_k_stride_bytes_for(self.cfg, self.cfg.head_dim_per_stage_kv) + ) + if cutlass.const_expr(is_v): + v_leading_byte_offset = ( + Int32(0) + if self.cfg.head_dim_per_stage_v == TCGEN05_BF16_K_BLOCK_WIDTH + else Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + ) + if cutlass.const_expr(self.cfg.is_fp8_qkv()): + v_leading_byte_offset = Int32(0) + desc = cprims.Tcgen05SmemDesc.build( + self._stage_base(stage_info), + leading_byte_offset=v_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=qkv_smem_swizzle(self.cfg), + ) + return desc + else: + k_leading_byte_offset = Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + if cutlass.const_expr(self.cfg.is_fp8_qkv()): + k_leading_byte_offset = Int32(self.cfg.kv_smem_tile_bytes) + qk_stage_idx = k_subtile_idx + if cutlass.const_expr( + self.cfg.rope_dim == 64 + and qk_stage_idx + == self.cfg.latent_dim // self.cfg.head_dim_per_stage_kv + ): + k_leading_byte_offset = Int32( + self.cfg.tile_size_kv + * self.cfg.rope_dim + * self.cfg.qkv_dtype_bytes + ) + return cprims.Tcgen05SmemDesc.build( + self._stage_base(stage_info), + leading_byte_offset=k_leading_byte_offset, + stride_byte_offset=Int32( + qkv_major_k_stride_bytes_for(self.cfg, self.cfg.rope_dim) + ), + layout=qkv_smem_swizzle_for_head_dim( + self.cfg, self.cfg.rope_dim + ), + ) + desc = cprims.Tcgen05SmemDesc.build( + self._stage_base(stage_info), + leading_byte_offset=k_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=qkv_smem_swizzle(self.cfg), + ) + return desc + + @consumer_work(returns=("kv_desc",)) + @cute.jit + def k_desc_0(self, stage_info: StageInfo, *, k_subtile_idx: cutlass.Constexpr[int]): + """Return the K descriptor for consumer instance 0.""" + # MMA has waited on the K stage; return a descriptor for qk_mma(). + return self._set_desc(stage_info, 0, 0, k_subtile_idx=k_subtile_idx) + + @consumer_work(returns=("kv_desc",)) + @cute.jit + def k_desc_1(self, stage_info: StageInfo, *, k_subtile_idx: cutlass.Constexpr[int]): + """Return the K descriptor for consumer instance 1.""" + # Same K SMEM payload, second interleaved QK instance. + return self._set_desc(stage_info, 0, 1, k_subtile_idx=k_subtile_idx) + + @consumer_work(returns=("v_desc_0",)) + @cute.jit + def v_desc_0(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int]): + """Return the V descriptor for consumer instance 0.""" + # PV MMA consumes this descriptor after V instance 0 is ready. + del v_subtile_idx + return self._set_desc(stage_info, 1, 0) + + @consumer_work(returns=("v_desc_1",)) + @cute.jit + def v_desc_1(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int]): + """Return the V descriptor for consumer instance 1.""" + # PV MMA consumes this descriptor after V instance 1 is ready. + del v_subtile_idx + return self._set_desc(stage_info, 1, 1) + + +# ===================================================================== +# SmemKResource — K-only SMEM view helper +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemKResource(MlaResource): + """K/V SMEM staging buffer and descriptor producer for QK/PV MMA.""" + + inst_id: cutlass.Constexpr[int] = 0 + is_v: cutlass.Constexpr[int] = 0 + + +# ===================================================================== +# SmemVResource — V-only SMEM view helper +# ===================================================================== + + +@dataclass(kw_only=True) +class SmemVResource(MlaResource): + """V-only view used by validation and resource graph naming.""" + + inst_id: cutlass.Constexpr[int] = 0 + is_v: cutlass.Constexpr[int] = 1 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py new file mode 100644 index 000000000000..f51036748809 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py @@ -0,0 +1,1922 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Correction, normalization, and output-store resource.""" + +from dataclasses import dataclass +from typing import Optional + +from cutlass.experimental import primitives as prims + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int8, Int32, Int64 +from cutlass.experimental import primitives as cprims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import SmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + producer_work, +) + +from ...helpers.constants import ( + SMEM_WORD_BYTE_SHIFT, + SMEM_WORD_BYTES, + TCGEN05_16X256B_REGS_PER_LOAD, + TCGEN05_16X256B_SHAPE, + WARPGROUP_THREADS, +) + + +from ...helpers.layout import ( + _TASK_CACHE_LANE_IDX, + _TASK_CACHE_TMEM_BASE_OFFSET, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _TASK_CACHE_WARP_IDX, + decode_gen_task_cache, + head_dim_cta_offset_v, + local_q_head_idx_for_scale, + num_fp8_output_regs, + num_o_reg_pairs, + num_o_repeats, + num_o_stsm_row_blocks, + num_o_tmem_loads_per_stage, + num_softmax_scale_groups, + o_stage_stsm_and_copy_offsets, + o_stage_tmem_col_offset, + smem_array, +) +from ...helpers.math import ( + ceil_div, + fadd2, + ffma2, + fmul2, + output_dtype, + pack_float2_to_bf16, + partial_output_dtype, +) +from ...helpers.ops import ( + fp8_quant_scale_rcp, + pack_float4_to_fp8_e4m3, + store_transposed_smem8b_x2, + store_transposed_smem8b_x4, + tcgen05_ld_16x32bx2_f32, + tcgen05_panel_addr, + tcgen05_second_panel_addr, + tcgen05_st_16x32bx2_f32, + vector_from_scalars, +) +from ...helpers.query import ( + flat_query_row_state, + public_query_flat_row, + split_o_element_offset, +) +from ...helpers.tile import ( + batch_idx_for_stage_cfg, + cta_idx_head_dim_v_for_stage, + cta_idx_kv_for_stage, + cta_idx_q_for_stage, + head_idx_for_stage, +) + +from .common import ( + MlaResource, +) + +# ===================================================================== +# TmemCorrResource — Correction and output store to GMEM +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemCorrResource(MlaResource): + """Fused throughput-latency 1CTA correction and output-store helper. + + Despite the historical name, this resource does not own TMEM. O is handed + from MmaTask to CorrectionTask through TmemOResource; this helper owns SMEM + scratch/staging and fuses correction, StoreO, and StoreLSE. This should be + cleaned up into clearer correction/store resource boundaries later. + """ + + inst_id: cutlass.Constexpr[int] = 0 + scale_softmax_log2: Float32 = None + output_scale: Float32 = None + o_tensor: object = None + lse_tensor: object = None + acc_o_tensor: object = None + acc_lse_tensor: object = None + cache_seqs: object = None + head_idx: object = None + batch_idx: object = None + cta_idx_q: object = None + cta_idx_kv: object = None + cta_idx_head_dim_v: object = None + store_barrier_id: cutlass.Constexpr[int] = 6 + sum_barrier_id: cutlass.Constexpr[int] = 7 + _sum_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + _cluster_reduction_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + _cluster_reduction_barrier_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + _smem_o: object = None + _smem_o_i32: object = None + _sum_scratch: object = None + _cluster_reduction_smem: object = None + _cluster_reduction_barrier: object = None + + def get_smem_requirements(self): + """Return correction/output SMEM scratch and optional cluster buffers.""" + if self.inst_id != 1: + return [] + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}_oStage", + size_bytes=self.cfg.o_smem_tile_bytes, + alignment=self.cfg.stensor_align, + ) + if self._sum_alloc is None: + self._sum_alloc = SmemAllocation( + name=f"{self.name}_sumScratch", + size_bytes=self.cfg.corr_scratch_bytes, + alignment=16, + ) + allocs = [self._alloc, self._sum_alloc] + if self.cfg.cluster_reduction_smem_bytes and self.acc_o_tensor is None: + if self._cluster_reduction_alloc is None: + self._cluster_reduction_alloc = SmemAllocation( + name=f"{self.name}_clusterReduction", + size_bytes=self.cfg.cluster_reduction_smem_bytes, + alignment=64, + ) + if self._cluster_reduction_barrier_alloc is None: + self._cluster_reduction_barrier_alloc = SmemAllocation( + name=f"{self.name}_clusterReductionBarrier", + size_bytes=8, + alignment=8, + ) + allocs.append(self._cluster_reduction_alloc) + allocs.append(self._cluster_reduction_barrier_alloc) + return allocs + + @cute.jit + def _init_store_state_from_context(self, context) -> None: + """Create correction/output SMEM views and initialize cluster state.""" + if cutlass.const_expr(self.inst_id != 1): + return + self._smem_o = smem_array( + context, + self._alloc, + output_dtype(self.cfg), + self.cfg.o_smem_tile_bytes // self.cfg.o_dtype_bytes, + ) + self._smem_o_i32 = smem_array( + context, + self._alloc, + Int32, + self.cfg.o_smem_tile_bytes // 4, + ) + self._sum_scratch = smem_array( + context, + self._sum_alloc, + Float32, + self.cfg.corr_scratch_bytes // self.cfg.acc_dtype_bytes, + ) + if cutlass.const_expr(self._cluster_reduction_alloc is not None): + self._cluster_reduction_smem = smem_array( + context, + self._cluster_reduction_alloc, + Int8, + self.cfg.cluster_reduction_smem_bytes, + ) + self._cluster_reduction_barrier = smem_array( + context, + self._cluster_reduction_barrier_alloc, + Int64, + 1, + ) + self._init_cluster_reduction_barrier() + + @cute.jit + def initialize_runtime_state_internal( + self, + context=None, + captured_schedule: cutlass.Constexpr[bool] = False, + ) -> None: + """Initialize correction/store SMEM views outside captured task bodies.""" + super().initialize_runtime_state_internal(context, captured_schedule) + # cluster state is initialized explicitly by the kernel entry before the + # cluster rendezvous. This guarantees that every remote transaction + # barrier is live before any peer can publish a partial. Non-cluster + # schedules retain the ordinary task-owned initialization path. + if cutlass.const_expr(self.cfg.use_cluster_reduction != 1): + self._init_store_state_from_context(context) + + @cute.jit + def create_cluster_function_variables(self, context) -> None: + """Bind cluster correction SMEM and initialize its transaction barrier.""" + + if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): + self._init_store_state_from_context(context) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_store_state(self, stage_info: StageInfo) -> None: + """Initialize correction/store task-local state for captured schedules.""" + # Producer aux work initializes the correction task's view of TMEM O. + # Runtime SMEM views are initialized once in initialize_runtime_state_internal(). + self._init_tmem_state(stage_info) + + @cute.jit + def _o_base_col(self): + """Return the first TMEM column owned by the O accumulator.""" + + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + return Int32(2 * self.cfg.tmem_s_cols) + return Int32(2 * self.cfg.tmem_s_cols + 2 * self.cfg.tmem_stats_cols) + + @cute.jit + def _init_cluster_reduction_barrier(self): + """Initialize the transaction barrier for this CTA's reduction slice.""" + if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): + cfg = self.cfg + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + rows_per_slice = cfg.cluster_reduction_rows_per_slice + num_slices_per_cta = ceil_div( + cfg.cluster_reduction_slices, cfg.num_ctas_per_seq_kv + ) + rows_per_cta = num_slices_per_cta * rows_per_slice + first_row = cta_rank * Int32(rows_per_cta) + remaining_rows = Int32(cfg.tile_size_q) - first_row + valid_rows = cute.math.max( + Int32(0), + cute.math.min(Int32(rows_per_cta), remaining_rows), + ) + bytes_per_row = Int32( + cfg.head_dim_per_cta_v * cfg.partial_o_dtype_bytes + cfg.acc_dtype_bytes + ) + expected_bytes = valid_rows * Int32(cfg.num_ctas_per_seq_kv) * bytes_per_row + if warp_idx == Int32(cfg.correction_warp_idx): + with cute.arch.elect_one(): + prims.mbarrier_init(self._cluster_reduction_barrier, 1) + prims.mbarrier_arrive_expect_tx( + self._cluster_reduction_barrier, expected_bytes + ) + + @cute.jit + def _cluster_rows_per_cta(self): + """Return how many Q rows this CTA owns during cluster reduction.""" + + num_slices_per_cta = ceil_div( + self.cfg.cluster_reduction_slices, self.cfg.num_ctas_per_seq_kv + ) + return Int32(num_slices_per_cta * self.cfg.cluster_reduction_rows_per_slice) + + @cute.jit + def _cluster_o_bytes(self): + """Return the O bytes contributed by all split CTAs for this owner.""" + + rows_per_cta = self._cluster_rows_per_cta() + return Int64( + Int32(self.cfg.num_ctas_per_seq_kv) + * rows_per_cta + * Int32(self.cfg.head_dim_per_cta_v * self.cfg.partial_o_dtype_bytes) + ) + + @cute.jit + def _cluster_remote_smem_ptr(self, owner_rank, byte_offset, dtype): + """Map a local cluster SMEM address to the CTA that owns the row slice.""" + + local_ptr = self._cluster_local_smem_ptr(byte_offset, dtype) + return prims.mapa(local_ptr, owner_rank) + + @cute.jit + def _cluster_local_smem_ptr(self, byte_offset, dtype): + """Return a typed pointer into this CTA's cluster reduction SMEM buffer.""" + + return cutlass.inttoptr( + self._cluster_reduction_smem.data_ptr().toint(Int64) + Int64(byte_offset), + mem_space=3, + dtype=dtype, + ) + + @cute.jit + def _store_partial_o_to_cluster_smem( + self, + local_row_idx, + cta_idx_kv, + local_head_dim_elem_offset, + local_head_dim_byte_offset, + partial_vec, + ): + """Send one partial O vector to the CTA that owns its reduction slice.""" + rows_per_cta = self._cluster_rows_per_cta() + owner_rank = local_row_idx // rows_per_cta + owner_row_idx = local_row_idx - owner_rank * rows_per_cta + elem_offset = ( + cta_idx_kv * rows_per_cta * Int32(self.cfg.head_dim_per_cta_v) + + owner_row_idx * Int32(self.cfg.head_dim_per_cta_v) + + local_head_dim_elem_offset + ) + byte_offset = Int64( + elem_offset * Int32(self.cfg.partial_o_dtype_bytes) + ) + Int64(local_head_dim_byte_offset) + remote_ptr = self._cluster_remote_smem_ptr(owner_rank, byte_offset, Int32) + remote_barrier = prims.mapa(self._cluster_reduction_barrier, owner_rank) + # Keep the vectorized publication on the public inline-PTX API so the + # operation remains one 16-byte store instead of four scalar stores. + cute.arch.inline_ptx( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[{$r0}], {{$r1}, {$r2}, {$r3}, {$r4}}, [{$r5}];", + read_only_args=[ + remote_ptr.ir_value(), + partial_vec[0], + partial_vec[1], + partial_vec[2], + partial_vec[3], + remote_barrier.ir_value(), + ], + ) + + @cute.jit + def _store_partial_lse_to_cluster_smem(self, local_row_idx, cta_idx_kv, lse_val): + """Send one partial log-sum-exp value to the matching reduction owner.""" + rows_per_cta = self._cluster_rows_per_cta() + owner_rank = local_row_idx // rows_per_cta + owner_row_idx = local_row_idx - owner_rank * rows_per_cta + stats_elem_offset = cta_idx_kv * rows_per_cta * Int32( + 2 + ) + owner_row_idx * Int32(2) + byte_offset = self._cluster_o_bytes() + Int64( + stats_elem_offset * Int32(self.cfg.acc_dtype_bytes) + ) + remote_ptr = self._cluster_remote_smem_ptr(owner_rank, byte_offset, Int32) + remote_barrier = prims.mapa(self._cluster_reduction_barrier, owner_rank) + cute.arch.inline_ptx( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 " + "[{$r0}], {$r1}, [{$r2}];", + read_only_args=[ + remote_ptr.ir_value(), + lse_val.bitcast(Int32), + remote_barrier.ir_value(), + ], + ) + + @cute.jit + def _store_o_slice_to_gmem( + self, + stage_info: StageInfo, + task_cache, + v_stage_idx: int, + head_dim_offset: Int32, + ): + """Store one corrected O head-dim slice to final or split-KV output.""" + if cutlass.const_expr(self.o_tensor is None and self.acc_o_tensor is None): + return + + cfg = self.cfg + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + head_base_idx = head_idx_for_stage(self.head_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + num_copy_segments = cfg.o_copy_segments_per_stage + for copy_segment_idx in cutlass.range_constexpr(num_copy_segments): + _, load_smem_offset, dst_row_idx, dst_col_offset = ( + o_stage_stsm_and_copy_offsets( + cfg, + warp_grp_thread_idx, + warp_idx, + lane_idx, + 0, + copy_segment_idx, + ) + ) + global_head_idx = head_base_idx + dst_row_idx + if dst_row_idx < Int32(cfg.tile_size_q) and global_head_idx < Int32( + cfg.num_heads_q + ): + storage_flat_query_row, _, _, _, valid_output_row = ( + flat_query_row_state( + global_head_idx, + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=self.cu_seqlens_q, + batch_idx=batch_idx, + ) + ) + smem_src = self._smem_o_i32.data_ptr( + load_smem_offset >> SMEM_WORD_BYTE_SHIFT + ) + if cutlass.const_expr(cfg.use_cluster_reduction == 1): + self._store_partial_o_to_cluster_smem( + dst_row_idx, + cta_idx_kv, + Int32(v_stage_idx * cfg.head_dim_per_stage_v), + dst_col_offset, + smem_src.load(count=4, alignment=16), + ) + elif valid_output_row: + if cutlass.const_expr(self.acc_o_tensor is not None): + base_elem_offset = split_o_element_offset( + cfg, + batch_idx, + cta_idx_q, + global_head_idx, + cta_idx_kv, + head_dim_offset + + Int32(v_stage_idx * cfg.head_dim_per_stage_v), + ) + base_ptr = self.acc_o_tensor.iterator.raw_ptr().toint(Int64) + element_bytes = cfg.partial_o_dtype_bytes + else: + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + self.cu_seqlens_q, + ) + base_elem_offset = ( + Int64(output_query_row) * Int64(cfg.head_dim_v) + + Int64(head_dim_offset) + + Int64(v_stage_idx * cfg.head_dim_per_stage_v) + ) + base_ptr = self.o_tensor.iterator.raw_ptr().toint(Int64) + element_bytes = cfg.o_dtype_bytes + byte_offset = Int64(base_elem_offset) * Int64( + element_bytes + ) + Int64(dst_col_offset) + dst_ptr = cutlass.inttoptr( + base_ptr + byte_offset, + mem_space=1, + dtype=Int32, + ) + dst_ptr.store( + smem_src.load(count=4, alignment=16), + alignment=16, + ) + + @cute.jit + def _store_lse_to_gmem( + self, stage_info: StageInfo, task_cache, final_max, reduced_sum + ): + """Store final LSE or split-KV partial LSE for one corrected tile.""" + if cutlass.const_expr(self.lse_tensor is None and self.acc_lse_tensor is None): + return + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, self.cfg, stage_info) + head_base_idx = head_idx_for_stage(self.head_idx, self.cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + cta_idx_head_dim_v = cta_idx_head_dim_v_for_stage( + self.cta_idx_head_dim_v, stage_info + ) + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + should_store_lse = lane_idx < Int32(16) + local_row_idx = warp_idx * Int32(16) + (lane_idx & Int32(0xF)) + head_idx = head_base_idx + local_row_idx + scale_idx = Int32(0) + else: + should_store_lse = warp_idx == Int32(0) and lane_idx < Int32( + 4 * num_softmax_scale_groups(self.cfg) + ) + col_group_idx = lane_idx & Int32(0x3) + scale_idx = lane_idx >> Int32(2) + local_row_idx = local_q_head_idx_for_scale( + self.cfg, col_group_idx, scale_idx + ) + head_idx = head_base_idx + local_row_idx + if should_store_lse: + # Swaps keeps at least four scale groups, so TileQ8 has eight + # padded rows in the correction register footprint. Do not let + # those rows publish split LSE into the next logical head tile. + if local_row_idx < Int32(self.cfg.tile_size_q) and head_idx < Int32( + self.cfg.num_heads_q + ): + storage_flat_query_row, _, _, _, valid_output_row = ( + flat_query_row_state( + head_idx, + cta_idx_q, + self.cfg.tile_size_q, + self.cfg.logical_num_heads_q, + self.cfg.logical_seq_len_q, + cu_seqlens_q=self.cu_seqlens_q, + batch_idx=batch_idx, + ) + ) + lse_sum = reduced_sum[scale_idx] + if cutlass.const_expr(self.cfg.is_fp8_qkv()): + # FP8 P is quantized as 448 * softmax(P) for BMM2. The + # online sum tracks the same scale, so undo it for the + # externally visible log-sum-exp value. + lse_sum = lse_sum * fp8_quant_scale_rcp() + lse_val = ( + cute.math.log2(lse_sum, fastmath=True) + + self.scale_softmax_log2 * final_max[scale_idx] + ) + if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): + self._store_partial_lse_to_cluster_smem( + local_row_idx, cta_idx_kv, lse_val + ) + elif cutlass.const_expr(self.acc_lse_tensor is not None): + if valid_output_row and cta_idx_head_dim_v == Int32(0): + split_kv = Int32(self.cfg.num_ctas_per_seq_kv) + elem_offset = ( + batch_idx + * Int32( + self.cfg.seq_len_q * self.cfg.num_heads_q * split_kv + ) + + cta_idx_q * Int32(self.cfg.num_heads_q * split_kv) + + head_idx * split_kv + + cta_idx_kv + ) + (self.acc_lse_tensor.iterator.raw_ptr() + elem_offset).store( + lse_val + ) + else: + if valid_output_row: + elem_offset = public_query_flat_row( + self.cfg, + storage_flat_query_row, + batch_idx, + self.cu_seqlens_q, + ) + (self.lse_tensor.iterator.raw_ptr() + elem_offset).store( + lse_val + ) + + @cute.jit + def _cluster_lse_byte_offset(self, split_idx, owner_row_idx): + rows_per_cta = self._cluster_rows_per_cta() + stats_elem_offset = split_idx * rows_per_cta * Int32(2) + owner_row_idx * Int32( + 2 + ) + return self._cluster_o_bytes() + Int64( + stats_elem_offset * Int32(self.cfg.acc_dtype_bytes) + ) + + @cute.jit + def _cluster_o_byte_offset(self, split_idx, owner_row_idx, dim_idx): + rows_per_cta = self._cluster_rows_per_cta() + elem_offset = ( + split_idx * rows_per_cta * Int32(self.cfg.head_dim_per_cta_v) + + owner_row_idx * Int32(self.cfg.head_dim_per_cta_v) + + dim_idx + ) + return Int64(elem_offset * Int32(self.cfg.partial_o_dtype_bytes)) + + @cute.jit + def _cluster_wait_transaction_barrier(self, warp_grp_thread_idx): + """Wait for peer cluster stores before the owner reads local DSMEM.""" + + # Give every correction lane one nonblocking acquire attempt. If the + # transaction is not already complete, only lane 0 keeps polling. The + # correction store barrier is free after the final O-staging phase and + # publishes lane 0's ready point to all 128 correction lanes before + # any of them reads the distributed-SMEM partials below. + cluster_transaction_ready = prims.mbarrier_try_wait_parity( + self._cluster_reduction_barrier, 0, time_limit=0 + ) + if warp_grp_thread_idx == Int32(0): + while not cluster_transaction_ready: + cluster_transaction_ready = prims.mbarrier_try_wait_parity( + self._cluster_reduction_barrier, 0, time_limit=10_000_000 + ) + prims.barrier_cta_sync( + barrier_id=self.store_barrier_id, + thread_count=WARPGROUP_THREADS, + ) + + @cute.jit + def publish_neutral_cluster_partial_and_reduce( + self, + batch_idx, + head_base_idx, + cta_idx_q, + cta_idx_kv, + cta_idx_head_dim_v, + ): + """Publish an empty split and execute this rank's static reduction work. + + Runtime split pruning only removes the K/V task graph. Every launched + cluster rank still contributes the configured partial byte count so owner + barriers, DSMEM offsets, and row ownership remain compile-time static. + The neutral pair ``O=0, LSE=-inf`` has no effect on online-softmax + reduction and is safe for mutable runtime sequence lengths. + """ + + if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): + cfg = self.cfg + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + is_correction_warp = (warp_idx >= Int32(cfg.correction_warp_idx)) & ( + warp_idx < Int32(cfg.correction_warp_idx + cfg.correction_num_warps) + ) + if is_correction_warp: + thread_idx, _, _ = cute.arch.thread_idx() + warp_grp_thread_idx = thread_idx - Int32(cfg.correction_warp_idx * 32) + base_vec_offset = warp_grp_thread_idx * Int32(8) + row_in_slice = base_vec_offset // Int32(cfg.head_dim_per_cta_v) + dim_idx = base_vec_offset - row_in_slice * Int32(cfg.head_dim_per_cta_v) + neutral_o = vector_from_scalars( + (Int32(0), Int32(0), Int32(0), Int32(0)), + dtype=Int32, + ) + for slice_idx in cutlass.range_constexpr(cfg.cluster_reduction_slices): + local_row_idx = ( + Int32(slice_idx * cfg.cluster_reduction_rows_per_slice) + + row_in_slice + ) + global_head_idx = head_base_idx + local_row_idx + if local_row_idx < Int32( + cfg.tile_size_q + ) and global_head_idx < Int32(cfg.num_heads_q): + self._store_partial_o_to_cluster_smem( + local_row_idx, + cta_idx_kv, + dim_idx, + Int32(0), + neutral_o, + ) + if dim_idx == Int32(0): + self._store_partial_lse_to_cluster_smem( + local_row_idx, + cta_idx_kv, + Float32(-Float32.inf), + ) + + head_dim_offset = head_dim_cta_offset_v(cfg, cta_idx_head_dim_v) + self._reduce_cluster_partials_and_store_for_tile( + batch_idx, + head_base_idx, + cta_idx_q, + cta_idx_kv, + cta_idx_head_dim_v, + warp_grp_thread_idx, + head_dim_offset, + ) + + @cute.jit + def _reduce_cluster_partials_and_store( + self, stage_info: StageInfo, task_cache, head_dim_offset + ): + """Reduce DSMEM split-KV partials for this CTA rank's row slices.""" + if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): + cfg = self.cfg + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + head_base_idx = head_idx_for_stage(self.head_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + cta_idx_head_dim_v = cta_idx_head_dim_v_for_stage( + self.cta_idx_head_dim_v, stage_info + ) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + self._reduce_cluster_partials_and_store_for_tile( + batch_idx, + head_base_idx, + cta_idx_q, + cta_idx_kv, + cta_idx_head_dim_v, + warp_grp_thread_idx, + head_dim_offset, + ) + + @cute.jit + def _reduce_cluster_partials_and_store_for_tile( + self, + batch_idx, + head_base_idx, + cta_idx_q, + cta_idx_kv, + cta_idx_head_dim_v, + warp_grp_thread_idx, + head_dim_offset, + ): + """Reduce one configured-S cluster row-owner slice.""" + if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): + cfg = self.cfg + rows_per_slice = cfg.cluster_reduction_rows_per_slice + num_slices_per_cta = ceil_div( + cfg.cluster_reduction_slices, cfg.num_ctas_per_seq_kv + ) + rows_per_cta = num_slices_per_cta * rows_per_slice + owner_first_row = cta_idx_kv * Int32(rows_per_cta) + valid_rows = cute.math.max( + Int32(0), + cute.math.min( + Int32(rows_per_cta), + Int32(cfg.tile_size_q) - owner_first_row, + ), + ) + + if valid_rows > Int32(0): + self._cluster_wait_transaction_barrier(warp_grp_thread_idx) + + base_vec_offset = warp_grp_thread_idx * Int32(8) + row_in_slice = base_vec_offset // Int32(cfg.head_dim_per_cta_v) + dim_idx = base_vec_offset - row_in_slice * Int32(cfg.head_dim_per_cta_v) + output_element_dtype = output_dtype(cfg) + partial_element_dtype = partial_output_dtype(cfg) + + for slice_offset in cutlass.range_constexpr(num_slices_per_cta): + owner_row_idx = Int32(slice_offset * rows_per_slice) + row_in_slice + local_row_idx = owner_first_row + owner_row_idx + if local_row_idx < Int32(cfg.tile_size_q): + global_head_idx = head_base_idx + local_row_idx + if global_head_idx < Int32(cfg.num_heads_q): + ( + storage_flat_query_row, + _, + _, + _, + valid_output_row, + ) = flat_query_row_state( + global_head_idx, + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=self.cu_seqlens_q, + batch_idx=batch_idx, + ) + lse_max = Float32(-Float32.inf) + local_lse = cutlass.Array( + Float32, + cfg.num_ctas_per_seq_kv, + space=cutlass.AddressSpace.rmem, + ) + for split_idx in cutlass.range_constexpr( + cfg.num_ctas_per_seq_kv + ): + lse_ptr = self._cluster_local_smem_ptr( + self._cluster_lse_byte_offset( + Int32(split_idx), owner_row_idx + ), + Int32, + ) + lse_val = lse_ptr.load().bitcast(Float32) + local_lse[split_idx] = lse_val + lse_max = cute.math.max(lse_max, lse_val, ftz=True) + + lse_max = ( + lse_max + if lse_max != Float32(-Float32.inf) + else Float32(0.0) + ) + lse_sum = Float32(0.0) + for split_idx in cutlass.range_constexpr( + cfg.num_ctas_per_seq_kv + ): + lse_sum += cute.math.exp2( + local_lse[split_idx] - lse_max, fastmath=True + ) + has_finite_mass = ( + lse_sum != Float32(0.0) and lse_sum == lse_sum + ) + is_fully_masked = lse_sum == Float32(0.0) + global_lse = ( + lse_max + cute.math.log2(lse_sum, fastmath=True) + if has_finite_mass + else Float32(Float32.inf) + ) + published_lse = ( + Float32(-Float32.inf) if is_fully_masked else global_lse + ) + + if ( + valid_output_row + and dim_idx == Int32(0) + and cta_idx_head_dim_v == Int32(0) + ): + if cutlass.const_expr(self.lse_tensor is not None): + lse_offset = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + self.cu_seqlens_q, + ) + ( + self.lse_tensor.iterator.raw_ptr() + lse_offset + ).store(published_lse) + + acc_vec = vector_from_scalars( + ( + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + ), + dtype=Float32, + ) + for split_idx in cutlass.range_constexpr( + cfg.num_ctas_per_seq_kv + ): + scale = cute.math.exp2( + local_lse[split_idx] - global_lse, fastmath=True + ) + partial_ptr = self._cluster_local_smem_ptr( + self._cluster_o_byte_offset( + Int32(split_idx), owner_row_idx, dim_idx + ), + partial_element_dtype, + ) + partial_vec = partial_ptr.load( + count=8, alignment=16 + ).to(Float32) + acc_vec = acc_vec + partial_vec * scale + + if cutlass.const_expr(self.o_tensor is not None): + if valid_output_row: + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + self.cu_seqlens_q, + ) + out_elem_offset = ( + Int64(output_query_row) * Int64(cfg.head_dim_v) + + Int64(head_dim_offset) + + Int64(dim_idx) + ) + out_byte_offset = out_elem_offset * Int64( + cfg.o_dtype_bytes + ) + if cutlass.const_expr(cfg.use_fp8_output == 1): + packed_o = cutlass.Array( + Int32, 2, space=cutlass.AddressSpace.rmem + ) + packed_o[0] = pack_float4_to_fp8_e4m3( + acc_vec[0], + acc_vec[1], + acc_vec[2], + acc_vec[3], + ) + packed_o[1] = pack_float4_to_fp8_e4m3( + acc_vec[4], + acc_vec[5], + acc_vec[6], + acc_vec[7], + ) + out_ptr = cutlass.inttoptr( + self.o_tensor.iterator.raw_ptr().toint( + Int64 + ) + + out_byte_offset, + mem_space=1, + dtype=Int32, + ) + out_ptr.store(packed_o.load(0, 2), alignment=8) + else: + out_ptr = cutlass.inttoptr( + self.o_tensor.iterator.raw_ptr().toint( + Int64 + ) + + out_byte_offset, + mem_space=1, + dtype=output_element_dtype, + ) + out_ptr.store( + acc_vec.to(output_element_dtype), + alignment=16, + ) + + @producer_work + @cute.jit + def correct_loop_and_store( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_new_max_arr, + inst1_sum_arr, + o_stage_idx, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ): + """Rescale the running O accumulator for a loop-stage softmax update.""" + del sum_arr, inst0_new_max_arr, inst0_sum_arr + del inst1_new_max_arr, inst1_sum_arr + del tail_o_stage_idx_0, tail_o_stage_idx_1 + self._rescale_loop_o( + stage_info, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + o_stage_idx=o_stage_idx, + ) + + @producer_work + @cute.jit + def correct_tail_and_store( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_new_max_arr, + inst1_sum_arr, + o_stage_idx, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ): + """Normalize final O, then store O/LSE or split-KV partials.""" + del old_max_arr, o_stage_idx + self._normalize_and_store_tail_o( + stage_info, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + ) + + @cute.jit + def _apply_loop_correction_in_tmem( + self, + stage_info: StageInfo, + *, + scale_vals, + o_stage_idx, + ): + """Apply one nonidentity loop correction to the live TMEM O stage.""" + cfg = self.cfg + task_cache = decode_gen_task_cache(stage_info) + tmem_row_base = task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + shape = TCGEN05_16X256B_SHAPE + q_repeats = num_o_repeats(cfg) + o_reg_pair_count = num_o_reg_pairs(cfg) + num_o_tmem_loads = num_o_tmem_loads_per_stage(cfg) + + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + for v_stage_idx in cutlass.range_constexpr(cfg.v_head_dim_stages): + base_addr = ( + tmem_row_base + + self._o_base_col() + + o_stage_tmem_col_offset(cfg, o_stage_idx, v_stage_idx) + ) + loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr, Float32), + num=cfg.head_dim_per_stage_v // 2, + offset=Int32(cfg.head_dim_per_stage_v // 2), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + scaled = cutlass.Array( + Float32, + cfg.head_dim_per_stage_v // 2, + space=cutlass.AddressSpace.rmem, + ) + for reg_idx in cutlass.range_constexpr(cfg.head_dim_per_stage_v // 2): + scaled[reg_idx] = loaded[reg_idx] * scale_vals[0] + scaled_vec = vector_from_scalars( + tuple( + scaled[reg_idx] + for reg_idx in range(cfg.head_dim_per_stage_v // 2) + ), + Float32, + ) + tcgen05_st_16x32bx2_f32( + prims.make_tmem_ptr(base_addr, Float32), + scaled_vec, + offset=Int32(cfg.head_dim_per_stage_v // 2), + ) + return + + for v_stage_idx in cutlass.range_constexpr(cfg.v_head_dim_stages): + base_addr = ( + tmem_row_base + + self._o_base_col() + + o_stage_tmem_col_offset(cfg, o_stage_idx, v_stage_idx) + ) + for chunk_idx in cutlass.range_constexpr(num_o_tmem_loads): + if cutlass.const_expr(cfg.tile_size_q == 8): + loaded_lo = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr, chunk_idx), Float32 + ), + num=1, + ) + loaded_hi = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_second_panel_addr( + tcgen05_panel_addr(base_addr, chunk_idx) + ), + Float32, + ), + num=1, + ) + else: + loaded = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr, chunk_idx), Float32 + ), + num=q_repeats, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + + if cutlass.const_expr(cfg.tile_size_q == 8): + scaled_lo = cutlass.Array( + Float32, 4, space=cutlass.AddressSpace.rmem + ) + scaled_hi = cutlass.Array( + Float32, 4, space=cutlass.AddressSpace.rmem + ) + for reg_pair_idx in cutlass.range_constexpr(o_reg_pair_count): + scale_base = ( + (reg_pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2 + ) * 2 + reg_base = (reg_pair_idx % 2) * 2 + if cutlass.const_expr(reg_pair_idx < 2): + scaled_lo[reg_base] = ( + loaded_lo[reg_base] * scale_vals[scale_base] + ) + scaled_lo[reg_base + 1] = ( + loaded_lo[reg_base + 1] * scale_vals[scale_base + 1] + ) + else: + scaled_hi[reg_base] = ( + loaded_hi[reg_base] * scale_vals[scale_base] + ) + scaled_hi[reg_base + 1] = ( + loaded_hi[reg_base + 1] * scale_vals[scale_base + 1] + ) + scaled_lo_vec = vector_from_scalars( + (scaled_lo[0], scaled_lo[1], scaled_lo[2], scaled_lo[3]), + Float32, + ) + scaled_hi_vec = vector_from_scalars( + (scaled_hi[0], scaled_hi[1], scaled_hi[2], scaled_hi[3]), + Float32, + ) + prims.tcgen05_st( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr, chunk_idx), Float32 + ), + scaled_lo_vec, + ) + prims.tcgen05_st( + shape, + prims.make_tmem_ptr( + tcgen05_second_panel_addr( + tcgen05_panel_addr(base_addr, chunk_idx) + ), + Float32, + ), + scaled_hi_vec, + ) + else: + scaled = cutlass.Array( + Float32, 4 * q_repeats, space=cutlass.AddressSpace.rmem + ) + for reg_pair_idx in cutlass.range_constexpr(o_reg_pair_count): + scale_base = ( + (reg_pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2 + ) * 2 + reg_base = reg_pair_idx * 2 + scaled[reg_base] = loaded[reg_base] * scale_vals[scale_base] + scaled[reg_base + 1] = ( + loaded[reg_base + 1] * scale_vals[scale_base + 1] + ) + scaled_vec = vector_from_scalars( + tuple(scaled[reg_idx] for reg_idx in range(4 * q_repeats)), + Float32, + ) + prims.tcgen05_st( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr, chunk_idx), Float32 + ), + scaled_vec, + ) + return + + @cute.jit + def _rescale_loop_o( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + o_stage_idx, + ): + """Rescale the running O accumulator after a loop softmax update.""" + cfg = self.cfg + num_scale_groups = num_softmax_scale_groups(cfg) + + # A loop correction is exactly identity when its row maximum did not + # change. Avoid exp2 for identity lanes, and skip collective TMEM + # traffic only when the entire correction warp agrees. + scale_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + old_max = old_max_arr[0] + new_max = new_max_arr[0] + scale_is_identity = old_max == new_max + scale_vals[0] = Float32(1.0) + if not scale_is_identity: + scale_vals[0] = cute.math.exp2( + self.scale_softmax_log2 * (old_max - new_max), + fastmath=True, + ) + lane_scales_are_identity = scale_is_identity + else: + lane_scales_are_identity = cutlass.Boolean(True) + for scale_base in cutlass.range_constexpr(0, num_scale_groups, 2): + old_max_0 = old_max_arr[scale_base] + old_max_1 = old_max_arr[scale_base + 1] + new_max_0 = new_max_arr[scale_base] + new_max_1 = new_max_arr[scale_base + 1] + scale_0_is_identity = old_max_0 == new_max_0 + scale_1_is_identity = old_max_1 == new_max_1 + scale_vals[scale_base] = Float32(1.0) + scale_vals[scale_base + 1] = Float32(1.0) + max_diff_pair = fadd2( + (old_max_0, old_max_1), + (-new_max_0, -new_max_1), + ) + scale_pair = fmul2( + (self.scale_softmax_log2, self.scale_softmax_log2), + max_diff_pair, + ) + if not scale_0_is_identity: + scale_vals[scale_base] = cute.math.exp2( + scale_pair[0], fastmath=True + ) + if not scale_1_is_identity: + scale_vals[scale_base + 1] = cute.math.exp2( + scale_pair[1], fastmath=True + ) + lane_scales_are_identity = ( + lane_scales_are_identity & scale_0_is_identity & scale_1_is_identity + ) + + skip_rescale = prims.vote_sync( + cute.arch.FULL_MASK, + lane_scales_are_identity, + prims.VoteSync.ALL, + ) + if not skip_rescale: + self._apply_loop_correction_in_tmem( + stage_info, + scale_vals=scale_vals, + o_stage_idx=o_stage_idx, + ) + # Preserve the correction task's TMEM ordering point even when this + # warp had no rescale transaction to issue. + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + if not skip_rescale: + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def _compute_tail_softmax_scales( + self, + task_cache, + *, + new_max_arr, + sum_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_new_max_arr, + inst1_sum_arr, + ): + """Compute final max/sum values and O normalization scales.""" + cfg = self.cfg + num_scale_groups = num_softmax_scale_groups(cfg) + final_max = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + final_sum = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + exp_scale0 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + exp_scale1 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + reduced_sum = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + final_scale0 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + final_scale1 = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + # Four column groups share one softmax row through separate lane + # subsets. The low two thread bits select the group-local scratch lane. + col_group_idx = warp_grp_thread_idx & Int32(0x3) + + for scale_idx in cutlass.range_constexpr(num_scale_groups): + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + inst_sum = sum_arr[scale_idx] + inst_max = new_max_arr[scale_idx] + final_max[scale_idx] = inst_max + exp_scale0[scale_idx] = Float32(1.0) + exp_scale1[scale_idx] = Float32(0.0) + final_sum[scale_idx] = inst_sum + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=16, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) + ) + reduced_sum[scale_idx] = final_sum[scale_idx] + else: + inst0_sum = inst0_sum_arr[scale_idx] + inst1_sum = inst1_sum_arr[scale_idx] + inst0_max = inst0_new_max_arr[scale_idx] + inst1_max = inst1_new_max_arr[scale_idx] + final_max[scale_idx] = cute.math.max(inst0_max, inst1_max, ftz=True) + exp_scale0[scale_idx] = cute.math.exp2( + self.scale_softmax_log2 * (inst0_max - final_max[scale_idx]), + fastmath=True, + ) + exp_scale1[scale_idx] = cute.math.exp2( + self.scale_softmax_log2 * (inst1_max - final_max[scale_idx]), + fastmath=True, + ) + final_sum[scale_idx] = ( + inst0_sum * exp_scale0[scale_idx] + + inst1_sum * exp_scale1[scale_idx] + ) + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=16, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) + ) + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=8, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) + ) + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=4, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) + ) + + if not cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + # The two softmax pipes first reduce within their owning warp, then + # publish one value per column group so the warpgroup can form the + # final normalization denominator before O is stored. + warp_store_base = warp_idx * Int32( + 4 * num_scale_groups + ) + col_group_idx * Int32(num_scale_groups) + if lane_idx < Int32(4): + self._sum_scratch.store( + tuple( + final_sum[scale_idx] for scale_idx in range(num_scale_groups) + ), + warp_store_base, + alignment=16 if num_scale_groups >= 4 else 8, + ) + prims.barrier_cta_sync( + barrier_id=self.sum_barrier_id, thread_count=WARPGROUP_THREADS + ) + + reduce_base = col_group_idx * Int32(num_scale_groups) + reduced_vec = self._sum_scratch.load( + reduce_base, + vector_size=num_scale_groups, + alignment=16 if num_scale_groups == 4 else 8, + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + reduced_sum[scale_idx] = reduced_vec[scale_idx] + for warp_offset in cutlass.range_constexpr(1, 4): + other_vec = self._sum_scratch.load( + reduce_base + warp_offset * Int32(4 * num_scale_groups), + vector_size=num_scale_groups, + alignment=16 if num_scale_groups == 4 else 8, + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + reduced_sum[scale_idx] += other_vec[scale_idx] + + for scale_idx in cutlass.range_constexpr(num_scale_groups): + norm_scale = Float32(0.0) + if reduced_sum[scale_idx] != Float32(0.0): + norm_scale = self.output_scale / reduced_sum[scale_idx] + final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] + final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] + return final_max, reduced_sum, final_scale0, final_scale1 + + @cute.jit + def _store_keeps_tail_o_stage( + self, + stage_info: StageInfo, + task_cache, + *, + base_addr0, + final_scale0, + v_stage_idx: int, + head_dim_offset, + ): + """Normalize and directly store one keeps-MMA-AB O stage.""" + cfg = self.cfg + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + head_base_idx = head_idx_for_stage(self.head_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + # Keeps-MMA-AB stores one head row per half-warp. The low 16 lanes own + # the row index; the upper lane bit selects the head-dim half. + local_row_idx = warp_idx * Int32(16) + (lane_idx & Int32(0xF)) + global_head_idx = head_base_idx + local_row_idx + storage_flat_query_row, _, _, _, valid_output_row = flat_query_row_state( + global_head_idx, + cta_idx_q, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + cu_seqlens_q=self.cu_seqlens_q, + batch_idx=batch_idx, + ) + half_warp_col_offset = (lane_idx >> Int32(4)) * Int32( + cfg.head_dim_per_stage_v // 2 + ) + + # ``tcgen05.ld.sync.aligned`` must remain convergent across the warp. + # Padded query rows still own valid physical TMEM rows, so load and + # normalize every row and predicate only the externally visible store. + for chunk_idx in cutlass.range_constexpr(0, cfg.head_dim_per_stage_v // 2, 8): + o_loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr0 + Int32(chunk_idx), Float32), + num=8, + offset=Int32(cfg.head_dim_per_stage_v // 2), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + final_vals = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) + packed_o = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + else: + packed_o = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for pair_idx in cutlass.range_constexpr(4): + reg_base = pair_idx * 2 + final_pair = fmul2( + (final_scale0[0], final_scale0[0]), + (o_loaded[reg_base], o_loaded[reg_base + 1]), + ) + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + final_vals[reg_base] = final_pair[0] + final_vals[reg_base + 1] = final_pair[1] + else: + packed_o[pair_idx] = pack_float2_to_bf16( + final_pair[0], final_pair[1] + ) + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + packed_o[0] = pack_float4_to_fp8_e4m3( + final_vals[0], + final_vals[1], + final_vals[2], + final_vals[3], + ) + packed_o[1] = pack_float4_to_fp8_e4m3( + final_vals[4], + final_vals[5], + final_vals[6], + final_vals[7], + ) + store_col_offset = ( + head_dim_offset + + Int32(v_stage_idx * cfg.head_dim_per_stage_v) + + half_warp_col_offset + + Int32(chunk_idx) + ) + if global_head_idx < Int32(cfg.num_heads_q) and valid_output_row: + if cutlass.const_expr(self.acc_o_tensor is not None): + base_elem_offset = split_o_element_offset( + cfg, + batch_idx, + cta_idx_q, + global_head_idx, + cta_idx_kv, + store_col_offset, + ) + base_ptr = self.acc_o_tensor.iterator.raw_ptr().toint(Int64) + element_bytes = cfg.partial_o_dtype_bytes + else: + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + self.cu_seqlens_q, + ) + base_elem_offset = Int64(output_query_row) * Int64( + cfg.head_dim_v + ) + Int64(store_col_offset) + base_ptr = self.o_tensor.iterator.raw_ptr().toint(Int64) + element_bytes = cfg.o_dtype_bytes + dst_ptr = cutlass.inttoptr( + base_ptr + Int64(base_elem_offset) * Int64(element_bytes), + mem_space=1, + dtype=Int32, + ) + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + dst_ptr.store(packed_o.load(0, 2), alignment=8) + else: + dst_ptr.store(packed_o.load(0, 4), alignment=16) + + @cute.jit + def _stage_tile32_tail_o_to_smem( + self, + task_cache, + *, + base_addr0, + base_addr1, + final_scale0, + final_scale1, + ): + """Stage the tileSizeQ=32 BF16/partial O tail path through SMEM.""" + cfg = self.cfg + shape = TCGEN05_16X256B_SHAPE + q_repeats = num_o_repeats(cfg) + o_reg_pair_count = num_o_reg_pairs(cfg) + num_o_tmem_loads = num_o_tmem_loads_per_stage(cfg) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + + for chunk_idx in cutlass.range_constexpr(num_o_tmem_loads): + o0_loaded = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr(tcgen05_panel_addr(base_addr0, chunk_idx), Float32), + num=q_repeats, + ) + o1_loaded = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr(tcgen05_panel_addr(base_addr1, chunk_idx), Float32), + num=q_repeats, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + regs_o = cutlass.Array( + Int32, + o_reg_pair_count, + space=cutlass.AddressSpace.rmem, + ) + for pair_idx in cutlass.range_constexpr(o_reg_pair_count): + scale_base = ((pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2) * 2 + reg_base = pair_idx * 2 + final_pair = ffma2( + ( + final_scale0[scale_base], + final_scale0[scale_base + 1], + ), + ( + o0_loaded[reg_base], + o0_loaded[reg_base + 1], + ), + fmul2( + ( + final_scale1[scale_base], + final_scale1[scale_base + 1], + ), + ( + o1_loaded[reg_base], + o1_loaded[reg_base + 1], + ), + ), + ) + regs_o[pair_idx] = pack_float2_to_bf16(final_pair[0], final_pair[1]) + + for store_row_block_idx in cutlass.range_constexpr( + num_o_stsm_row_blocks(cfg) + ): + smem_offset_bytes, _, _, _ = o_stage_stsm_and_copy_offsets( + cfg, + warp_grp_thread_idx, + warp_idx, + lane_idx, + chunk_idx, + store_row_block_idx, + ) + smem_dst = self._smem_o_i32.data_ptr( + smem_offset_bytes >> SMEM_WORD_BYTE_SHIFT + ) + reg_store_offset = store_row_block_idx * TCGEN05_16X256B_REGS_PER_LOAD + store_regs = (regs_o.data_ptr() + reg_store_offset).load( + count=TCGEN05_16X256B_REGS_PER_LOAD, + alignment=SMEM_WORD_BYTES, + ) + prims.stmatrix( + smem_dst, + store_regs, + prims.MMALayout.COL, + shape=prims.StoreShape.M8N8, + ) + + @cute.jit + def _store_fp8_tail_vals_to_smem(self, task_cache, final_vals): + """Pack normalized FP8 O values and stage them into store SMEM.""" + cfg = self.cfg + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + regs_fp8 = cutlass.Array( + Int32, + num_fp8_output_regs(cfg), + space=cutlass.AddressSpace.rmem, + ) + for packed_idx in cutlass.range_constexpr(num_fp8_output_regs(cfg)): + val_base = packed_idx * TCGEN05_16X256B_REGS_PER_LOAD + regs_fp8[packed_idx] = pack_float4_to_fp8_e4m3( + final_vals[val_base], + final_vals[val_base + 1], + final_vals[val_base + 2], + final_vals[val_base + 3], + ) + if cutlass.const_expr(num_fp8_output_regs(cfg) == 2): + store_transposed_smem8b_x2( + self._smem_o_i32, + regs_fp8[0], + regs_fp8[1], + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.head_dim_per_stage_v, + ) + else: + for stsm_chunk_idx in cutlass.range_constexpr( + num_fp8_output_regs(cfg) // 4 + ): + reg_base = stsm_chunk_idx * TCGEN05_16X256B_REGS_PER_LOAD + store_transposed_smem8b_x4( + self._smem_o_i32, + regs_fp8[reg_base], + regs_fp8[reg_base + 1], + regs_fp8[reg_base + 2], + regs_fp8[reg_base + 3], + warp_grp_thread_idx, + cfg.tile_size_q, + cfg.head_dim_per_stage_v, + stsm_idx=stsm_chunk_idx, + ) + + @cute.jit + def _store_bf16_tail_regs_to_smem(self, task_cache, regs_o): + """Stage normalized BF16/partial O registers into store SMEM.""" + cfg = self.cfg + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + for store_chunk_idx in cutlass.range_constexpr(cfg.o_copy_segments_per_stage): + stsm_group_idx = store_chunk_idx // num_o_stsm_row_blocks(cfg) + copy_segment_idx = store_chunk_idx % num_o_stsm_row_blocks(cfg) + if cutlass.const_expr(cfg.tile_size_q == 8): + stsm_group_idx = 0 + copy_segment_idx = store_chunk_idx + smem_offset_bytes, _, _, _ = o_stage_stsm_and_copy_offsets( + cfg, + warp_grp_thread_idx, + warp_idx, + lane_idx, + stsm_group_idx, + copy_segment_idx, + ) + smem_dst = self._smem_o_i32.data_ptr( + smem_offset_bytes >> SMEM_WORD_BYTE_SHIFT + ) + reg_store_offset = store_chunk_idx * TCGEN05_16X256B_REGS_PER_LOAD + store_regs = (regs_o.data_ptr() + reg_store_offset).load( + count=TCGEN05_16X256B_REGS_PER_LOAD, + alignment=SMEM_WORD_BYTES, + ) + prims.stmatrix( + smem_dst, + store_regs, + prims.MMALayout.COL, + shape=prims.StoreShape.M8N8, + ) + + @cute.jit + def _stage_generic_tail_o_to_smem( + self, + task_cache, + *, + base_addr0, + base_addr1, + final_scale0, + final_scale1, + ): + """Stage the generic swaps-MMA-AB O tail path through SMEM.""" + cfg = self.cfg + shape = TCGEN05_16X256B_SHAPE + q_repeats = num_o_repeats(cfg) + o_reg_pair_count = num_o_reg_pairs(cfg) + num_o_tmem_loads = num_o_tmem_loads_per_stage(cfg) + + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + final_vals = cutlass.Array( + Float32, + o_reg_pair_count * num_o_tmem_loads * 2, + space=cutlass.AddressSpace.rmem, + ) + else: + regs_o = cutlass.Array( + Int32, + o_reg_pair_count * num_o_tmem_loads, + space=cutlass.AddressSpace.rmem, + ) + for chunk_idx in cutlass.range_constexpr(num_o_tmem_loads): + if cutlass.const_expr(cfg.tile_size_q == 8): + o0_loaded_lo = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr0, chunk_idx), Float32 + ), + num=1, + ) + o1_loaded_lo = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr1, chunk_idx), Float32 + ), + num=1, + ) + o0_loaded_hi = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_second_panel_addr( + tcgen05_panel_addr(base_addr0, chunk_idx) + ), + Float32, + ), + num=1, + ) + o1_loaded_hi = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_second_panel_addr( + tcgen05_panel_addr(base_addr1, chunk_idx) + ), + Float32, + ), + num=1, + ) + else: + o0_loaded = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr0, chunk_idx), Float32 + ), + num=q_repeats, + ) + o1_loaded = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr1, chunk_idx), Float32 + ), + num=q_repeats, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + for pair_idx in cutlass.range_constexpr(o_reg_pair_count): + scale_base = ((pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2) * 2 + if cutlass.const_expr(cfg.tile_size_q == 8): + reg_base = (pair_idx % 2) * 2 + o0_pair = ( + o0_loaded_lo[reg_base], + o0_loaded_lo[reg_base + 1], + ) + o1_pair = ( + o1_loaded_lo[reg_base], + o1_loaded_lo[reg_base + 1], + ) + if cutlass.const_expr(pair_idx >= 2): + o0_pair = ( + o0_loaded_hi[reg_base], + o0_loaded_hi[reg_base + 1], + ) + o1_pair = ( + o1_loaded_hi[reg_base], + o1_loaded_hi[reg_base + 1], + ) + else: + reg_base = pair_idx * 2 + o0_pair = ( + o0_loaded[reg_base], + o0_loaded[reg_base + 1], + ) + o1_pair = ( + o1_loaded[reg_base], + o1_loaded[reg_base + 1], + ) + final_pair = ffma2( + ( + final_scale0[scale_base], + final_scale0[scale_base + 1], + ), + o0_pair, + fmul2( + ( + final_scale1[scale_base], + final_scale1[scale_base + 1], + ), + o1_pair, + ), + ) + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + final_base = chunk_idx * o_reg_pair_count * 2 + pair_idx * 2 + final_vals[final_base] = final_pair[0] + final_vals[final_base + 1] = final_pair[1] + else: + regs_o[chunk_idx * o_reg_pair_count + pair_idx] = ( + pack_float2_to_bf16(final_pair[0], final_pair[1]) + ) + + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + self._store_fp8_tail_vals_to_smem(task_cache, final_vals) + else: + self._store_bf16_tail_regs_to_smem(task_cache, regs_o) + + @cute.jit + def _stage_swaps_tail_o_to_smem( + self, + task_cache, + *, + base_addr0, + base_addr1, + final_scale0, + final_scale1, + ): + """Select the swaps-MMA-AB O staging path for one V stage.""" + if cutlass.const_expr( + self.cfg.tile_size_q == 32 + and ( + self.cfg.use_fp8_output != 1 + or self.acc_o_tensor is not None + or self.cfg.use_cluster_reduction == 1 + ) + ): + self._stage_tile32_tail_o_to_smem( + task_cache, + base_addr0=base_addr0, + base_addr1=base_addr1, + final_scale0=final_scale0, + final_scale1=final_scale1, + ) + else: + self._stage_generic_tail_o_to_smem( + task_cache, + base_addr0=base_addr0, + base_addr1=base_addr1, + final_scale0=final_scale0, + final_scale1=final_scale1, + ) + + @cute.jit + def _normalize_and_store_tail_o( + self, + stage_info: StageInfo, + *, + new_max_arr, + sum_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_new_max_arr, + inst1_sum_arr, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ): + """Normalize final O and store O/LSE or split-KV partials.""" + if cutlass.const_expr(self.inst_id == 1): + cfg = self.cfg + task_cache = decode_gen_task_cache(stage_info) + tmem_row_base = task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + cta_idx_head_dim_v = cta_idx_head_dim_v_for_stage( + self.cta_idx_head_dim_v, stage_info + ) + head_dim_offset = head_dim_cta_offset_v(cfg, cta_idx_head_dim_v) + final_max, reduced_sum, final_scale0, final_scale1 = ( + self._compute_tail_softmax_scales( + task_cache, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + ) + self._store_lse_to_gmem(stage_info, task_cache, final_max, reduced_sum) + + for v_stage_idx in cutlass.range_constexpr(cfg.v_head_dim_stages): + base_addr0 = ( + tmem_row_base + + self._o_base_col() + + o_stage_tmem_col_offset(cfg, tail_o_stage_idx_0, v_stage_idx) + ) + base_addr1 = ( + tmem_row_base + + self._o_base_col() + + o_stage_tmem_col_offset(cfg, tail_o_stage_idx_1, v_stage_idx) + ) + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + self._store_keeps_tail_o_stage( + stage_info, + task_cache, + base_addr0=base_addr0, + final_scale0=final_scale0, + v_stage_idx=v_stage_idx, + head_dim_offset=head_dim_offset, + ) + else: + self._stage_swaps_tail_o_to_smem( + task_cache, + base_addr0=base_addr0, + base_addr1=base_addr1, + final_scale0=final_scale0, + final_scale1=final_scale1, + ) + cute.arch.fence_view_async_shared() + prims.barrier_cta_sync( + barrier_id=self.store_barrier_id, + thread_count=WARPGROUP_THREADS, + ) + self._store_o_slice_to_gmem( + stage_info, task_cache, v_stage_idx, head_dim_offset + ) + prims.barrier_cta_sync( + barrier_id=self.store_barrier_id, + thread_count=WARPGROUP_THREADS, + ) + self._reduce_cluster_partials_and_store( + stage_info, task_cache, head_dim_offset + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py new file mode 100644 index 000000000000..eebeb725d40f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py @@ -0,0 +1,441 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TMEM output-accumulator resource for PV MMA.""" + +from dataclasses import dataclass +from typing import ClassVar + +from cutlass.experimental import primitives as prims + +import cutlass +import cutlass.cute as cute +from cutlass import Boolean, Float32, Int32 +from cutlass.experimental import primitives as cprims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import TmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...helpers.constants import ( + WARP_LANES, +) + +from ...helpers.layout import ( + _TASK_CACHE_TMEM_BASE_OFFSET, + SMEM_ROW_BYTES, + decode_gen_task_cache, + o_stage_tmem_col_offset, + q_p_desc_k_block_wrap_bytes, +) +from ...helpers.math import ( + mma_k_step_for_qkv, + mma_kind_for_qkv, + qkv_dtype, + qkv_smem_swizzle, +) +from ...helpers.ops import ( + freeze_smem_descriptor, + tcgen05_panel_addr, +) +from ...helpers.tile import ( + batch_idx_for_stage_cfg, + cta_idx_q_for_stage, + runtime_local_kv_tiles, + runtime_seq_len_kv_from_task_cache, +) + +from .common import ( + TCGEN05_BF16_SWIZZLE_STRIDE_BYTES, + MlaResource, +) + +# ===================================================================== +# TmemOResource — O accumulator in TMEM, UmmaProducerAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemOResource(MlaResource): + """TMEM O accumulator resource that corrects and stores final outputs.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("o_stage_idx", Int32, Int32(0), "Current O TMEM stage."), + ("tail_o_stage_idx_0", Int32, Int32(0), "Final O stage for instance 0."), + ("tail_o_stage_idx_1", Int32, Int32(1), "Final O stage for instance 1."), + ) + p0_ref: object = None + p1_ref: object = None + p_tmem_ref: object = None + cache_seqs: object = None + batch_idx: object = None + cta_idx_q: object = None + o_stage_idx: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + tail_o_stage_idx_0: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + tail_o_stage_idx_1: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def get_tmem_requirements(self): + if self._tmem_alloc is None: + o_buffer_cols = self.cfg.tmem_o_buffer_cols * self.cfg.v_head_dim_stages + if ( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.head_dim_per_cta_v > 256 + ): + o_buffer_cols = 2 * self.cfg.tmem_o_buffer_cols + self._tmem_alloc = TmemAllocation( + name=f"{self.name}", + num_columns=self.cfg.o_stages * o_buffer_cols, + ) + return [self._tmem_alloc] + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_state(self, stage_info: StageInfo) -> None: + """Initialize O TMEM state needed by the MMA producer path.""" + # Producer aux work initializes the O accumulator TMEM view before PV + # MMA starts writing into it. + self._init_tmem_state(stage_info) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1), + ) + @cute.jit + def init_stage_state(self, stage_info: StageInfo): + """Initialize O-stage index variables for the first work tile.""" + # Consumer aux work creates the stage-tracking variables used by the + # correction task after it waits for O. + self._init_tmem_state(stage_info) + return Int32(0), Int32(0), Int32(1) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1), + ) + @cute.jit + def init_stage_work_tile_state(self, stage_info: StageInfo): + """Initialize O-stage index variables for a persistent work tile.""" + # Reset O stage bookkeeping for each persistent work tile. + del stage_info + return Int32(0), Int32(0), Int32(1) + + @cute.jit + def _tail_has_prior_o(self, stage_info: StageInfo): + """Return whether tail PV should accumulate into existing O state.""" + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, self.cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + self.cfg, + decode_gen_task_cache(stage_info), + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + return runtime_local_kv_tiles(self.cfg, seq_len_kv) > Int32( + self.cfg.num_insts_kv + ) + + @cute.jit + def _p_desc_for_inst(self, p_inst: int): + """Build the P SMEM descriptor for one softmax/PV pipe instance.""" + if cutlass.const_expr(p_inst == 0): + return cprims.Tcgen05SmemDesc.build( + self.p0_ref._smem_p, + leading_byte_offset=Int32(self.cfg.tile_size_q * SMEM_ROW_BYTES), + stride_byte_offset=TCGEN05_BF16_SWIZZLE_STRIDE_BYTES, + layout=qkv_smem_swizzle(self.cfg), + ) + return cprims.Tcgen05SmemDesc.build( + self.p1_ref._smem_p, + leading_byte_offset=Int32(self.cfg.tile_size_q * SMEM_ROW_BYTES), + stride_byte_offset=TCGEN05_BF16_SWIZZLE_STRIDE_BYTES, + layout=qkv_smem_swizzle(self.cfg), + ) + + @cute.jit + def _pv_o_base_addr( + self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] + ): + """Return the TMEM O base address for this PV stage.""" + cfg = self.cfg + v_stage_idx = v_subtile_idx % cfg.v_head_dim_stages + task_cache = decode_gen_task_cache(stage_info) + return ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._tmem_alloc.offset) + + o_stage_tmem_col_offset(cfg, stage_info.stage_idx, v_stage_idx) + ) + + @cute.jit + def _issue_smem_p_pv_mma( + self, + stage_info: StageInfo, + p_desc, + v_desc, + scale_d, + *, + v_subtile_idx: cutlass.Constexpr[int], + ): + """Issue PV MMA with P sourced from SMEM and O accumulated in TMEM.""" + cfg = self.cfg + idesc = cprims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + a_major=1, + b_major=0, + n_dim=cfg.tile_size_q, + m_dim=cfg.head_dim_per_stage_v, + ) + if prims.elect_sync(): + tmem_col = prims.make_tmem_ptr( + self._pv_o_base_addr(stage_info, v_subtile_idx=v_subtile_idx), + Float32, + ) + mma_k_step = mma_k_step_for_qkv(cfg) + k_block_count = cfg.tile_size_kv // mma_k_step + k_blocks_per_smem_row = 128 // (mma_k_step * cfg.qkv_dtype_bytes) + for k_block in cutlass.range_constexpr(k_block_count): + prims.tcgen05_mma( + mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + v_desc, + p_desc, + idesc, + scale_d, + ) + scale_d = Boolean(True) + if cutlass.const_expr(k_block + 1 < k_block_count): + v_desc = v_desc + Int32(256 if cfg.is_fp8_qkv() else 128) + if cutlass.const_expr((k_block + 1) % k_blocks_per_smem_row == 0): + p_desc = p_desc.advance_start_address( + Int32(q_p_desc_k_block_wrap_bytes(cfg)) + ) + else: + p_desc = p_desc.advance_start_address(Int32(16 * 2)) + + @producer_work + @cute.jit + def pv_mma_loop_0( + self, + stage_info: StageInfo, + *, + v_desc_0, + v_subtile_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ): + """Issue loop PV MMA instance 0 from P0/V0.""" + del is_tail + self._issue_smem_p_pv_mma( + stage_info, + self._p_desc_for_inst(0), + freeze_smem_descriptor(v_desc_0), + stage_info.loop_offset != Int32(0), + v_subtile_idx=v_subtile_idx, + ) + + @producer_work + @cute.jit + def pv_mma_tail_0( + self, + stage_info: StageInfo, + *, + v_desc_0, + v_subtile_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = True, + ): + """Issue tail PV MMA instance 0 from P0/V0.""" + del is_tail + self._issue_smem_p_pv_mma( + stage_info, + self._p_desc_for_inst(0), + freeze_smem_descriptor(v_desc_0), + self._tail_has_prior_o(stage_info), + v_subtile_idx=v_subtile_idx, + ) + + @producer_work + @cute.jit + def pv_mma_loop_1( + self, + stage_info: StageInfo, + *, + v_desc_1, + v_subtile_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ): + """Issue loop PV MMA instance 1 from P1/V1.""" + del is_tail + self._issue_smem_p_pv_mma( + stage_info, + self._p_desc_for_inst(1), + freeze_smem_descriptor(v_desc_1), + stage_info.loop_offset != Int32(0), + v_subtile_idx=v_subtile_idx, + ) + + @producer_work + @cute.jit + def pv_mma_tail_1( + self, + stage_info: StageInfo, + *, + v_desc_1, + v_subtile_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = True, + ): + """Issue tail PV MMA instance 1 from P1/V1.""" + del is_tail + self._issue_smem_p_pv_mma( + stage_info, + self._p_desc_for_inst(1), + freeze_smem_descriptor(v_desc_1), + self._tail_has_prior_o(stage_info), + v_subtile_idx=v_subtile_idx, + ) + + @producer_work + @cute.jit + def pv_mma_loop_tmem_p( + self, + stage_info: StageInfo, + *, + p_stage_idx, + v_desc_0, + v_subtile_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ): + """Issue loop PV MMA with P read from TMEM instead of SMEM.""" + del is_tail + self._issue_tmem_p_pv_mma( + stage_info, + p_stage_idx=p_stage_idx, + v_desc=freeze_smem_descriptor(v_desc_0), + scale_d=stage_info.loop_offset != Int32(0), + v_subtile_idx=v_subtile_idx, + ) + + @producer_work + @cute.jit + def pv_mma_tail_tmem_p( + self, + stage_info: StageInfo, + *, + p_stage_idx, + v_desc_0, + v_subtile_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = True, + ): + """Issue tail PV MMA with P read from TMEM instead of SMEM.""" + del is_tail + self._issue_tmem_p_pv_mma( + stage_info, + p_stage_idx=p_stage_idx, + v_desc=freeze_smem_descriptor(v_desc_0), + scale_d=self._tail_has_prior_o(stage_info), + v_subtile_idx=v_subtile_idx, + ) + + @cute.jit + def _issue_tmem_p_pv_mma( + self, + stage_info: StageInfo, + *, + p_stage_idx, + v_desc, + scale_d, + v_subtile_idx: cutlass.Constexpr[int], + ): + """Issue PV MMA for keeps-MMA-AB with P sourced from TMEM.""" + cfg = self.cfg + v_stage_idx = v_subtile_idx % cfg.v_head_dim_stages + task_cache = decode_gen_task_cache(stage_info) + p_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self.p_tmem_ref.tmem_alias_ref._tmem_alloc.offset) + + Int32(WARP_LANES) + + Int32(p_stage_idx) * Int32(cfg.tmem_s_cols) + ) + if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): + p_addr = tcgen05_panel_addr(p_addr, v_stage_idx // 2) + idesc = cprims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + a_major=0, + b_major=1, + n_dim=cfg.head_dim_per_stage_v, + m_dim=cfg.tile_size_q, + ) + if prims.elect_sync(): + base_addr = self._pv_o_base_addr(stage_info, v_subtile_idx=v_subtile_idx) + for k_block in cutlass.range_constexpr( + cfg.tile_size_kv // mma_k_step_for_qkv(cfg) + ): + prims.tcgen05_mma( + mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + prims.make_tmem_ptr(base_addr, Float32), + prims.make_tmem_ptr(p_addr, qkv_dtype(cfg)), + v_desc, + idesc, + scale_d, + ) + scale_d = Boolean(True) + if cutlass.const_expr( + k_block + 1 < cfg.tile_size_kv // mma_k_step_for_qkv(cfg) + ): + p_addr = p_addr + Int32(8) + v_desc = v_desc + Int32(256 if cfg.is_fp8_qkv() else 128) + + @consumer_work( + returns=(o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1), + ) + @cute.jit + def o_stage( + self, + stage_info: StageInfo, + *, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + inst_idx: cutlass.Constexpr[int], + is_tail: cutlass.Constexpr[bool] = False, + ): + """Publish the TMEM O stage index for correction/output.""" + # Consumer work for O: correction has waited for PV MMA to finish. It + # receives the current O stage, plus tail-stage bookkeeping so it can + # combine the two interleaved softmax/MMA instances. + o_stage_idx = stage_info.stage_idx + # The tail stage is the only point where the correction task needs to + # remember both final O stages for cross-instance normalization. + if cutlass.const_expr(is_tail): + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + tail_o_stage_idx_0 = stage_info.stage_idx + tail_o_stage_idx_1 = stage_info.stage_idx + elif cutlass.const_expr(inst_idx == 0): + tail_o_stage_idx_0 = stage_info.stage_idx + else: + tail_o_stage_idx_1 = stage_info.stage_idx + return o_stage_idx, tail_o_stage_idx_0, tail_o_stage_idx_1 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py new file mode 100644 index 000000000000..91dd9802cbff --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py @@ -0,0 +1,316 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TMEM probability resource for the keeps-MMA-AB PV operand.""" + +from dataclasses import dataclass +from typing import ClassVar, Optional + +from cutlass.experimental import primitives as prims + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...helpers.constants import ( + WARP_LANES, +) + + +from ...helpers.layout import ( + _TASK_CACHE_TMEM_BASE_OFFSET, + decode_gen_task_cache, + num_packed_p_regs, + num_softmax_scale_groups, +) +from ...helpers.math import ( + fadd2, + fmul2, + neg_max_f32, + pack_float2_to_bf16, +) +from ...helpers.ops import ( + float_to_u32_bits, + fp8_log2_quant_scale, + pack_float4_to_fp8_e4m3, + softmax_sum_state_ptr, + tcgen05_second_panel_addr, + tcgen05_store_p_16x32bx2_x16, + tcgen05_store_p_fp8_16x32bx2_x16, + u32_bits_to_float, +) + +from .common import ( + MlaResource, +) + +# ===================================================================== +# TmemPResource — P in TMEM for keeps-MMA-AB PV MMA +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemPResource(MlaResource): + """TMEM probability tile exchanged from softmax to PV MMA. + + The keeps-MMA-AB schedule uses one P pipe: softmax writes BF16 P directly + into the score TMEM stage, and MmaTask consumes that TMEM stage with + ``tcgen05.mma`` A-from-TMEM. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("p_stage_idx", Int32, Int32(0), "Current TMEM P pipeline stage."), + ) + inst_id: cutlass.Constexpr[int] = 0 + scale_softmax_log2: Float32 = None + tmem_alias_ref: Optional[MlaResource] = None + p_stage_idx: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_stage_idx) + @cute.jit + def init_stage_state(self, stage_info: StageInfo): + """Initialize the TMEM P stage index for the first work tile.""" + # Keeps-MMA-AB uses TMEM P instead of SMEM P. Consumer aux state tracks + # which score/TMEM stage the PV MMA consumer should read. + self._init_tmem_state(stage_info) + return Int32(0) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_stage_idx) + @cute.jit + def init_stage_work_tile_state(self, stage_info: StageInfo): + """Initialize the TMEM P stage index for a persistent work tile.""" + # Reset the TMEM P stage index when a persistent CTA advances to a new + # work tile. + del stage_info + return Int32(0) + + @cute.jit + def _p_base_addr(self, stage_info: StageInfo, task_cache): + p_base = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self.tmem_alias_ref._tmem_alloc.offset) + + stage_info.stage_idx * Int32(self.cfg.tmem_s_cols) + ) + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + p_base = p_base + Int32(WARP_LANES) + return p_base + + @producer_work + @cute.jit + def materialize_p( + self, + stage_info: StageInfo, + *, + new_max_arr, + s_arr, + local_sum_arr, + ): + """Materialize grouped-head softmax probabilities into TMEM P.""" + # Producer work for TmemPResource: softmax converts S to P and stores it + # directly into the TMEM score/P stage. PV MMA consumes the returned + # p_stage_idx through the loop/tail TMEM-P producer labels. + cfg = self.cfg + task_cache = decode_gen_task_cache(stage_info) + + num_scale_groups = num_softmax_scale_groups(cfg) + neg_scaled_max = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + local_sums = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + for idx in cutlass.range_constexpr(num_scale_groups): + new_max = new_max_arr[idx] + safe_new_max = new_max + if safe_new_max == neg_max_f32(): + safe_new_max = Float32(0.0) + neg_scaled_max[idx] = -self.scale_softmax_log2 * safe_new_max + local_sums[idx] = Float32(0.0) + + packed_p_reg_count = num_packed_p_regs(cfg) + regs_p = cutlass.Array( + Int32, packed_p_reg_count, space=cutlass.AddressSpace.rmem + ) + if cutlass.const_expr( + cfg.is_fp8_qkv() and cfg.kernel_variant == "keeps_mma_ab" + ): + log2_scale_pair = (self.scale_softmax_log2, self.scale_softmax_log2) + neg_scaled_pair = ( + neg_scaled_max[0] + fp8_log2_quant_scale(), + neg_scaled_max[0] + fp8_log2_quant_scale(), + ) + has_finite_max = new_max_arr[0] != neg_max_f32() + for packed_idx in cutlass.range_constexpr(packed_p_reg_count): + s_base = packed_idx * 4 + p0 = Float32(0.0) + p1 = Float32(0.0) + p2 = Float32(0.0) + p3 = Float32(0.0) + if has_finite_max: + scaled01 = fadd2( + fmul2( + ( + s_arr[s_base + 0], + s_arr[s_base + 1], + ), + log2_scale_pair, + ), + neg_scaled_pair, + ) + scaled23 = fadd2( + fmul2( + ( + s_arr[s_base + 2], + s_arr[s_base + 3], + ), + log2_scale_pair, + ), + neg_scaled_pair, + ) + p0 = cute.math.exp2(scaled01[0], fastmath=True) + p1 = cute.math.exp2(scaled01[1], fastmath=True) + p2 = cute.math.exp2(scaled23[0], fastmath=True) + p3 = cute.math.exp2(scaled23[1], fastmath=True) + local_sums[0] += p0 + local_sums[0] += p1 + local_sums[0] += p2 + local_sums[0] += p3 + regs_p[packed_idx] = pack_float4_to_fp8_e4m3(p0, p1, p2, p3) + elif cutlass.const_expr(cfg.is_fp8_qkv()): + for packed_idx in cutlass.range_constexpr(packed_p_reg_count): + p_vals = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for elem_idx in cutlass.range_constexpr(4): + s_idx = packed_idx * 4 + elem_idx + pair_idx = s_idx // 2 + scale_base = ( + (pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2 + ) * 2 + scale_idx = scale_base + (s_idx % 2) + p_val = Float32(0.0) + if new_max_arr[scale_idx] != neg_max_f32(): + p_val = cute.math.exp2( + s_arr[s_idx] * self.scale_softmax_log2 + + neg_scaled_max[scale_idx] + + fp8_log2_quant_scale(), + fastmath=True, + ) + p_vals[elem_idx] = p_val + local_sums[scale_idx] += p_val + regs_p[packed_idx] = pack_float4_to_fp8_e4m3( + p_vals[0], p_vals[1], p_vals[2], p_vals[3] + ) + else: + for pair_idx in cutlass.range_constexpr(packed_p_reg_count): + s0 = pair_idx * 2 + s1 = s0 + 1 + p0 = Float32(0.0) + p1 = Float32(0.0) + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + new_max = new_max_arr[0] + if new_max != neg_max_f32(): + p0 = cute.math.exp2( + s_arr[s0] * self.scale_softmax_log2 + neg_scaled_max[0], + fastmath=True, + ) + p1 = cute.math.exp2( + s_arr[s1] * self.scale_softmax_log2 + neg_scaled_max[0], + fastmath=True, + ) + local_sums[0] += p0 + local_sums[0] += p1 + else: + scale0 = ((pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2) * 2 + scale1 = scale0 + 1 + new_max0 = new_max_arr[scale0] + new_max1 = new_max_arr[scale1] + if new_max0 != neg_max_f32(): + p0 = cute.math.exp2( + s_arr[s0] * self.scale_softmax_log2 + + neg_scaled_max[scale0], + fastmath=True, + ) + if new_max1 != neg_max_f32(): + p1 = cute.math.exp2( + s_arr[s1] * self.scale_softmax_log2 + + neg_scaled_max[scale1], + fastmath=True, + ) + local_sums[scale0] += p0 + local_sums[scale1] += p1 + regs_p[pair_idx] = pack_float2_to_bf16(p0, p1) + + for scale_idx in cutlass.range_constexpr(num_scale_groups): + self.tmem_alias_ref._p_local_sum_arr[scale_idx] = local_sums[scale_idx] + local_sum_arr[scale_idx] = local_sums[scale_idx] + + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + state_idx = cute.arch.thread_idx()[0] + state_ptr = self.tmem_alias_ref._softmax_scratch.data_ptr(state_idx) + old_max = u32_bits_to_float(state_ptr.load(is_volatile=True, alignment=4)) + old_sum = u32_bits_to_float( + softmax_sum_state_ptr(state_ptr).load(is_volatile=True, alignment=4) + ) + new_max = new_max_arr[0] + exp_scale = cute.math.exp2( + self.scale_softmax_log2 * (old_max - new_max), + fastmath=True, + ) + updated_sum = exp_scale * old_sum + local_sums[0] + state_ptr.store( + float_to_u32_bits(new_max), + is_volatile=True, + alignment=4, + ) + softmax_sum_state_ptr(state_ptr).store( + float_to_u32_bits(updated_sum), + is_volatile=True, + alignment=4, + ) + + p_base = self._p_base_addr(stage_info, task_cache) + if cutlass.const_expr(cfg.is_fp8_qkv()): + tcgen05_store_p_fp8_16x32bx2_x16(p_base, regs_p) + if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): + tcgen05_store_p_fp8_16x32bx2_x16( + tcgen05_second_panel_addr(p_base), regs_p + ) + else: + tcgen05_store_p_16x32bx2_x16(p_base, regs_p, 0) + tcgen05_store_p_16x32bx2_x16(p_base + Int32(16), regs_p, 16) + if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base), regs_p, 0 + ) + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base) + Int32(16), regs_p, 16 + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + + @consumer_work(returns=p_stage_idx) + @cute.jit + def p_stage(self, stage_info: StageInfo): + """Publish the TMEM P pipeline stage consumed by PV MMA.""" + # Consumer work returns only the live TMEM stage index; the payload is + # the P data already stored in the aliased TMEM allocation. + return stage_info.stage_idx diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py new file mode 100644 index 000000000000..88c7aa259d4f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py @@ -0,0 +1,817 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TMEM score resources for QK MMA and online softmax.""" + +from dataclasses import dataclass +from typing import ClassVar, Optional + +from cutlass.experimental import primitives as prims + +import cutlass +import cutlass.cute as cute +from cutlass import Boolean, Float32, Int32, Int64, Uint32 +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import SmemAllocation, TmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...helpers.constants import ( + OCTET_LANES, + QUAD_LANE_MASK, + QUAD_LANE_SHIFT, + SCORE_ROWS_PER_Q_PAIR, + SCORE_TOKENS_PER_QK_GROUP, + SMEM_WORD_BYTES, + TCGEN05_16X256B_SHAPE, + TCGEN05_DESC_NEXT_K_BLOCK_UNITS, + TCGEN05_DESC_WRAPPED_K_BLOCK_UNITS, + WARP_LANES, + WARPGROUP_THREADS, + WARPGROUP_WARPS, +) + + +from ...helpers.layout import ( + _TASK_CACHE_LANE_IDX, + _TASK_CACHE_TMEM_BASE_OFFSET, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _TASK_CACHE_WARP_IDX, + decode_gen_task_cache, + num_q_repeats, + num_s_regs_per_thread, + num_softmax_scale_groups, + q_p_desc_k_block_wrap_units, + q_stage_smem_element_offset, + smem_array, + softmax_scratch_words, +) +from ...helpers.math import ( + ffma2, + float_to_u32_for_atomic_max, + init_softmax_scratch_u32, + mma_k_step_for_qkv, + mma_kind_for_qkv, + neg_max_f32, + qkv_dtype, + smem_atomic_max_u32, + u32_to_float_for_atomic_max, +) +from ...helpers.mask import MaskType +from ...helpers.ops import ( + float_to_u32_bits, + freeze_smem_descriptor, + softmax_sum_state_ptr, + tcgen05_ld_16x32bx2_f32, + tcgen05_second_panel_addr, + u32_bits_to_float, +) +from ...helpers.stage import MlaStage +from ...helpers.tile import ( + batch_idx_for_stage_cfg, + cta_idx_kv_for_stage, + cta_idx_q_for_stage, + global_kv_tile_idx, + runtime_seq_len_kv_for_query_row, + runtime_seq_len_kv_from_task_cache, + softmax_kv_tile_idx, +) + +from .common import ( + MlaResource, +) + +# ===================================================================== +# TmemSResource — S scores in TMEM, UmmaProducerAsync pipeline +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemSResource(MlaResource): + """TMEM score resource plus task-local online-softmax state.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("old_max_arr", cutlass.Array, None, "Previous running softmax maxima."), + ("sum_arr", cutlass.Array, None, "Running softmax denominators."), + ("new_max_arr", cutlass.Array, None, "Current running softmax maxima."), + ("local_sum_arr", cutlass.Array, None, "Local denominator contributions."), + ("s_arr", cutlass.Array, None, "Loaded S scores for the current tile."), + ) + inst_id: cutlass.Constexpr[int] = 0 + scale_softmax_log2: Float32 = None + p_ref: Optional[MemoryResource] = None + global_ref: Optional[MemoryResource] = None + cache_seqs: object = None + head_idx: object = None + batch_idx: object = None + cta_idx_q: object = None + cta_idx_kv: object = None + sync_barrier_id: cutlass.Constexpr[int] = 0 + q_desc_current: object = None + q_desc_rope_current: object = None + new_max_state: object = None + sum_state: object = None + _scratch_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + _softmax_scratch: object = None + _p_local_sum_arr: object = None + old_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + sum_arr: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + new_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + local_sum_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + s_arr: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + @cute.jit + def _reset_softmax_state(self): + """Reset per-resource softmax max and sum state arrays.""" + num_scale_groups = num_softmax_scale_groups(self.cfg) + self.new_max_state = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + self.sum_state = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + for idx in cutlass.range_constexpr(num_scale_groups): + self.new_max_state[idx] = neg_max_f32() + self.sum_state[idx] = Float32(0.0) + + def get_smem_requirements(self): + """Return the softmax scratch SMEM allocation.""" + if self._scratch_alloc is None: + self._scratch_alloc = SmemAllocation( + name=f"{self.name}_softmaxScratch", + size_bytes=softmax_scratch_words(self.cfg) * SMEM_WORD_BYTES, + alignment=16, + ) + return [self._scratch_alloc] + + def get_tmem_requirements(self): + """Return the TMEM allocation for score tiles.""" + if self._tmem_alloc is None: + num_stages = ( + self.pipeline_config.num_stages + if self.cfg.kernel_variant == "keeps_mma_ab" + and self.pipeline_config is not None + else 1 + ) + self._tmem_alloc = TmemAllocation( + name=f"{self.name}", + num_columns=self.cfg.tmem_s_cols * num_stages, + ) + return [self._tmem_alloc] + + @cute.jit + def _make_initial_softmax_vars(self): + """Create fresh softmax state arrays for the current work tile.""" + self._reset_softmax_state() + num_scale_groups = num_softmax_scale_groups(self.cfg) + num_s_regs = num_s_regs_per_thread(self.cfg) + self._p_local_sum_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + old_max_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + sum_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + new_max_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + local_sum_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + s_arr = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) + for idx in cutlass.range_constexpr(num_scale_groups): + old_max_arr[idx] = neg_max_f32() + sum_arr[idx] = Float32(0.0) + new_max_arr[idx] = neg_max_f32() + local_sum_arr[idx] = Float32(0.0) + self._p_local_sum_arr[idx] = Float32(0.0) + for idx in cutlass.range_constexpr(num_s_regs): + s_arr[idx] = neg_max_f32() + return old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr), + ) + @cute.jit + def init_softmax_state(self, stage_info: StageInfo): + """Create softmax arrays and scratch state for the first work tile.""" + # Consumer aux work belongs to the softmax task. It creates the local + # arrays that will be threaded through update_softmax(), + # materialize_p(), and update_softmax_sum(). + context = stage_info.context + self._init_tmem_state(stage_info) + self._softmax_scratch = smem_array( + context, + self._scratch_alloc, + Uint32, + softmax_scratch_words(self.cfg), + ) + self.q_desc_current = Int64(0) + self.q_desc_rope_current = Int64(0) + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + thread_idx = cute.arch.thread_idx()[0] + state_ptr = self._softmax_scratch.data_ptr(thread_idx) + state_ptr.store( + float_to_u32_bits(neg_max_f32()), + is_volatile=True, + alignment=4, + ) + softmax_sum_state_ptr(state_ptr).store( + float_to_u32_bits(Float32(0.0)), + is_volatile=True, + alignment=4, + ) + return self._make_initial_softmax_vars() + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr), + ) + @cute.jit + def init_softmax_work_tile_state(self, stage_info: StageInfo): + """Create fresh softmax arrays for each persistent work tile.""" + # Persistent schedules reuse the same task graph for multiple work + # tiles, so the consumer-side softmax state must be reset per tile. + del stage_info + return self._make_initial_softmax_vars() + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def reset_softmax_work_tile_state(self, stage_info: StageInfo) -> None: + """Reset per-resource softmax state for a persistent work tile.""" + # Producer aux work mirrors the consumer reset for schedules where this + # resource later produces QK scores into TMEM. + del stage_info + self._make_initial_softmax_vars() + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def set_q_desc(self, stage_info: StageInfo, *, q_desc, q_desc_rope): + """Cache the Q descriptor once while the Q SMEM stage is live.""" + # q_desc() returns descriptors from the SmemQ consumer side. This + # producer aux method stores them in the score resource so each QK MMA + # call can reuse them without adding another schedule edge. + del stage_info + self.q_desc_current = q_desc + self.q_desc_rope_current = q_desc_rope + + @producer_work + @cute.jit + def qk_mma( + self, + stage_info: StageInfo, + *, + kv_desc, + k_subtile_idx: cutlass.Constexpr[int], + ): + """Issue the staged QK MMA for one 128-wide MLA head-dim slice.""" + # Producer work for TmemSResource: the MMA warp writes score fragments + # into the acquired TMEM score stage. The paired softmax task consumes + # this TMEM stage through update_softmax(). + cfg = self.cfg + qk_stage_idx = k_subtile_idx + q_stage_offset_bytes = Int32( + q_stage_smem_element_offset(cfg, qk_stage_idx) * cfg.qkv_dtype_bytes + ) + q_desc = freeze_smem_descriptor(self.q_desc_current) + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + if cutlass.const_expr( + qk_stage_idx == cfg.latent_dim // cfg.head_dim_per_stage_kv + ): + q_desc = freeze_smem_descriptor(self.q_desc_rope_current) + kv_desc = freeze_smem_descriptor(kv_desc) + if cutlass.const_expr( + not cfg.is_fp8_qkv() + or cfg.rope_dim != 64 + or qk_stage_idx != cfg.latent_dim // cfg.head_dim_per_stage_kv + ): + q_desc = q_desc + (q_stage_offset_bytes >> 4) + + task_cache = decode_gen_task_cache(stage_info) + stage_col_offset = Int32(0) + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + stage_col_offset = stage_info.stage_idx * Int32(cfg.tmem_s_cols) + tmem_ptr = prims.make_tmem_ptr( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._tmem_alloc.offset) + + stage_col_offset, + Float32, + ) + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + n_dim=cfg.tile_size_kv, + m_dim=cfg.tile_size_q, + ) + else: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + n_dim=cfg.tile_size_q, + m_dim=cfg.tile_size_kv, + ) + k_block_count = cfg.qk_head_stage_width(qk_stage_idx) // mma_k_step_for_qkv(cfg) + if prims.elect_sync(): + scale_d = Boolean(qk_stage_idx != 0) + for k_block in cutlass.range_constexpr(k_block_count): + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + prims.tcgen05_mma( + mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_ptr, + q_desc, + kv_desc, + idesc, + scale_d, + ) + else: + prims.tcgen05_mma( + mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_ptr, + kv_desc, + q_desc, + idesc, + scale_d, + ) + scale_d = Boolean(True) + if cutlass.const_expr(k_block + 1 < k_block_count): + if cutlass.const_expr(k_block == 3): + kv_desc = kv_desc + Int32(TCGEN05_DESC_WRAPPED_K_BLOCK_UNITS) + q_desc = q_desc + Int32(q_p_desc_k_block_wrap_units(cfg)) + else: + kv_desc = kv_desc + Int32(TCGEN05_DESC_NEXT_K_BLOCK_UNITS) + q_desc = q_desc + Int32(TCGEN05_DESC_NEXT_K_BLOCK_UNITS) + + @consumer_work(returns=(old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr)) + @cute.jit + def update_softmax( + self, + stage_info: StageInfo, + *, + old_max_arr, + sum_arr, + new_max_arr, + local_sum_arr, + s_arr, + section: cutlass.Constexpr[MlaStage], + ): + """Read S from TMEM and update the flat-query softmax state.""" + # Consumer work for TmemSResource: softmax waits for the QK MMA score + # stage, loads S from TMEM, applies masks, and returns updated local + # softmax state to the captured schedule. + cfg = self.cfg + task_cache = decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + warp_grp_thread_idx = cute.arch.thread_idx()[0] + q_repeats = num_q_repeats(cfg) + num_scale_groups = num_softmax_scale_groups(cfg) + old_max_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + sum_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + new_max_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + local_max_vals = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + s_vals = cutlass.Array( + Float32, num_s_regs_per_thread(cfg), space=cutlass.AddressSpace.rmem + ) + + for idx in cutlass.range_constexpr(num_scale_groups): + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + state_ptr = self._softmax_scratch.data_ptr(warp_grp_thread_idx) + old_max_vals[idx] = u32_bits_to_float( + state_ptr.load(is_volatile=True, alignment=4) + ) + sum_vals[idx] = u32_bits_to_float( + softmax_sum_state_ptr(state_ptr).load(is_volatile=True, alignment=4) + ) + else: + old_max_vals[idx] = new_max_arr[idx] + sum_vals[idx] = sum_arr[idx] + new_max_vals[idx] = old_max_vals[idx] + local_max_vals[idx] = neg_max_f32() + for idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + s_vals[idx] = neg_max_f32() + + # Normalize both score layouts into one local S register array before + # masking: keeps-MMA-AB reads a single TMEM panel, swaps-MMA-AB reads + # two 16x256b panels. + stage_col_offset = Int32(0) + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + stage_col_offset = stage_info.stage_idx * Int32(cfg.tmem_s_cols) + base_addr = ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._tmem_alloc.offset) + + stage_col_offset + ) + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr, Float32), + num=cfg.tile_size_kv // 2, + offset=Int32(cfg.tile_size_kv // 2), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + s_vals[reg_idx] = loaded[reg_idx] + else: + loaded0 = prims.tcgen05_ld( + TCGEN05_16X256B_SHAPE, + prims.make_tmem_ptr(base_addr, Float32), + num=q_repeats, + ) + loaded1 = prims.tcgen05_ld( + TCGEN05_16X256B_SHAPE, + prims.make_tmem_ptr(tcgen05_second_panel_addr(base_addr), Float32), + num=q_repeats, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for repeat_idx in cutlass.range_constexpr(q_repeats): + ld_base = repeat_idx * 4 + s_vals[ld_base + 0] = loaded0[ld_base + 0] + s_vals[ld_base + 1] = loaded0[ld_base + 1] + s_vals[ld_base + 2] = loaded0[ld_base + 2] + s_vals[ld_base + 3] = loaded0[ld_base + 3] + s_vals[q_repeats * 4 + ld_base + 0] = loaded1[ld_base + 0] + s_vals[q_repeats * 4 + ld_base + 1] = loaded1[ld_base + 1] + s_vals[q_repeats * 4 + ld_base + 2] = loaded1[ld_base + 2] + s_vals[q_repeats * 4 + ld_base + 3] = loaded1[ld_base + 3] + + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + cfg, + task_cache, + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + local_tile_idx = softmax_kv_tile_idx(cfg, stage_info, self.inst_id) + tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) + tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) + next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) + should_apply_dense_mask = ( + (seq_len_kv % Int32(cfg.tile_size_kv)) != Int32(0) + ) or (next_tile_offset_k > seq_len_kv) + needs_row_causal_mask = cutlass.const_expr( + cfg.mask_type == MaskType.CAUSAL.value and cfg.logical_seq_len_q > 1 + ) + if cutlass.const_expr(needs_row_causal_mask): + min_seq_len_kv = runtime_seq_len_kv_for_query_row( + cfg, + self.cache_seqs, + batch_idx, + cta_idx_q, + Int32(0), + self.cu_seqlens_q, + ) + should_apply_dense_mask = should_apply_dense_mask or ( + next_tile_offset_k > min_seq_len_kv + ) + if should_apply_dense_mask: + # The CTA domain follows the latest logical query row in the flat + # tile. Earlier rows can have a narrower bottom-right causal limit. + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + local_col_base = (lane_idx >> Int32(4)) * Int32(cfg.tile_size_kv // 2) + local_row_idx = warp_idx * Int32(16) + (lane_idx & Int32(0xF)) + row_seq_len_kv = seq_len_kv + if cutlass.const_expr(needs_row_causal_mask): + row_seq_len_kv = runtime_seq_len_kv_for_query_row( + cfg, + self.cache_seqs, + batch_idx, + cta_idx_q, + local_row_idx, + self.cu_seqlens_q, + ) + for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + token_idx = tile_offset_k + local_col_base + Int32(reg_idx) + if token_idx >= row_seq_len_kv: + s_vals[reg_idx] = neg_max_f32() + else: + local_idx_k0 = warp_idx * Int32(WARP_LANES) + ( + lane_idx >> Int32(QUAD_LANE_SHIFT) + ) + if cutlass.const_expr(cfg.kernel_variant != "keeps_mma_ab"): + # Score TMEM is read as two panels: each repeat contributes + # the first two 8-token groups, then the second panel carries + # the next two groups. + local_q_pair = (lane_idx & Int32(QUAD_LANE_MASK)) * Int32( + SCORE_ROWS_PER_Q_PAIR + ) + for repeat_idx in cutlass.range_constexpr(q_repeats): + row_in_tile_0 = ( + Int32(repeat_idx * SCORE_TOKENS_PER_QK_GROUP) + local_q_pair + ) + row_in_tile_1 = row_in_tile_0 + Int32(1) + seq_len_kv_0 = seq_len_kv + seq_len_kv_1 = seq_len_kv + if cutlass.const_expr(needs_row_causal_mask): + seq_len_kv_0 = runtime_seq_len_kv_for_query_row( + cfg, + self.cache_seqs, + batch_idx, + cta_idx_q, + row_in_tile_0, + self.cu_seqlens_q, + ) + seq_len_kv_1 = runtime_seq_len_kv_for_query_row( + cfg, + self.cache_seqs, + batch_idx, + cta_idx_q, + row_in_tile_1, + self.cu_seqlens_q, + ) + s_base = repeat_idx * 4 + s_second_panel_base = q_repeats * 4 + s_base + token_idx = tile_offset_k + local_idx_k0 + if token_idx >= seq_len_kv_0: + s_vals[s_base + 0] = neg_max_f32() + if token_idx >= seq_len_kv_1: + s_vals[s_base + 1] = neg_max_f32() + token_idx = ( + tile_offset_k + local_idx_k0 + Int32(SCORE_TOKENS_PER_QK_GROUP) + ) + if token_idx >= seq_len_kv_0: + s_vals[s_base + 2] = neg_max_f32() + if token_idx >= seq_len_kv_1: + s_vals[s_base + 3] = neg_max_f32() + token_idx = ( + tile_offset_k + + local_idx_k0 + + Int32(2 * SCORE_TOKENS_PER_QK_GROUP) + ) + if token_idx >= seq_len_kv_0: + s_vals[s_second_panel_base + 0] = neg_max_f32() + if token_idx >= seq_len_kv_1: + s_vals[s_second_panel_base + 1] = neg_max_f32() + token_idx = ( + tile_offset_k + + local_idx_k0 + + Int32(3 * SCORE_TOKENS_PER_QK_GROUP) + ) + if token_idx >= seq_len_kv_0: + s_vals[s_second_panel_base + 2] = neg_max_f32() + if token_idx >= seq_len_kv_1: + s_vals[s_second_panel_base + 3] = neg_max_f32() + + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + # Keeps-MMA-AB keeps the softmax state in scratch words indexed by + # CTA thread, so only a warp-level max is needed here. + local_max = old_max_vals[0] + for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + local_max = cute.math.max(local_max, s_vals[reg_idx], ftz=True) + local_max = cute.math.max( + local_max, + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=local_max, + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + new_max_vals[0] = local_max + self.new_max_state[0] = local_max + else: + # Swaps-MMA-AB reduces per-scale maxima first within each warp and + # then across the four softmax columns through SMEM atomics. + for scale_idx in cutlass.range_constexpr(num_scale_groups): + s_base = Int32((scale_idx // 2) * 4 + (scale_idx & 1)) + s_stride = Int32(q_repeats * 4) + local_max = cute.math.max( + cute.math.max(s_vals[s_base + 0], s_vals[s_base + 2], ftz=True), + cute.math.max( + s_vals[s_base + s_stride], + s_vals[s_base + s_stride + Int32(2)], + ftz=True, + ), + ftz=True, + ) + local_max = cute.math.max(local_max, old_max_vals[scale_idx], ftz=True) + local_max = cute.math.max( + local_max, + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=local_max, + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + local_max = cute.math.max( + local_max, + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=local_max, + offset=8, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + local_max_vals[scale_idx] = local_max + + # This is state initialization, not schedule dispatch: the first + # softmax tile in a work tile must clear the shared atomic scratch. + # The sync covers one four-warp softmax group. + if cutlass.const_expr(section == MlaStage.Head): + init_softmax_scratch_u32( + self._softmax_scratch, + warp_grp_thread_idx, + WARPGROUP_WARPS * num_softmax_scale_groups(cfg), + ) + prims.barrier_cta_sync( + barrier_id=self.sync_barrier_id, thread_count=WARPGROUP_THREADS + ) + elif cutlass.const_expr(section == MlaStage.Loop): + if stage_info.loop_offset == stage_info.loop_start: + init_softmax_scratch_u32( + self._softmax_scratch, + warp_grp_thread_idx, + WARPGROUP_WARPS * num_softmax_scale_groups(cfg), + ) + prims.barrier_cta_sync( + barrier_id=self.sync_barrier_id, + thread_count=WARPGROUP_THREADS, + ) + + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + col_group_idx = lane_idx & Int32(QUAD_LANE_MASK) + if lane_idx < Int32(OCTET_LANES): + atomic_reduce_base = col_group_idx * Int32(num_scale_groups) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + smem_atomic_max_u32( + self._softmax_scratch.data_ptr( + atomic_reduce_base + Int32(scale_idx) + ), + float_to_u32_for_atomic_max(local_max_vals[scale_idx]), + ) + prims.barrier_cta_sync( + barrier_id=self.sync_barrier_id, thread_count=WARPGROUP_THREADS + ) + + reduced_max_ptr = self._softmax_scratch.data_ptr( + col_group_idx * Int32(num_scale_groups) + ) + reduced_max = reduced_max_ptr.load( + count=num_scale_groups, + alignment=16 if num_scale_groups == 4 else 8, + ) + for scale_idx in cutlass.range_constexpr(num_scale_groups): + new_max_vals[scale_idx] = u32_to_float_for_atomic_max( + reduced_max[scale_idx] + ) + self.new_max_state[scale_idx] = new_max_vals[scale_idx] + for scale_idx in cutlass.range_constexpr(num_scale_groups): + old_max_arr[scale_idx] = old_max_vals[scale_idx] + sum_arr[scale_idx] = sum_vals[scale_idx] + new_max_arr[scale_idx] = new_max_vals[scale_idx] + for idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + s_arr[idx] = s_vals[idx] + return old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr), + ) + @cute.jit + def update_softmax_sum( + self, + stage_info: StageInfo, + *, + old_max_arr, + sum_arr, + new_max_arr, + local_sum_arr, + s_arr, + ): + """Apply the online-softmax sum correction after P is materialized.""" + # This consumer aux step runs after SmemP/TmemP has produced P and + # published local_sum_arr. It updates the running denominator for the + # next score tile without consuming a new score payload. + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + state_idx = cute.arch.thread_idx()[0] + state_ptr = self._softmax_scratch.data_ptr(state_idx) + old_max = old_max_arr[0] + new_max = u32_bits_to_float(state_ptr.load(is_volatile=True, alignment=4)) + updated_sum = u32_bits_to_float( + softmax_sum_state_ptr(state_ptr).load(is_volatile=True, alignment=4) + ) + prims.barrier_cta_sync( + barrier_id=self.sync_barrier_id, + thread_count=WARPGROUP_THREADS, + ) + self.new_max_state[0] = new_max + self.sum_state[0] = updated_sum + old_max_arr[0] = old_max + sum_arr[0] = updated_sum + new_max_arr[0] = new_max + return old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr + for scale_base in cutlass.range_constexpr( + 0, num_softmax_scale_groups(self.cfg), 2 + ): + old_max_0 = old_max_arr[scale_base] + old_max_1 = old_max_arr[scale_base + 1] + new_max_0 = new_max_arr[scale_base] + new_max_1 = new_max_arr[scale_base + 1] + local_sum_0 = local_sum_arr[scale_base] + local_sum_1 = local_sum_arr[scale_base + 1] + if cutlass.const_expr(self.p_ref is not None): + local_sum_0 = self._p_local_sum_arr[scale_base] + local_sum_1 = self._p_local_sum_arr[scale_base + 1] + sum_0 = sum_arr[scale_base] + sum_1 = sum_arr[scale_base + 1] + + exp_scale0 = Float32(0.0) + exp_scale1 = Float32(0.0) + if (old_max_0 != neg_max_f32()) and (new_max_0 != neg_max_f32()): + exp_scale0 = cute.math.exp2( + self.scale_softmax_log2 * (old_max_0 - new_max_0), + fastmath=True, + ) + if (old_max_1 != neg_max_f32()) and (new_max_1 != neg_max_f32()): + exp_scale1 = cute.math.exp2( + self.scale_softmax_log2 * (old_max_1 - new_max_1), + fastmath=True, + ) + updated_sums = ffma2( + (exp_scale0, exp_scale1), + (sum_0, sum_1), + (local_sum_0, local_sum_1), + ) + self.sum_state[scale_base] = updated_sums[0] + self.sum_state[scale_base + 1] = updated_sums[1] + sum_arr[scale_base] = updated_sums[0] + sum_arr[scale_base + 1] = updated_sums[1] + old_max_arr[scale_base] = old_max_0 + old_max_arr[scale_base + 1] = old_max_1 + new_max_arr[scale_base] = new_max_0 + new_max_arr[scale_base + 1] = new_max_1 + local_sum_arr[scale_base] = local_sum_0 + local_sum_arr[scale_base + 1] = local_sum_1 + return old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr + + +# ===================================================================== +# TmemSKeepsResource — Keeps-MMA-AB S scores in TMEM +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemSKeepsResource(TmemSResource): + """Keeps-MMA-AB score resource with a single softmax scale group.""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py new file mode 100644 index 000000000000..d9c83631fbec --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py @@ -0,0 +1,563 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TMEM resources for local and global softmax statistics.""" + +from dataclasses import dataclass +from typing import ClassVar, Optional + +from cutlass.experimental import primitives as prims + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import TmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...helpers.constants import ( + TCGEN05_32B_SHAPE, +) + +from ...helpers.layout import ( + _TASK_CACHE_TMEM_BASE_OFFSET, + decode_gen_task_cache, + num_softmax_scale_groups, +) +from ...helpers.math import ( + neg_max_f32, +) +from ...helpers.ops import ( + vector_from_scalars, +) + +from .common import ( + MlaResource, +) + +# ===================================================================== +# TmemSoftmaxLocalResource — Local softmax stats in TMEM +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemSoftmaxLocalResource(MlaResource): + """TMEM scratch resource for local softmax max/sum statistics.""" + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("old_max_arr", cutlass.Array, None, "Previous softmax maxima read from TMEM."), + ("new_max_arr", cutlass.Array, None, "Current softmax maxima read from TMEM."), + ("sum_arr", cutlass.Array, None, "Softmax denominators read from TMEM."), + ("inst0_old_max_arr", cutlass.Array, None, "Instance 0 previous maxima."), + ("inst0_new_max_arr", cutlass.Array, None, "Instance 0 current maxima."), + ("inst0_sum_arr", cutlass.Array, None, "Instance 0 denominator sums."), + ("inst1_old_max_arr", cutlass.Array, None, "Instance 1 previous maxima."), + ("inst1_new_max_arr", cutlass.Array, None, "Instance 1 current maxima."), + ("inst1_sum_arr", cutlass.Array, None, "Instance 1 denominator sums."), + ) + inst_id: cutlass.Constexpr[int] = 0 + tmem_alias_ref: Optional[MlaResource] = None + old_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + new_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + sum_arr: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + inst0_old_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + inst0_new_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + inst0_sum_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + inst1_old_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + inst1_new_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + inst1_sum_arr: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + + def get_tmem_requirements(self): + """Return the TMEM allocation for local softmax statistics.""" + if self.tmem_alias_ref is not None: + return [] + if self._tmem_alloc is None: + num_stages = ( + self.pipeline_config.num_stages + if self.cfg.kernel_variant == "keeps_mma_ab" + and self.pipeline_config is not None + else 1 + ) + self._tmem_alloc = TmemAllocation( + name=f"{self.name}", + num_columns=self.cfg.tmem_stats_cols * num_stages, + ) + return [self._tmem_alloc] + + @cute.jit + def _stats_base_addr(self, stage_info: StageInfo, task_cache): + """Return the TMEM address that stores this stage's softmax stats.""" + if cutlass.const_expr(self.tmem_alias_ref is not None): + return ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self.tmem_alias_ref._tmem_alloc.offset) + + stage_info.stage_idx * Int32(self.cfg.tmem_s_cols) + ) + stats_stage_offset = Int32(0) + if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + stats_stage_offset = stage_info.stage_idx * Int32(self.cfg.tmem_stats_cols) + return ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._tmem_alloc.offset) + + stats_stage_offset + ) + + @cute.jit + def _make_initial_stats_vars(self): + """Create and initialize softmax statistic arrays.""" + num_scale_groups = num_softmax_scale_groups(self.cfg) + old_max_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + new_max_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + sum_arr = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) + inst0_old_max_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + inst0_new_max_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + inst0_sum_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + inst1_old_max_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + inst1_new_max_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + inst1_sum_arr = cutlass.Array( + Float32, + num_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + for idx in cutlass.range_constexpr(num_scale_groups): + old_max_arr[idx] = neg_max_f32() + new_max_arr[idx] = neg_max_f32() + sum_arr[idx] = Float32(0.0) + inst0_old_max_arr[idx] = neg_max_f32() + inst0_new_max_arr[idx] = neg_max_f32() + inst0_sum_arr[idx] = Float32(0.0) + inst1_old_max_arr[idx] = neg_max_f32() + inst1_new_max_arr[idx] = neg_max_f32() + inst1_sum_arr[idx] = Float32(0.0) + return ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ), + ) + @cute.jit + def init_stats_state(self, stage_info: StageInfo): + """Create statistic variables for the first work tile.""" + # Consumer aux work creates the arrays that hold stats loaded from + # TMEM. The arrays are later populated by the loop/tail stat loads. + self._init_tmem_state(stage_info) + return self._make_initial_stats_vars() + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ), + ) + @cute.jit + def init_stats_work_tile_state(self, stage_info: StageInfo): + """Create statistic variables for each persistent work tile.""" + # Persistent CTAs reuse the same resource instance, so local stats must + # be reset when the work tile changes. + del stage_info + return self._make_initial_stats_vars() + + @producer_work + @cute.jit + def store_loop_stats( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst_idx: cutlass.Constexpr[int], + ): + """Store loop softmax statistics into TMEM for correction.""" + # Producer work for stats: softmax writes old/new max or sum/new max + # snapshots into the acquired TMEM stats stage. Correction consumes the + # same stage through the matching loop/tail load function. + cfg = self.cfg + num_scale_groups = num_softmax_scale_groups(cfg) + task_cache = decode_gen_task_cache(stage_info) + stats_base = self._stats_base_addr(stage_info, task_cache) + stats_ptr = prims.make_tmem_ptr(stats_base, Float32) + skip_loop_stats_store = cutlass.const_expr(inst_idx == 0) and ( + stage_info.loop_end == Int32(1) + ) + if not skip_loop_stats_store: + if cutlass.const_expr(inst_idx == 0): + stats = vector_from_scalars( + tuple(old_max_arr[idx] for idx in range(num_scale_groups)) + + tuple(new_max_arr[idx] for idx in range(num_scale_groups)), + Float32, + ) + prims.tcgen05_st( + TCGEN05_32B_SHAPE, + stats_ptr, + stats, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + elif cutlass.const_expr(inst_idx == 1): + stats = vector_from_scalars( + tuple(sum_arr[idx] for idx in range(num_scale_groups)) + + tuple(new_max_arr[idx] for idx in range(num_scale_groups)), + Float32, + ) + prims.tcgen05_st( + TCGEN05_32B_SHAPE, + stats_ptr, + stats, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def _stats_tuple( + self, + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ): + """Return the task-local stat arrays in schedule order.""" + return ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) + + @cute.jit + def _load_stats_payload(self, stage_info: StageInfo): + """Load the raw stat payload from the current TMEM stats stage.""" + cfg = self.cfg + task_cache = decode_gen_task_cache(stage_info) + stats_base = self._stats_base_addr(stage_info, task_cache) + stats_ptr = prims.make_tmem_ptr(stats_base, Float32) + loaded = prims.tcgen05_ld( + TCGEN05_32B_SHAPE, + stats_ptr, + num=num_softmax_scale_groups(cfg) * 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + return loaded + + @cute.jit + def _store_loop_stats_payload( + self, + loaded, + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ): + """Update task-local arrays from loop old/new-max stats.""" + cfg = self.cfg + for scale_idx in cutlass.range_constexpr(num_softmax_scale_groups(cfg)): + loaded_new_max = loaded[num_softmax_scale_groups(cfg) + scale_idx] + new_max_arr[scale_idx] = loaded_new_max + loaded_old_max = loaded[scale_idx] + old_max_arr[scale_idx] = loaded_old_max + if cutlass.const_expr(self.inst_id == 0): + inst0_old_max_arr[scale_idx] = loaded_old_max + inst0_new_max_arr[scale_idx] = loaded_new_max + else: + inst1_old_max_arr[scale_idx] = loaded_old_max + inst1_new_max_arr[scale_idx] = loaded_new_max + return self._stats_tuple( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) + + @cute.jit + def _store_tail_stats_payload( + self, + loaded, + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ): + """Update task-local arrays from tail sum/new-max stats.""" + cfg = self.cfg + for scale_idx in cutlass.range_constexpr(num_softmax_scale_groups(cfg)): + loaded_new_max = loaded[num_softmax_scale_groups(cfg) + scale_idx] + loaded_sum = loaded[scale_idx] + new_max_arr[scale_idx] = loaded_new_max + sum_arr[scale_idx] = loaded_sum + if cutlass.const_expr(self.inst_id == 0): + inst0_new_max_arr[scale_idx] = loaded_new_max + inst0_sum_arr[scale_idx] = loaded_sum + else: + inst1_new_max_arr[scale_idx] = loaded_new_max + inst1_sum_arr[scale_idx] = loaded_sum + return self._stats_tuple( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) + + @consumer_work( + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ), + ) + @cute.jit + def load_initial_stats( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ): + """Return initial stats before any loop-stage TMEM payload exists.""" + del stage_info + return self._stats_tuple( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) + + @consumer_work( + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ), + ) + @cute.jit + def load_loop_stats( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ): + """Load loop old/new-max statistics for O rescaling.""" + loaded = self._load_stats_payload(stage_info) + return self._store_loop_stats_payload( + loaded, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + + @consumer_work( + returns=( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ), + ) + @cute.jit + def load_tail_stats( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ): + """Load tail sum/new-max statistics for final normalization.""" + loaded = self._load_stats_payload(stage_info) + return self._store_tail_stats_payload( + loaded, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + + +# ===================================================================== +# TmemSoftmaxGlobalResource — Global softmax dependency marker +# ===================================================================== + + +@dataclass(kw_only=True) +class TmemSoftmaxGlobalResource(MlaResource): + """Named dependency edge from local softmax stats to correction.""" + + inst_id: cutlass.Constexpr[int] = 0 + + @producer_work + @cute.jit + def track_global(self, stage_info: StageInfo): + """No-op placeholder for the global softmax dependency edge.""" + # This resource has no payload. Its producer work gives TS a named + # dependency edge between the two softmax instances and correction. + del stage_info diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py new file mode 100644 index 000000000000..e3d2e64eb937 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py @@ -0,0 +1,1707 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Captured task schedules for the throughput-latency 1CTA MLA TS path.""" + +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.experimental import primitives as prims + +from cutlass.experimental.task_scheduling.resources import WorkQueue +from cutlass.experimental.task_scheduling.schedule_builder import ( + domain_loop, + schedule, + work_tile_loop, +) +from cutlass.experimental.task_scheduling.task import Task + +from ..helpers.constants import ( + WARP_LANES, + WARP_LANE_MASK, + WARP_LANE_SHIFT, + WARPGROUP_WARPS, +) +from ..helpers.schedule import ( + page_offsets_produce, + staged_kv_load, + staged_pv_mma, + staged_pv_mma_tmem_p, + staged_qk_mma, + runtime_work_tile_skip_if, + schedule_token_throttle_head, + schedule_token_throttle_tail, + work_tile_schedule_loop, +) +from ..helpers.stage import MlaStage +from ..helpers.tile import ( + runtime_base_seq_len_kv, + runtime_local_kv_tiles, + runtime_seq_len_kv_from_task_cache, +) + + +class MlaDecodeTask(Task): + """MLA decode task with cached thread state and Q-aware KV domain logic.""" + + def __init__(self, **kwargs): + """Initialize MLA-specific task parameters and cached thread indices.""" + self.cfg = kwargs.pop("cfg", None) + self.seqlens_kv = kwargs.pop("seqlens_kv", None) + self.cu_seqlens_q = kwargs.pop("cu_seqlens_q", None) + self.domain_bias = kwargs.pop("domain_bias", 0) + super().__init__(**kwargs) + self._tmem_base_offset = Int32(0) + self._warp_grp_thread_idx = Int32(0) + self._local_warp_idx = Int32(0) + self._lane_idx = Int32(0) + self._seq_len_kv = Int32(0) + + def init_variables(self, context=None): + """Initialize per-task thread cache and TMEM base state.""" + super().init_variables(context) + tidx, _, _ = cute.arch.thread_idx() + warp_grp_start = Int32( + (self.warp_idx // WARPGROUP_WARPS) * WARPGROUP_WARPS * WARP_LANES + ) + self._warp_grp_thread_idx = tidx - warp_grp_start + self._local_warp_idx = self._warp_grp_thread_idx >> Int32(WARP_LANE_SHIFT) + self._lane_idx = self._warp_grp_thread_idx & Int32(WARP_LANE_MASK) + if context is not None and context.tmem_ptr_i32 is not None: + loaded = Int32(context.tmem_ptr_i32.load()) + self._tmem_base_offset = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=loaded, + offset=0, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + + @cute.jit + def make_task_cache(self): + """Return the compact tuple of cached per-task thread values.""" + return ( + self._tmem_base_offset, + self._warp_grp_thread_idx, + self._local_warp_idx, + self._lane_idx, + self._seq_len_kv, + Int32(0), + Int32(0), + Int32(0), + ) + + def get_domain(self, tile_coord): + """Return the runtime loop domain for the current 1CTA work tile.""" + if self.cfg is None: + return self.domain + + # Persistent 1CTA coordinates combine batch and head tile in z; + # non-persistent grids keep z as batch. cache_seqs remains batch-owned. + if isinstance(tile_coord[2], int): + if self.cfg.use_persistent_scheduler == 1: + batch_idx = tile_coord[2] // self.cfg.num_ctas_for_all_heads + else: + batch_idx = tile_coord[2] + else: + batch_idx = Int32(tile_coord[2]) + if self.cfg.use_persistent_scheduler == 1: + batch_idx = batch_idx // Int32(self.cfg.num_ctas_for_all_heads) + + # Task-domain ownership must match K/V, softmax, and reduction: dense + # uses the full runtime K length, while causal uses the largest + # logical-Q-visible length in this physical flat-query tile. Using the + # raw batch length for causal can move an all-masked tile into tail and + # replace a real loop accumulator at a tile boundary. + cta_idx_q = tile_coord[0] + if self.cfg.use_persistent_scheduler != 1: + # Nonpersistent grid X combines (head tile, Q tile, KV split), with + # KV split innermost. Persistent work queues already expose Q as + # tile coordinate 0 and must not be decoded a second time. + cta_idx_q = (tile_coord[0] // Int32(self.cfg.num_ctas_per_seq_kv)) % Int32( + self.cfg.num_ctas_per_seq_q + ) + # Cache the raw batch length once per task/work tile. Resource work + # methods derive their CTA-visible length from this task cache instead + # of independently reloading cache_seqs throughout the hot loop. + self._seq_len_kv = runtime_base_seq_len_kv( + self.cfg, + self.seqlens_kv, + batch_idx, + ) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + self.cfg, + self.make_task_cache(), + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + total_kv_tiles = runtime_local_kv_tiles(self.cfg, seq_len_kv) + remaining_kv_tiles = cute.math.max( + total_kv_tiles - Int32(self.cfg.num_insts_kv), Int32(0) + ) + num_insts_kv = Int32(self.cfg.num_insts_kv) + loop_domain = (remaining_kv_tiles + num_insts_kv - Int32(1)) // num_insts_kv + return loop_domain + Int32(self.domain_bias) + + +def create_throughput_latency_softmax_task_impl( + tmem_s, + tmem_softmax_local, + p_resource, + tmem_softmax_global, + work_queue: WorkQueue | None, + cfg, + *, + inst_id: int, + domain, + store_stats_before_p: bool = True, + task_name: str | None = None, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create one softmax task for a throughput-latency 1CTA score pipe. + + The task consumes TMEM scores, updates online-softmax state, materializes + P for PV MMA, and stores per-iteration statistics for correction. + """ + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def softmax_schedule( + tmem_s, + tmem_softmax_local, + p_resource, + tmem_softmax_global, + work_queue=None, + ): + """Captured softmax schedule for one of the two interleaved instances.""" + ( + old_max_arr, + sum_arr, + new_max_arr, + local_sum_arr, + s_arr, + ) = tmem_s.init_softmax_state() + if store_stats_before_p: + p_resource.init_materialize_state() + + def store_stats(inst_idx: int): + """Store local online-softmax stats for correction.""" + tmem_softmax_local.acquire() + tmem_softmax_local.store_loop_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst_idx=inst_idx, + ) + tmem_softmax_local.commit() + + def init_work_tile_state(): + return tmem_s.init_softmax_work_tile_state() + + with work_tile_schedule_loop( + work_queue, + skip_if=work_tile_skip_if, + non_skippable_prelude=( + init_work_tile_state if work_queue is not None else None + ), + ) as (_, work_tile_state): + if work_queue is not None: + ( + old_max_arr, + sum_arr, + new_max_arr, + local_sum_arr, + s_arr, + ) = work_tile_state + with domain_loop(0, domain, 1) as d: + tmem_s.wait() + ( + old_max_arr, + sum_arr, + new_max_arr, + local_sum_arr, + s_arr, + ) = tmem_s.update_softmax( + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + local_sum_arr=local_sum_arr, + s_arr=s_arr, + section=MlaStage.Loop, + ) + tmem_s.release() + if store_stats_before_p: + store_stats(0) + # AsyncUmma protects P until PV MMA releases this stage. + p_resource.acquire() + p_resource.materialize_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + local_sum_arr=local_sum_arr, + ) + p_resource.commit() + tmem_softmax_global.track_global() + ( + old_max_arr, + sum_arr, + new_max_arr, + local_sum_arr, + s_arr, + ) = tmem_s.update_softmax_sum( + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + local_sum_arr=local_sum_arr, + s_arr=s_arr, + ) + if not store_stats_before_p: + store_stats(0) + with d.last_iter(): + store_stats(1) + + captured_schedule = ( + softmax_schedule( + tmem_s, + tmem_softmax_local, + p_resource, + tmem_softmax_global, + ) + if work_queue is None + else softmax_schedule( + tmem_s, + tmem_softmax_local, + p_resource, + tmem_softmax_global, + work_queue, + ) + ) + src = [tmem_s] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_softmax_local, p_resource, tmem_softmax_global], + cfg=cfg, + warp_idx=cfg.softmax0_warp_idx if inst_id == 0 else cfg.softmax1_warp_idx, + num_warps=cfg.softmax0_num_warps if inst_id == 0 else cfg.softmax1_num_warps, + schedule=captured_schedule, + num_registers=cfg.softmax_regs, + name=task_name or f"Softmax{inst_id}Task", + **kw, + ) + + +def create_keeps_mma_ab_softmax_task( + tmem_s, + tmem_softmax_local, + tmem_p, + tmem_softmax_global, + work_queue: WorkQueue | None, + cfg, + *, + domain, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the single-pipe softmax/P task for the keeps-MMA-AB path.""" + return create_throughput_latency_softmax_task_impl( + tmem_s, + tmem_softmax_local, + tmem_p, + tmem_softmax_global, + work_queue, + cfg, + inst_id=0, + domain=domain, + store_stats_before_p=False, + task_name="SoftmaxTask", + task_class=task_class, + **kw, + ) + + +def create_keeps_mma_ab_correction_task( + tmem_softmax_local, + tmem_o, + tmem_corr, + work_queue: WorkQueue | None, + cfg, + *, + domain, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the single-pipe correction/store task for keeps-MMA-AB.""" + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def correction_schedule( + tmem_softmax_local, + tmem_o, + tmem_corr, + work_queue=None, + ): + """Captured keeps-MMA-AB correction schedule for loop and tail.""" + + def load_initial_local_stats(local_state): + """Wait for the keeps-MMA-AB pipe and return its initial stats.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_softmax_local.wait() + local_state = tmem_softmax_local.load_initial_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + tmem_softmax_local.release() + return local_state + + def load_loop_local_stats(local_state): + """Wait for the keeps-MMA-AB pipe and load loop stats.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_softmax_local.wait() + local_state = tmem_softmax_local.load_loop_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + tmem_softmax_local.release() + return local_state + + def load_tail_local_stats(local_state): + """Wait for the keeps-MMA-AB pipe and load tail stats.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_softmax_local.wait() + local_state = tmem_softmax_local.load_tail_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + tmem_softmax_local.release() + return local_state + + def correct_loop_o(local_state, tail_o_stage_idx_0, tail_o_stage_idx_1): + """Rescale one keeps-MMA-AB loop O stage.""" + local_state = load_loop_local_stats(local_state) + ( + old_max_arr, + new_max_arr, + sum_arr, + _inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + _inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_o.wait() + ( + o_stage_idx, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = tmem_o.o_stage( + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst_idx=0, + ) + tmem_corr.correct_loop_and_store( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + ) + tmem_o.release() + return local_state, tail_o_stage_idx_0, tail_o_stage_idx_1 + + def correct_tail_o(local_state, tail_o_stage_idx_0, tail_o_stage_idx_1): + """Normalize and store the final keeps-MMA-AB O stage.""" + local_state = load_tail_local_stats(local_state) + ( + old_max_arr, + new_max_arr, + sum_arr, + _inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + _inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_o.wait() + ( + o_stage_idx, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = tmem_o.o_stage( + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst_idx=0, + is_tail=True, + ) + tmem_corr.correct_tail_and_store( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + ) + tmem_o.release() + return local_state, tail_o_stage_idx_0, tail_o_stage_idx_1 + + local_state = tmem_softmax_local.init_stats_state() + _, tail_o_stage_idx_0, tail_o_stage_idx_1 = tmem_o.init_stage_state() + tmem_corr.init_store_state() + + def init_work_tile_state(): + local_state = tmem_softmax_local.init_stats_work_tile_state() + _, tail_stage_0, tail_stage_1 = tmem_o.init_stage_work_tile_state() + return local_state, tail_stage_0, tail_stage_1 + + with work_tile_schedule_loop( + work_queue, + skip_if=work_tile_skip_if, + non_skippable_prelude=( + init_work_tile_state if work_queue is not None else None + ), + ) as (_, work_tile_state): + if work_queue is not None: + ( + local_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = work_tile_state + local_state = load_initial_local_stats(local_state) + with domain_loop(0, domain, 1): + ( + local_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = correct_loop_o(local_state, tail_o_stage_idx_0, tail_o_stage_idx_1) + + ( + local_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = correct_tail_o(local_state, tail_o_stage_idx_0, tail_o_stage_idx_1) + + schedule_result = ( + correction_schedule(tmem_softmax_local, tmem_o, tmem_corr) + if work_queue is None + else correction_schedule( + tmem_softmax_local, + tmem_o, + tmem_corr, + work_queue, + ) + ) + src = [tmem_softmax_local, tmem_o] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_corr], + cfg=cfg, + warp_idx=cfg.correction_warp_idx, + num_warps=cfg.correction_num_warps, + schedule=schedule_result, + num_registers=cfg.correction_regs, + name="CorrectionTask", + **kw, + ) + + +def create_throughput_latency_softmax0_task( + tmem_s0, + tmem_softmax_local0, + smem_p0, + tmem_softmax_global0, + work_queue: WorkQueue | None, + cfg, + *, + domain, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the softmax/P producer task for score instance 0.""" + return create_throughput_latency_softmax_task_impl( + tmem_s0, + tmem_softmax_local0, + smem_p0, + tmem_softmax_global0, + work_queue, + cfg, + inst_id=0, + domain=domain, + task_class=task_class, + **kw, + ) + + +def create_throughput_latency_softmax1_task( + tmem_s1, + tmem_softmax_local1, + smem_p1, + tmem_softmax_global1, + work_queue: WorkQueue | None, + cfg, + *, + domain, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the softmax/P producer task for score instance 1.""" + return create_throughput_latency_softmax_task_impl( + tmem_s1, + tmem_softmax_local1, + smem_p1, + tmem_softmax_global1, + work_queue, + cfg, + inst_id=1, + domain=domain, + task_class=task_class, + **kw, + ) + + +def create_throughput_latency_correction_task( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + work_queue: WorkQueue | None, + cfg, + *, + domain, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the correction task that normalizes and stores 1CTA output. + + The task consumes local softmax statistics and TMEM O stages from both + interleaved MMA pipes, applies online-softmax correction, and stores O/LSE + or split-KV partials. + """ + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + work_queue=None, + ): + """Captured correction schedule for loop and tail O-stage draining.""" + + def load_initial_local_stats(tmem_softmax_local, local_state): + """Wait for one softmax pipe and return its initial stat state.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_softmax_local.wait() + local_state = tmem_softmax_local.load_initial_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + tmem_softmax_local.release() + return local_state + + def load_loop_local_stats(tmem_softmax_local, local_state): + """Wait for one softmax pipe and load loop old/new stats.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_softmax_local.wait() + local_state = tmem_softmax_local.load_loop_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + tmem_softmax_local.release() + return local_state + + def load_tail_local_stats(tmem_softmax_local, local_state): + """Wait for one softmax pipe and load tail sum/new-max stats.""" + ( + old_max_arr, + new_max_arr, + sum_arr, + inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_softmax_local.wait() + local_state = tmem_softmax_local.load_tail_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_old_max_arr=inst0_old_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_old_max_arr=inst1_old_max_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) + tmem_softmax_local.release() + return local_state + + def correct_loop_o( + tmem_softmax_local, + tmem_corr, + local_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + o_stage_inst_idx, + ): + """Rescale one loop O accumulator using local softmax stats.""" + local_state = load_loop_local_stats(tmem_softmax_local, local_state) + ( + old_max_arr, + new_max_arr, + sum_arr, + _inst0_old_max_arr, + inst0_new_max_arr, + inst0_sum_arr, + _inst1_old_max_arr, + inst1_new_max_arr, + inst1_sum_arr, + ) = local_state + tmem_o.wait() + ( + o_stage_idx, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = tmem_o.o_stage( + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst_idx=o_stage_inst_idx, + ) + tmem_corr.correct_loop_and_store( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + ) + tmem_o.release() + return local_state, tail_o_stage_idx_0, tail_o_stage_idx_1 + + local0_state = tmem_softmax_local0.init_stats_state() + local1_state = tmem_softmax_local1.init_stats_state() + _, tail_o_stage_idx_0, tail_o_stage_idx_1 = tmem_o.init_stage_state() + tmem_corr0.init_store_state() + tmem_corr1.init_store_state() + + def init_work_tile_state(): + local0_state = tmem_softmax_local0.init_stats_work_tile_state() + local1_state = tmem_softmax_local1.init_stats_work_tile_state() + _, tail_stage_0, tail_stage_1 = tmem_o.init_stage_work_tile_state() + return local0_state, local1_state, tail_stage_0, tail_stage_1 + + with work_tile_schedule_loop( + work_queue, + skip_if=work_tile_skip_if, + non_skippable_prelude=( + init_work_tile_state if work_queue is not None else None + ), + ) as (_, work_tile_state): + if work_queue is not None: + ( + local0_state, + local1_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = work_tile_state + local0_state = load_initial_local_stats(tmem_softmax_local0, local0_state) + local1_state = load_initial_local_stats(tmem_softmax_local1, local1_state) + + with domain_loop(0, domain, 1): + ( + local0_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = correct_loop_o( + tmem_softmax_local0, + tmem_corr0, + local0_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + 0, + ) + ( + local1_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = correct_loop_o( + tmem_softmax_local1, + tmem_corr1, + local1_state, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + 1, + ) + + stats0 = load_tail_local_stats(tmem_softmax_local0, local0_state) + ( + _old_max_arr0, + _new_max_arr0, + _sum_arr0, + _inst0_old_max_arr0, + inst0_new_max_arr0, + inst0_sum_arr0, + _inst1_old_max_arr0, + _inst1_new_max_arr0, + _inst1_sum_arr0, + ) = stats0 + tmem_o.wait() + ( + o_stage_idx, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = tmem_o.o_stage( + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst_idx=0, + is_tail=True, + ) + # Loop-body correction already rescaled stream 0 to the final + # stream-0 max before the tail PV0 accumulation. Keep the tail + # stage index for final stream combination, but avoid a second + # rescale of the same O0 data. + del o_stage_idx + stats1 = load_tail_local_stats(tmem_softmax_local1, local1_state) + ( + old_max_arr1, + new_max_arr1, + sum_arr1, + _inst0_old_max_arr1, + _inst0_new_max_arr1, + _inst0_sum_arr1, + _inst1_old_max_arr1, + inst1_new_max_arr1, + inst1_sum_arr1, + ) = stats1 + tmem_o.wait() + ( + o_stage_idx, + tail_o_stage_idx_0, + tail_o_stage_idx_1, + ) = tmem_o.o_stage( + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + inst_idx=1, + is_tail=True, + ) + tmem_corr1.correct_tail_and_store( + old_max_arr=old_max_arr1, + new_max_arr=new_max_arr1, + sum_arr=sum_arr1, + inst0_new_max_arr=inst0_new_max_arr0, + inst0_sum_arr=inst0_sum_arr0, + inst1_new_max_arr=inst1_new_max_arr1, + inst1_sum_arr=inst1_sum_arr1, + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_o_stage_idx_0, + tail_o_stage_idx_1=tail_o_stage_idx_1, + ) + tmem_o.release() + tmem_o.release() + + captured_schedule = ( + correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + ) + if work_queue is None + else correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + work_queue, + ) + ) + src = [tmem_softmax_local0, tmem_softmax_local1, tmem_o] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_corr0, tmem_corr1], + cfg=cfg, + warp_idx=cfg.correction_warp_idx, + num_warps=cfg.correction_num_warps, + schedule=captured_schedule, + num_registers=cfg.correction_regs, + name="CorrectionTask", + **kw, + ) + + +def _make_page_offsets_task( + smem_page_offsets, + work_queue: WorkQueue | None, + cfg, + *, + schedule_result, + warp_idx=None, + num_warps=None, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create a page-offset staging task from a captured schedule result.""" + src = [] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[smem_page_offsets], + cfg=cfg, + warp_idx=cfg.page_offsets_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.page_offsets_num_warps if num_warps is None else num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_regs, + name="LoadPageTableTask", + **kw, + ) + + +def create_load_page_offsets_task( + smem_page_offsets, + work_queue: WorkQueue | None, + cfg, + *, + domain, + warp_idx=None, + num_warps=None, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the page-offset staging task for the selected 1CTA cadence. + + Swaps-MMA-AB stages only K0/K1 offsets because each delayed V tile shares + its K tile's page map and reuses the register-cached IDs. Keeps-MMA-AB + retains its independent K/V offset cadence. The produced offset stages are + consumed by the generic load task. + """ + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def page_offsets_schedule(smem_page_offsets, work_queue=None): + """Stage page offsets in the selected K-before-V cadence.""" + smem_page_offsets.init_load_state() + + with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): + if cfg.kernel_variant == "keeps_mma_ab": + page_offsets_produce( + smem_page_offsets, "load_k0", section=MlaStage.Head + ) + with domain_loop(0, domain, 1): + page_offsets_produce( + smem_page_offsets, "load_k0", section=MlaStage.Loop + ) + page_offsets_produce( + smem_page_offsets, "load_v0", section=MlaStage.Loop + ) + page_offsets_produce( + smem_page_offsets, "load_v0", section=MlaStage.Tail + ) + else: + page_offsets_produce( + smem_page_offsets, "load_k0", section=MlaStage.Head + ) + page_offsets_produce( + smem_page_offsets, "load_k1", section=MlaStage.Head + ) + + with domain_loop(0, domain, 1): + page_offsets_produce( + smem_page_offsets, "load_k0", section=MlaStage.Loop + ) + page_offsets_produce( + smem_page_offsets, "load_k1", section=MlaStage.Loop + ) + + schedule_result = ( + page_offsets_schedule(smem_page_offsets) + if work_queue is None + else page_offsets_schedule(smem_page_offsets, work_queue) + ) + return _make_page_offsets_task( + smem_page_offsets, + work_queue, + cfg, + schedule_result=schedule_result, + warp_idx=warp_idx, + num_warps=num_warps, + task_class=task_class, + **kw, + ) + + +def _staged_kv_load_with_reused_page_ids( + smem_kv, + *, + head_dim_stages, + producer_label, + section: MlaStage, + smem_page_offsets, + cached_page_ids, + page_id_slot: int, + consume_offsets: bool, +): + """Stage one swaps-MMA K/V tile using its K-owned page-ID cache slot.""" + if consume_offsets: + smem_page_offsets.wait() + cached_page_ids = smem_page_offsets.read_offsets( + cached_page_ids=cached_page_ids, + cache_slot=page_id_slot, + ) + for stage_idx in range(head_dim_stages): + smem_kv.acquire() + getattr(smem_kv, producer_label)( + stage_idx=stage_idx, + section=section, + cached_page_ids=cached_page_ids, + page_id_slot=page_id_slot, + ) + smem_kv.commit() + if consume_offsets: + smem_page_offsets.release() + return cached_page_ids + + +def create_throughput_latency_load_task( + smem_q, + smem_kv, + work_queue: WorkQueue | None, + schedule_token_throttle, + cfg, + *, + domain, + smem_page_offsets=None, + use_page_offsets=False, + warp_idx=None, + num_warps=None, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the Q/K/V load task. + + The load warp stages Q once, then follows the selected K-before-V cadence. + Swaps-MMA-AB uses two interleaved K/V streams, while keeps-MMA-AB uses a + single K stream followed by the deferred V stream. + """ + page_offsets = smem_page_offsets if use_page_offsets else None + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + def load_schedule_body( + smem_q, + smem_kv, + smem_page_offsets, + work_queue, + schedule_token_throttle, + ): + """Shared captured load schedule body for optional page offsets/WQ paths.""" + smem_q.init_load_state() + smem_kv.init_load_state() + if smem_page_offsets is not None: + cached_page_ids = smem_page_offsets.init_read_state() + else: + cached_page_ids = None + reuse_delayed_v_page_ids = ( + smem_page_offsets is not None and cfg.kernel_variant != "keeps_mma_ab" + ) + + def load_kv( + *, + head_dim_stages, + producer_label, + section, + page_id_slot=0, + reuse_cached_page_ids=False, + ): + """Select the native paged swaps cache or the existing load path.""" + nonlocal cached_page_ids + if reuse_delayed_v_page_ids: + cached_page_ids = _staged_kv_load_with_reused_page_ids( + smem_kv, + head_dim_stages=head_dim_stages, + producer_label=producer_label, + section=section, + smem_page_offsets=smem_page_offsets, + cached_page_ids=cached_page_ids, + page_id_slot=page_id_slot, + consume_offsets=not reuse_cached_page_ids, + ) + else: + cached_page_ids = staged_kv_load( + smem_kv, + head_dim_stages=head_dim_stages, + producer_label=producer_label, + section=section, + smem_page_offsets=smem_page_offsets, + cached_page_ids=cached_page_ids, + ) + + with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): + schedule_token_throttle_head(schedule_token_throttle) + smem_q.acquire() + smem_q.load_q() + smem_q.commit() + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Head, + page_id_slot=0, + ) + if cfg.kernel_variant != "keeps_mma_ab": + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Head, + page_id_slot=1, + ) + + with domain_loop(0, domain, 1): + if cfg.kernel_variant == "keeps_mma_ab": + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Loop, + ) + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v0", + section=MlaStage.Loop, + ) + else: + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v0", + section=MlaStage.Loop, + page_id_slot=0, + reuse_cached_page_ids=True, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Loop, + page_id_slot=0, + ) + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v1", + section=MlaStage.Loop, + page_id_slot=1, + reuse_cached_page_ids=True, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Loop, + page_id_slot=1, + ) + + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v0", + section=MlaStage.Tail, + page_id_slot=0, + reuse_cached_page_ids=True, + ) + if cfg.kernel_variant != "keeps_mma_ab": + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v1", + section=MlaStage.Tail, + page_id_slot=1, + reuse_cached_page_ids=True, + ) + + @schedule + def load_schedule(smem_q, smem_kv): + """Captured load schedule without persistent work queue.""" + load_schedule_body(smem_q, smem_kv, None, None, None) + + @schedule + def load_wq_schedule(smem_q, smem_kv, work_queue): + """Captured load schedule with a persistent work queue.""" + load_schedule_body(smem_q, smem_kv, None, work_queue, None) + + @schedule + def load_wq_throttle_schedule( + smem_q, + smem_kv, + work_queue, + schedule_token_throttle, + ): + """Captured load schedule with persistent work queue throttling.""" + load_schedule_body( + smem_q, + smem_kv, + None, + work_queue, + schedule_token_throttle, + ) + + @schedule + def load_page_offsets_schedule( + smem_q, + smem_kv, + smem_page_offsets, + ): + """Captured load schedule with precomputed page offsets.""" + load_schedule_body( + smem_q, + smem_kv, + smem_page_offsets, + None, + None, + ) + + @schedule + def load_page_offsets_wq_schedule( + smem_q, + smem_kv, + smem_page_offsets, + work_queue, + ): + """Captured load schedule with page offsets and persistent work queue.""" + load_schedule_body( + smem_q, + smem_kv, + smem_page_offsets, + work_queue, + None, + ) + + @schedule + def load_page_offsets_wq_throttle_schedule( + smem_q, + smem_kv, + smem_page_offsets, + work_queue, + schedule_token_throttle, + ): + """Captured load schedule with page offsets, work queue, and throttling.""" + load_schedule_body( + smem_q, + smem_kv, + smem_page_offsets, + work_queue, + schedule_token_throttle, + ) + + if page_offsets is None: + if work_queue is None: + captured_schedule = load_schedule(smem_q, smem_kv) + elif schedule_token_throttle is None: + captured_schedule = load_wq_schedule(smem_q, smem_kv, work_queue) + else: + captured_schedule = load_wq_throttle_schedule( + smem_q, + smem_kv, + work_queue, + schedule_token_throttle, + ) + elif work_queue is None: + captured_schedule = load_page_offsets_schedule( + smem_q, + smem_kv, + page_offsets, + ) + elif schedule_token_throttle is None: + captured_schedule = load_page_offsets_wq_schedule( + smem_q, + smem_kv, + page_offsets, + work_queue, + ) + else: + captured_schedule = load_page_offsets_wq_throttle_schedule( + smem_q, + smem_kv, + page_offsets, + work_queue, + schedule_token_throttle, + ) + src = [] + if use_page_offsets: + src.append(smem_page_offsets) + if work_queue is not None: + src.append(work_queue) + dst = [smem_q, smem_kv] + if schedule_token_throttle is not None: + dst.append(schedule_token_throttle) + return task_class( + src_resources=src, + dst_resources=dst, + cfg=cfg, + warp_idx=cfg.load_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.load_num_warps if num_warps is None else num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_regs, + name="LoadTmaTask", + **kw, + ) + + +def create_throughput_latency_scheduler_task( + work_queue: WorkQueue, + schedule_token_throttle, + cfg, + *, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the CLC dynamic persistent scheduler task.""" + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def scheduler_schedule(work_queue, schedule_token_throttle=None): + """Fetch and publish the next dynamic persistent work tile.""" + + with work_tile_loop(work_queue, skip_if=work_tile_skip_if) as work_tiles: + if work_tile_skip_if is None: + with domain_loop(0, 0, 1): + pass + schedule_token_throttle_tail(schedule_token_throttle) + else: + # Runtime-empty packed-Q tiles do not publish a load token, but + # the scheduler must still fetch and advance past them. + with work_tiles.skippable(): + with domain_loop(0, 0, 1): + pass + schedule_token_throttle_tail(schedule_token_throttle) + work_queue.acquire() + work_queue.fetch_work_tile() + work_queue.commit() + work_queue.wait() + work_queue.get_and_advance_work_tile() + work_queue.release() + + captured_schedule = ( + scheduler_schedule(work_queue) + if schedule_token_throttle is None + else scheduler_schedule(work_queue, schedule_token_throttle) + ) + src = [work_queue] + if schedule_token_throttle is not None: + src.append(schedule_token_throttle) + return task_class( + src_resources=src, + dst_resources=[work_queue], + cfg=cfg, + warp_idx=cfg.scheduler_warp_idx, + num_warps=cfg.scheduler_num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_regs, + name="SchedulerTask", + **kw, + ) + + +def create_padding_task( + cfg, + work_queue: WorkQueue | None = None, + *, + warp_idx=None, + num_warps=None, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the padding task used to reserve otherwise idle warps.""" + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def padding_schedule(work_queue=None): + """Captured no-op schedule with optional persistent work-queue tail.""" + # Keep the task-scheduling and domain-loop scopes distinct for CuTe DSL. + with work_tile_schedule_loop( # noqa: SIM117 + work_queue, skip_if=work_tile_skip_if + ): + with domain_loop(0, 1, 1): + pass + + captured_schedule = ( + padding_schedule() if work_queue is None else padding_schedule(work_queue) + ) + src = [work_queue] if work_queue is not None else [] + return task_class( + src_resources=src, + dst_resources=[], + cfg=cfg, + warp_idx=cfg.padding_warp_idx if warp_idx is None else warp_idx, + num_warps=cfg.padding_num_warps if num_warps is None else num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_regs, + name="PaddingTask", + **kw, + ) + + +def create_throughput_latency_mma_task( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue: WorkQueue | None, + cfg, + *, + domain, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the MMA task for interleaved QK and PV stages. + + The task consumes Q/K/V SMEM descriptors, produces two TMEM score pipes, + materializes PV output into TMEM O, and drains the final V tiles in tail. + """ + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def mma_schedule( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue=None, + ): + """Captured MMA schedule with K-prefetch and V-tail drain.""" + smem_q.init_descriptor_state() + smem_kv.init_descriptor_state() + smem_p0.init_descriptor_state() + smem_p1.init_descriptor_state() + tmem_o.init_mma_state() + + with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): + if work_queue is not None: + tmem_s0.reset_softmax_work_tile_state() + tmem_s1.reset_softmax_work_tile_state() + smem_q.wait() + q_desc, q_desc_rope = smem_q.q_desc() + tmem_s0.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) + tmem_s1.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + ) + staged_qk_mma( + smem_kv, + tmem_s1, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_1", + ) + with domain_loop(0, domain, 1): + tmem_s0.acquire() + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_0", + producer_label="pv_mma_loop_0", + ) + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + include_acquire=False, + ) + tmem_s1.acquire() + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_1", + producer_label="pv_mma_loop_1", + ) + staged_qk_mma( + smem_kv, + tmem_s1, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_1", + include_acquire=False, + ) + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_0", + producer_label="pv_mma_tail_0", + ) + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_1", + producer_label="pv_mma_tail_1", + ) + smem_q.release() + + captured_schedule = ( + mma_schedule( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + ) + if work_queue is None + else mma_schedule( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue, + ) + ) + src = [smem_q, smem_kv, smem_p0, smem_p1] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_s0, tmem_s1, tmem_o], + cfg=cfg, + warp_idx=cfg.mma_warp_idx, + num_warps=cfg.mma_num_warps, + schedule=captured_schedule, + num_registers=cfg.mma_load_regs, + name="MmaTask", + **kw, + ) + + +def create_keeps_mma_ab_mma_task( + smem_q, + smem_kv, + tmem_s, + tmem_p, + tmem_o, + work_queue: WorkQueue | None, + cfg, + *, + domain, + task_class=MlaDecodeTask, + **kw, +) -> Task: + """Create the keeps-MMA-AB single-pipe MMA task. + + HEAD computes QK for K[0]. LOOP computes QK for K[n] before PV for + K[n-1], and TAIL drains PV for K[last]. + """ + + work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + + @schedule + def mma_schedule( + smem_q, + smem_kv, + tmem_s, + tmem_p, + tmem_o, + work_queue=None, + ): + """Captured keeps-MMA-AB schedule with TMEM P.""" + smem_q.init_descriptor_state() + smem_kv.init_descriptor_state() + tmem_o.init_mma_state() + tmem_p.init_stage_state() + + with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): + if work_queue is not None: + tmem_s.reset_softmax_work_tile_state() + tmem_p.init_stage_work_tile_state() + smem_q.wait() + q_desc, q_desc_rope = smem_q.q_desc() + tmem_s.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) + staged_qk_mma( + smem_kv, + tmem_s, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + ) + with domain_loop(0, domain, 1): + staged_qk_mma( + smem_kv, + tmem_s, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + ) + staged_pv_mma_tmem_p( + smem_kv, + tmem_p, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_0", + producer_label="pv_mma_loop_tmem_p", + ) + staged_pv_mma_tmem_p( + smem_kv, + tmem_p, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_0", + producer_label="pv_mma_tail_tmem_p", + ) + smem_q.release() + + schedule_result = ( + mma_schedule(smem_q, smem_kv, tmem_s, tmem_p, tmem_o) + if work_queue is None + else mma_schedule(smem_q, smem_kv, tmem_s, tmem_p, tmem_o, work_queue) + ) + src = [smem_q, smem_kv, tmem_p] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[tmem_s, tmem_o], + cfg=cfg, + warp_idx=cfg.mma_warp_idx, + num_warps=cfg.mma_num_warps, + schedule=schedule_result, + num_registers=cfg.mma_load_regs, + name="MmaTask", + **kw, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/placeholder_helpers.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/placeholder_helpers.py new file mode 100644 index 000000000000..1747dead7aa1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/placeholder_helpers.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Trace-time placeholder constructors shared by the FMHA TS kernels.""" + +import cutlass +from cutlass.experimental import primitives as prims + + +def _shape_tuple(shape: int | tuple[int, ...]) -> tuple[int, ...]: + """Normalize scalar and tuple shapes for ``cutlass.Array`` construction.""" + if isinstance(shape, tuple): + return shape + return (shape,) + + +def _placeholder_smem_array( + dtype: type, shape: int | tuple[int, ...] = 1 +) -> cutlass.Array | None: + """Build a typed shared-memory-view placeholder only when an MLIR context exists.""" + try: + return cutlass.Array( + cutlass.Int64(0), + dtype=dtype, + shape=_shape_tuple(shape), + addrspace=3, + ) + except (RuntimeError, ValueError): + return None + + +def _placeholder_local_array( + dtype: type, shape: int | tuple[int, ...] = 1, alignment: int | None = None +) -> cutlass.Array | None: + """Build a typed local-memory placeholder only when an MLIR context exists.""" + try: + if alignment is None: + return cutlass.Array(dtype, shape, space=cutlass.AddressSpace.rmem) + return cutlass.Array( + dtype, shape, space=cutlass.AddressSpace.rmem, alignment=alignment + ) + except (RuntimeError, ValueError): + return None + + +def _placeholder_tmem_ptr() -> cutlass.Array | None: + """Build a typed tensor-memory pointer placeholder only when an MLIR context exists.""" + try: + return prims.make_tmem_ptr(cutlass.Int32(0), cutlass.Int8) + except RuntimeError: + return None diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py new file mode 100644 index 000000000000..01bd9e49152d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Common arithmetic for standalone split-KV GMEM reducers. + +Every producer using this interface publishes the same mathematical state: + +* one FP32 log2-LSE value per split and output row; and +* one normalized 16-bit partial-O vector for that split and row. + +Reducer policies still own row addressing, active-split discovery, cluster +topology, PDL ordering, and final-output stores. Those details differ between +FMHA, MLA 1CTA, and MLA 2CTA and are deliberately kept out of this module. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float16, Float32, Int32 + + +@cute.jit +def finalize_log2_sum_exp(frame: Float32, exp_sum: Float32) -> Float32: + """Return ``frame + log2(exp_sum)`` or the neutral ``-inf`` LSE.""" + + has_mass = exp_sum == exp_sum and exp_sum != Float32(0.0) + return ( + frame + cute.math.log2(exp_sum, fastmath=True) + if has_mass + else Float32(-Float32.inf) + ) + + +@cute.jit +def normalized_lse_weight(partial_lse: Float32, global_lse: Float32) -> Float32: + """Return the weight of one normalized partial in the global output.""" + + has_mass = global_lse != Float32(-Float32.inf) + return ( + Float32( + cute.math.exp2( + partial_lse - global_lse, + fastmath=True, + ) + ) + if has_mass + else Float32(0.0) + ) + + +@cute.jit +def merge_log2_lse( + lhs_lse: Float32, + rhs_lse: Float32, +) -> tuple[Float32, Float32, Float32]: + """Merge two normalized states and return LSE plus output weights.""" + + neg_inf = Float32(-Float32.inf) + frame = cute.math.max(lhs_lse, rhs_lse, ftz=True) + has_mass = frame != neg_inf + lhs_exp = ( + Float32(cute.math.exp2(lhs_lse - frame, fastmath=True)) + if has_mass + else Float32(0.0) + ) + rhs_exp = ( + Float32(cute.math.exp2(rhs_lse - frame, fastmath=True)) + if has_mass + else Float32(0.0) + ) + exp_sum = lhs_exp + rhs_exp + merged_lse = finalize_log2_sum_exp(frame, exp_sum) + inv_sum = Float32(1.0) / exp_sum if has_mass else Float32(0.0) + return merged_lse, lhs_exp * inv_sum, rhs_exp * inv_sum + + +@cute.jit +def unpack_normalized_vec8( + regs_i32: cutlass.Array, + use_bf16_partial: cutlass.Constexpr[bool], +) -> cutlass.Array: + """Decode one packed 16-byte normalized partial-O vector to FP32.""" + + regs_vec = cutlass.Vector.from_elements( + (regs_i32[0], regs_i32[1], regs_i32[2], regs_i32[3]), + Int32, + ) + if cutlass.const_expr(use_bf16_partial): + return regs_vec.bitcast(BFloat16).to(Float32) + return regs_vec.bitcast(Float16).to(Float32) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/stage.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/stage.py new file mode 100644 index 000000000000..faea48d7c22d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/stage.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Kernel-local schedule-section tag for FMHA decode/context work bodies. + +Some FMHA work methods need to know which schedule section (head/loop/tail) a +call belongs to. Rather than depend on the task-scheduling framework's +``ScheduleStageType``, the FMHA schedules pass this small kernel-local enum +explicitly as a compile-time constant (the ``section`` work argument), so bodies +branch on it with ``cutlass.const_expr(section == FmhaStage.Loop)``. +""" + +import enum + + +class FmhaStage(enum.Enum): + """Schedule section of a single FMHA decode/context work call.""" + + Head = 0 + Loop = 1 + Tail = 2 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tcgen05_compat.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tcgen05_compat.py new file mode 100644 index 000000000000..f852b1daec46 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tcgen05_compat.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Narrow compatibility helpers for tcgen05 primitives used by PrimTS kernels.""" + +import cutlass +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + + +@dsl_user_op +def tcgen05_mma_ws( + mma_kind, + d, + a, + b, + idesc, + enable_input_d, + *, + loc=None, + ip=None, +) -> None: + """Issue WS MMA across the CUTLASS DSL 4.7 keyword mismatch. + + CUTLASS DSL 4.7 exposes ``col_b_zero_mask`` in the public wrapper but + forwards it under the rejected name ``zero_col_mask``. Prefer the public + primitive so newer releases stay on their supported API; import the private + binding only after observing that exact compatibility failure. + """ + + try: + prims.tcgen05_mma_ws( + mma_kind, + d, + a, + b, + idesc, + enable_input_d, + col_b_zero_mask=None, + loc=loc, + ip=ip, + ) + return + except TypeError as error: + if "zero_col_mask" not in str(error): + raise + + from cutlass.experimental.primitives import nvvm_wrapper as prims_nvvm + + prims_nvvm._assert_tensor_mem(d, "tcgen05.mma.ws") + prims_nvvm._nvvm.tcgen05_mma_ws( + prims_nvvm._TCGEN05_MMA_KIND_TO_DIALECT[mma_kind], + d, + a, + cutlass.Int64(b), + cutlass.Int32(idesc), + cutlass.Boolean(enable_input_d), + collector_b_buffer=None, + collector_op=None, + col_b_zero_mask=None, + loc=loc, + ip=ip, + ) + + +__all__ = ["tcgen05_mma_ws"] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tensor_map.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tensor_map.py new file mode 100644 index 000000000000..1501d57317e3 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tensor_map.py @@ -0,0 +1,567 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ragged tensor-map helpers for FlashInfer Attention-TS kernels. + +The helpers do not interact with the CUDA driver or runtime directly. They +orchestrate DSL/IR-level work and delegate to the tensor-map primitives in +``cutlass.experimental.cuda`` (such as :func:`create_tensor_map_tiled`). +""" + +from typing import Sequence, Tuple, Type + +import cutlass.cute as cute +from cutlass.cute import depth, leading_dim +from cutlass.cute.testing import assert_ as runtime_assert +from cutlass.cutlass_dsl import ( + Int8, + Int32, + Int64, + Numeric, + dsl_user_op, + is_dynamic_expression, +) + +from cutlass.experimental.cuda.tensor_map import ( + TensorMap, + TensorMapDataFormat, + TensorMapDataType, + TensorMapFloatOOBFill, + TensorMapInterleave, + TensorMapL2Promotion, + TensorMapSwizzle, + create_tensor_map_tiled, + create_tensor_map_tiled_from_view, # noqa: F401 - public re-export + get_dsl_type_to_tensormap_type, + get_tensormap_type_to_dsl_type, +) + + +# Module-private constants used by the ragged-TMA helpers below. The +# descriptor splits the ragged axis into three TMA dimensions of size +# (box, TmaDimMax, TmaDimMax) whose contributions to the global address +# sum to a multiple of 2^64 (so they wrap to zero) for any element type +# whose width is at least 4 bits: +# LargeN * XLargeN * elem_bytes = 2^30 * 2^35 * elem_bytes +# = 2^65 * elem_bytes +# ≡ 0 (mod 2^64) for elem_bytes ≥ 1/2. +# Match the large-dimension sentinels used by the reference ragged tensor-map +# implementation. Kept private — callers never need to reference these +# directly; they are an implementation detail of +# the splice the helpers build. +_RAGGED_LARGE_N = 1 << 30 +_RAGGED_XLARGE_N = 1 << 35 +_TMA_DIM_MAX = 1 << 31 + +_LEGACY_TMA_TYPE_TO_FORMAT: dict[TensorMapDataType, TensorMapDataFormat] = { + TensorMapDataType.uint8: TensorMapDataFormat.BYTE, + TensorMapDataType.uint16: TensorMapDataFormat.DEFAULT, + TensorMapDataType.uint32: TensorMapDataFormat.DEFAULT, + TensorMapDataType.int32: TensorMapDataFormat.DEFAULT, + TensorMapDataType.uint64: TensorMapDataFormat.DEFAULT, + TensorMapDataType.int64: TensorMapDataFormat.DEFAULT, + TensorMapDataType.float16: TensorMapDataFormat.DEFAULT, + TensorMapDataType.float32: TensorMapDataFormat.DEFAULT, + TensorMapDataType.float64: TensorMapDataFormat.DEFAULT, + TensorMapDataType.bfloat16: TensorMapDataFormat.DEFAULT, + TensorMapDataType.float32_ftz: TensorMapDataFormat.F32_FTZ, + TensorMapDataType.tfloat32: TensorMapDataFormat.DEFAULT, + TensorMapDataType.tfloat32_ftz: TensorMapDataFormat.TF32_FTZ, + TensorMapDataType.f416u4_align8b: TensorMapDataFormat.B4X16, + TensorMapDataType.f416u4_align16b: TensorMapDataFormat.B4X16_P64, + TensorMapDataType.f416u6_align16b: TensorMapDataFormat.B6X16_P32, +} + + +@dsl_user_op +def create_tensor_map_ragged( + global_address: Int64 | int, + tma_format: TensorMapDataType | Type[Numeric], + global_dims: Sequence[Int32 | int], + global_strides: Sequence[Int64 | int], + box_dims: Sequence[Int8 | int], + *, + ragged_dim_idx: int, + interleave: TensorMapInterleave | None = None, + swizzle: TensorMapSwizzle | None = None, + l2_promotion: TensorMapL2Promotion | None = None, + oob_fill: TensorMapFloatOOBFill | None = None, + loc=None, + ip=None, +) -> TensorMap: + """Low-level: build a ragged TMA descriptor from explicit parameters. + + Same role for ragged TMA that + :func:`cutlass.experimental.cuda.create_tensor_map_tiled` plays + for dense TMA: takes raw column-major-ordered dims / strides / box, + inserts the synthetic out-of-bounds dimension splice + around ``ragged_dim_idx``, and delegates to ``create_tensor_map_tiled``. + Prefer :func:`create_tensor_map_ragged_from_tensor` when a + ``cute.Tensor`` is available. + + The splice replaces the ragged axis at TMA position + ``ragged_dim_idx`` with three TMA dimensions of size + ``(box_dims[ragged_dim_idx], TmaDimMax, TmaDimMax)``. The + synthetic strides ``(S, XLargeN - S, S)`` (element units) are + chosen so that ``LargeN * XLargeN * elem_bytes ≡ 0 (mod 2^64)`` + for any element type whose width is at least 4 bits. At kernel + time, the matching :func:`transform_ragged_coords` helper folds a + runtime ``ragged_extent`` into the coords so out-of-range elements + fall past the TMA box boundary and are filled per ``oob_fill``. + + :param global_address: Device pointer to the first element of the + global tensor. + :type global_address: Int64 or int + :param tma_format: Element data type (e.g. + ``TensorMapDataType.float16`` or ``cutlass.Float16``). + :type tma_format: TensorMapDataType or CUTLASS dtype + :param global_dims: Shape of the **unspliced** global tensor in + TMA (column-major) order, length ``R ∈ {2, 3}``. + :type global_dims: Sequence[Int32 or int] + :param global_strides: Inter-dimension strides of the unspliced + tensor in **16-byte units**, length ``R - 1`` (same convention + as ``create_tensor_map_tiled``: innermost stride is implicit). + :type global_strides: Sequence[Int64 or int] + :param box_dims: Tile (box) dimensions in column-major order, + length ``R``. ``box_dims[ragged_dim_idx]`` doubles as the + wraparound period of the synthetic splice. + :type box_dims: Sequence[Int8 or int] + :param ragged_dim_idx: TMA-order index of the ragged axis. Must + satisfy ``1 ≤ ragged_dim_idx ≤ R - 1`` — the innermost + (stride-1) axis is rejected because the wraparound stride + would truncate to 0 under the 16-byte-unit convention for + sub-128-bit element types. + :type ragged_dim_idx: int + :param interleave: Interleave mode, defaults to ``none``. + :param swizzle: Swizzle mode, defaults to ``none``. + :param l2_promotion: L2 promotion hint, defaults to ``none``. + :param oob_fill: OOB fill mode. Default is format-aware: + floating TMA formats (``float16``/``float32``/``float64``/ + ``bfloat16``/``float32_ftz``/``tfloat32``/``tfloat32_ftz``) + get ``nan_request_zero_fma``; everything else (integer types + and packed sub-byte float types) gets ``none``, because + ``nan_request_zero_fma`` is only legal for full-precision + IEEE float TMA formats and the hardware would otherwise fault. + :raises ValueError: For rank not in ``{2, 3}``, ``ragged_dim_idx`` + out of range or pointing at the innermost axis, or + ``box_dims`` / ``global_strides`` of wrong length. + :return: A :class:`TensorMap` of TMA rank ``R + 2``. + :rtype: TensorMap + """ + rank = len(global_dims) + if not 2 <= rank <= 3: + raise ValueError( + f"create_tensor_map_ragged supports rank 2 or 3 " + f"(resulting TMA rank ≤ 5); got input rank {rank}. Rank 1 " + f"is rejected because the unit element-stride of the ragged " + f"axis truncates to 0 in the 16-byte-unit stride convention " + f"for sub-128-bit element types." + ) + if len(box_dims) != rank: + raise ValueError( + f"box_dims length {len(box_dims)} does not match global_dims rank {rank}" + ) + if len(global_strides) != rank - 1: + raise ValueError( + f"global_strides length {len(global_strides)} != rank-1 = {rank - 1}" + ) + if not 1 <= ragged_dim_idx <= rank - 1: + raise ValueError( + "ragged_dim_idx must satisfy 1 ≤ idx ≤ rank-1 " + f"(innermost axis at TMA position 0 is not allowed); " + f"got ragged_dim_idx={ragged_dim_idx} for rank-{rank} input" + ) + + # Coerce CUTLASS dtypes (e.g. ``cutlass.Float16``) to their TensorMapDataType + # representative so the format-aware OOB-fill default below works + # uniformly with whichever spelling the caller passed. + fmt_resolved = ( + tma_format + if isinstance(tma_format, TensorMapDataType) + else get_dsl_type_to_tensormap_type(tma_format) + ) + elem_bits = fmt_resolved.bit_width + + # Splice global_dims: replace the ragged entry with (box, MaxDim, MaxDim). + box_at_ragged = box_dims[ragged_dim_idx] + spliced_global_dims = ( + list(global_dims[:ragged_dim_idx]) + + [box_at_ragged, _TMA_DIM_MAX, _TMA_DIM_MAX] + + list(global_dims[ragged_dim_idx + 1 :]) + ) + + # Splice box_dims: replace the ragged entry with (box, 1, 1). + spliced_box_dims = ( + list(box_dims[:ragged_dim_idx]) + + [box_at_ragged, 1, 1] + + list(box_dims[ragged_dim_idx + 1 :]) + ) + + # Splice strides (16-byte units). `global_strides[i]` is associated + # with `global_dims[i+1]` (the innermost stride is implicit), so the + # ragged-axis stride lives at `global_strides[ragged_dim_idx - 1]`. + # The synthetic wraparound stride ``XLargeN`` is converted from + # element units to 16-byte units here so the formula below stays + # uniform with the pass-through entries. + s_ragged_16b = global_strides[ragged_dim_idx - 1] + xlarge_n_16b = _RAGGED_XLARGE_N * elem_bits // 128 + spliced_global_strides = ( + list(global_strides[: ragged_dim_idx - 1]) + + [s_ragged_16b, xlarge_n_16b - s_ragged_16b, s_ragged_16b] + + list(global_strides[ragged_dim_idx:]) + ) + + if oob_fill is None: + # nan_request_zero_fma is only legal for full-precision IEEE + # float TMA formats; integer types and packed sub-byte float + # formats (uint8 for FP8, f416u4, ...) must use `none` or the + # descriptor faults at issue. + _FLOAT_TMA_FORMATS = ( + TensorMapDataType.float16, + TensorMapDataType.float32, + TensorMapDataType.float64, + TensorMapDataType.bfloat16, + TensorMapDataType.float32_ftz, + TensorMapDataType.tfloat32, + TensorMapDataType.tfloat32_ftz, + ) + oob_fill = ( + TensorMapFloatOOBFill.nan_request_zero_fma + if fmt_resolved in _FLOAT_TMA_FORMATS + else TensorMapFloatOOBFill.none + ) + + descriptor_dtype = ( + tma_format + if isinstance(tma_format, type) + else get_tensormap_type_to_dsl_type(fmt_resolved) + ) + descriptor_format = ( + None + if isinstance(tma_format, type) + else _LEGACY_TMA_TYPE_TO_FORMAT[fmt_resolved] + ) + + return create_tensor_map_tiled( + global_address=global_address, + dtype=descriptor_dtype, + tma_format=descriptor_format, + global_dims=spliced_global_dims, + global_strides=spliced_global_strides, + box_dims=spliced_box_dims, + interleave=interleave, + swizzle=swizzle, + l2_promotion=l2_promotion, + oob_fill=oob_fill, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def create_tensor_map_ragged_from_tensor( + tensor: cute.Tensor, + box_dims: Tuple[Int8 | int, ...], + *, + ragged_dim: int, + stride_order: Tuple[int, ...] | None = None, + interleave: TensorMapInterleave | None = None, + swizzle: TensorMapSwizzle | None = None, + l2_promotion: TensorMapL2Promotion | None = None, + oob_fill: TensorMapFloatOOBFill | None = None, + tma_format: TensorMapDataType | None = None, + loc=None, + ip=None, +) -> TensorMap: + """Build a ragged TMA descriptor from a :class:`cute.Tensor`. + + Thin wrapper around :func:`create_tensor_map_ragged`: derives + ``global_address`` / ``global_dims`` / ``global_strides`` / + ``tma_format`` from ``tensor`` (auto-inferring stride order from + the tensor's strides unless ``stride_order`` is given), then + delegates the splice + descriptor encoding. Pairs with + :func:`transform_ragged_coords` on the kernel side. This is the + synthetic out-of-bounds pattern exposed as a convenience helper. + + :param tensor: Input ``cute.Tensor`` of rank 2 or 3 with a + flattened layout and a stride-1 leading dimension. + :type tensor: cute.Tensor + :param box_dims: Tile dimensions in **tensor mode order** (same + convention as + ``create_tensor_map_tiled_from_view``). + ``box_dims[ragged_dim]`` doubles as the wraparound period of + the synthetic splice. + :type box_dims: tuple[Int8 or int, ...] + :param ragged_dim: Tensor mode index of the ragged axis (the axis + whose runtime length varies per CTA). Must not be the + innermost (stride-1) axis when rank > 1. + :type ragged_dim: int + :param stride_order: Explicit dimension order from innermost to + outermost (same semantics as + ``create_tensor_map_tiled_from_view``). + :type stride_order: tuple[int, ...], optional + :param interleave: Interleave mode, defaults to ``none``. + :param swizzle: Swizzle mode, defaults to ``none``. + :param l2_promotion: L2 promotion hint, defaults to ``none``. + :param oob_fill: OOB fill mode; see + :func:`create_tensor_map_ragged` for the format-aware default. + :param tma_format: Override the element data type, defaults to + auto-detect from ``tensor.element_type``. + :raises ValueError: For unsupported rank (``R ∉ {2, 3}``), + invalid ``ragged_dim``, ragged axis being the innermost, or + ambiguous stride ordering, or static TMA strides that are + smaller than 16 bytes or not 16-byte aligned. + :return: A :class:`TensorMap` of TMA rank ``R + 2``. + :rtype: TensorMap + + Example — fp16 ``[outer_padded, inner]`` row-major, ragged on + the outer axis:: + + desc = create_tensor_map_ragged_from_tensor( + t, + box_dims=(tile_outer, inner), # tensor mode order + ragged_dim=0, # outer (ragged) axis + swizzle=TensorMapSwizzle.s128b, + ) + """ + if depth(tensor) > 1: + raise ValueError( + f"Expected tensor to have flattened layout, got {tensor.layout}" + ) + + rank = len(tensor.shape) + if not 0 <= ragged_dim < rank: + raise ValueError(f"ragged_dim={ragged_dim} out of range for rank-{rank} tensor") + + leading_mode = leading_dim(tensor.shape, tensor.stride) + if leading_mode is None or not isinstance(leading_mode, int): + raise ValueError( + "Expected tensor to have a leading (stride-1) dimension, but got " + f"tensor layout {tensor.layout}" + ) + + tensor_shapes = list(tensor.shape) + tensor_strides = list(tensor.stride) + box_dims_list = list(box_dims) + if len(box_dims_list) != rank: + raise ValueError( + f"box_dims rank {len(box_dims_list)} does not match tensor rank {rank}" + ) + + if stride_order is not None: + order = list(stride_order) + else: + if any(is_dynamic_expression(s) for s in tensor_strides): + raise ValueError( + f"Cannot infer a unique stride order from tensor strides " + f"{tensor_strides} due to dynamic strides. Please provide " + "`stride_order` explicitly." + ) + if len(set(tensor_strides)) < rank: + raise ValueError( + f"Cannot infer a unique stride order from tensor strides " + f"{tensor_strides} due to duplicate strides. Please provide " + "`stride_order` explicitly." + ) + order = sorted(range(rank), key=lambda i: tensor_strides[i]) + + try: + ragged_pos = order.index(ragged_dim) + except ValueError as e: + raise ValueError( + f"ragged_dim={ragged_dim} not present in stride_order {order}" + ) from e + + # Reorder shapes / strides / box into TMA column-major order. + tma_dims = [tensor_shapes[order[j]] for j in range(rank)] + tma_strides_orig = [tensor_strides[order[j]] for j in range(rank)] + tma_box = [box_dims_list[order[j]] for j in range(rank)] + + # Convert to the 16-byte-unit stride convention the low-level helper + # (and create_tensor_map_tiled) expects: length rank - 1, drops the + # innermost stride. + elem_bits = tensor.element_type.width + global_strides = [] + for i in range(rank - 1): + element_stride = tma_strides_orig[i + 1] + stride_bits = element_stride * elem_bits + if is_dynamic_expression(stride_bits): + runtime_assert( + stride_bits >= 128, + f"TMA dimension {i + 1} stride must be at least 16 bytes", + ) + runtime_assert( + stride_bits % 128 == 0, + f"TMA dimension {i + 1} stride must be a multiple of 16 bytes", + ) + elif stride_bits < 128 or stride_bits % 128 != 0: + raise ValueError( + f"stride {element_stride} for TMA dimension {i + 1} must be " + f"a positive multiple of 16 bytes for a {elem_bits}-bit element type" + ) + global_strides.append(stride_bits // 128) + + fmt = ( + tma_format + if tma_format is not None + else get_dsl_type_to_tensormap_type(tensor.element_type) + ) + + return create_tensor_map_ragged( + global_address=tensor.iterator.toint(), + tma_format=fmt, + global_dims=tma_dims, + global_strides=global_strides, + box_dims=tma_box, + ragged_dim_idx=ragged_pos, + interleave=interleave, + swizzle=swizzle, + l2_promotion=l2_promotion, + oob_fill=oob_fill, + loc=loc, + ip=ip, + ) + + +def transform_ragged_coords( + coords: Sequence[Int32 | int], + *, + ragged_dim_idx: int, + ragged_box_size: Int32 | int, + ragged_extent: Int32 | int, +) -> Tuple[Int32, ...]: + """Expand a logical coordinate tuple into the rank-``R + 2`` form + expected by a descriptor built with + :func:`create_tensor_map_ragged_from_tensor`. + + Call from inside ``@cute.kernel``. Given the kernel's original + coordinate tuple (rank ``R = len(coords)`` in TMA order, the same + rank you would pass to ``cp_async_bulk_tensor_*`` for the dense + path) plus a runtime ``ragged_extent`` (number of valid elements + along the ragged axis), this inserts the two synthetic LargeN + coordinates and rewrites the ragged-axis coordinate so + out-of-range elements fall past the TMA box boundary. The math + follows the reference coordinate transformation, with an explicit + ``ragged_extent == 0`` carve-out so a + fully-empty tile is treated as fully OOB rather than fully + in-bounds):: + + ext = clamp(ragged_extent, 0, ragged_box_size) + ext_mod = ext % ragged_box_size + dist = ragged_box_size - ext_mod + d_mod = dist % ragged_box_size # (ext == box ⇒ 0) + if ext == 0: + d_mod = ragged_box_size # all OOB carve-out + bal = -d_mod + out[ragged_dim_idx + 0] = d_mod + out[ragged_dim_idx + 1] = LargeN + out[ragged_dim_idx + 2] = coords[ragged_dim_idx] + LargeN + bal + + All other axes pass through unchanged. + + Behavior across the ``ragged_extent`` value range, with + ``box = ragged_box_size``: + + +-------------------+------------------------------------------+-----------+ + | ``ragged_extent`` | code path | outcome | + +===================+==========================================+===========+ + | ``0`` | empty-tile carve-out → ``d_mod = box`` | all OOB | + +-------------------+------------------------------------------+-----------+ + | ``0 < ext < box`` | normal formula → ``d_mod = box - ext`` | partial | + +-------------------+------------------------------------------+-----------+ + | ``ext == box`` | normal formula → ``d_mod = 0`` | all valid | + +-------------------+------------------------------------------+-----------+ + | ``ext > box`` | upper-clamp → ``ext = box`` → ``d_mod=0``| all valid | + +-------------------+------------------------------------------+-----------+ + | ``ext < 0`` | lower-clamp → ``ext = 0`` → carve-out | all OOB | + +-------------------+------------------------------------------+-----------+ + + :param coords: Original coordinates in TMA order. Length is the + original tensor rank ``R``; must be 2 or 3 (matching the + descriptor helper's rank constraint). + :type coords: Sequence[Int32 or int] + :param ragged_dim_idx: TMA-order index of the ragged axis (i.e., + ``stride_order.index(ragged_dim)`` for the descriptor that + this coord tuple drives). Must be ≥ 1. + :type ragged_dim_idx: int + :param ragged_box_size: Box size along the ragged axis (same + as ``box_dims[ragged_dim]`` passed to the descriptor helper). + Used as the wraparound modulus. + :type ragged_box_size: Int32 or int + :param ragged_extent: Runtime number of valid elements along the + ragged axis starting at this TMA coordinate's ragged-axis origin. + This is a tile-local remaining length, typically computed as + ``ragged_limit - coords[ragged_dim_idx]``. It is not an absolute + limit coordinate or the total logical length of the ragged axis. + Any value is accepted; the helper clamps to ``[0, ragged_box_size]`` + internally so callers can pass a raw per-TMA-tile difference + (e.g., ``mn_limit - tile_origin``) without preconditioning. + ``ragged_extent == 0`` produces all-OOB coordinates (TMA + load fills the whole tile per ``oob_fill``; TMA store is a + full no-op). ``ragged_extent >= ragged_box_size`` produces + all-in-bounds coordinates. + :type ragged_extent: Int32 or int + :return: Expanded coordinate tuple of length ``len(coords) + 2``. + :rtype: tuple[Int32, ...] + """ + rank = len(coords) + if not 2 <= rank <= 3: + raise ValueError(f"transform_ragged_coords supports rank 2 or 3; got {rank}") + if not 0 <= ragged_dim_idx < rank: + raise ValueError( + f"ragged_dim_idx={ragged_dim_idx} out of range for rank-{rank}" + ) + if ragged_dim_idx == 0: + raise ValueError("ragged axis cannot be the innermost (stride-1) TMA dimension") + + box = Int32(ragged_box_size) + ext = Int32(ragged_extent) + # Clamp ext to [0, box]. The OOB-trick formula is only + # well-defined in that range — `ext > box` would treat the + # excess like a partial tile and falsely mark some rows OOB; a + # negative `ext` (e.g. when the caller passes + # `mn_limit - tile_origin` for a tile fully past the limit) + # likewise produces nonsense. Doing the clamp here lets every + # call site pass a raw per-CTA difference and stay readable. + is_neg = Int32(ext < Int32(0)) + ext = ext - ext * is_neg # ext < 0 ⇒ 0 + is_over = Int32(ext > box) + ext = ext + (box - ext) * is_over # ext > box ⇒ box + + ext_mod = ext % box + dist = box - ext_mod + d_mod = dist % box + # When `ragged_extent == 0`, the formula above yields d_mod = 0 + # (same as a fully-utilized tile). Override so a fully-empty + # tile gets d_mod = ragged_box_size, pushing the entire box past + # the TMA bound so all lanes are OOB-handled. + is_empty = Int32(ext == Int32(0)) + d_mod = d_mod + box * is_empty + bal = Int32(0) - d_mod + large_n = Int32(_RAGGED_LARGE_N) + + orig = Int32(coords[ragged_dim_idx]) + expanded = ( + tuple(Int32(c) for c in coords[:ragged_dim_idx]) + + (d_mod, large_n, orig + large_n + bal) + + tuple(Int32(c) for c in coords[ragged_dim_idx + 1 :]) + ) + return expanded + + +__all__ = [ + "create_tensor_map_ragged", + "create_tensor_map_ragged_from_tensor", + "transform_ragged_coords", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py new file mode 100644 index 000000000000..9098a7c58a80 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -0,0 +1,2110 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task-scheduled paged MLA decode with a plan/run lifecycle.""" + +from collections.abc import Callable +from dataclasses import dataclass, field +import functools +from typing import Any, Literal, Optional, cast + +import torch + +from flashinfer.api_logging import flashinfer_api +from flashinfer.trace.templates.attention import ( + prims_ts_decode_mla_one_shot_trace_dispatch, + prims_ts_decode_mla_trace_dispatch, +) + +from ._tensor_aliasing import ( + _validate_out_does_not_overlap_inputs, + _validate_tensor_does_not_overlap_inputs, +) +from .decode import ( + _WorkspaceSection, + _align_up, + _append_workspace_section, + _dtype_key, + _resolve_cuda_device, + _validate_16byte_alignment, + _validate_mask, + _validate_page_size, + _validate_positive_int, + _validate_runtime_device, + _validate_scale, + _validate_workspace_buffer, + _workspace_section_view, +) + + +_COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 2" +_MLA_LATENT_DIM = 512 +_MLA_ROPE_DIM = 64 +_MLA_QUERY_DIM = _MLA_LATENT_DIM + _MLA_ROPE_DIM +_SUPPORTED_INPUT_DTYPES = (torch.bfloat16, torch.float8_e4m3fn) +_SUPPORTED_OUTPUT_DTYPES = (torch.bfloat16,) +_INT32_MAX = 2**31 - 1 +# The largest public 1CTA schedule can pad one K/V split group across 128 +# splits, two 128-token K/V instructions apiece. Reserve that complete span so +# every padded tile boundary remains representable as signed Int32. +_MLA_MAX_KV_COORDINATE_SPAN = 128 * 2 * 128 +_MLA_MAX_KV_LEN = _INT32_MAX - (_MLA_MAX_KV_COORDINATE_SPAN - 1) + + +@dataclass(frozen=True) +class _MLADecodeLaunchSpec: + """Automatic MLA policy and scratch geometry for one plan.""" + + kernel: Any + policy: tuple[tuple[str, object], ...] + kernel_workspace_bytes: int + split_kv: int + + +@dataclass(frozen=True) +class _MLADecodeCompileSpec: + """Batch-independent identity and implementation for one MLA compile.""" + + device_index: int + kernel_signature: tuple[object, ...] + num_heads: int + kv_lora_rank: int + qk_rope_head_dim: int + page_size: int + q_dtype_key: str + output_dtype_key: str + max_seq_len_q: int + packed_query: bool + has_kernel_workspace: bool + split_kv: int + kernel: Any = field(compare=False, hash=False, repr=False) + + +@dataclass(frozen=True) +class _MLAWorkspaceLayout: + """Private MLA scratch layout; only ``total_bytes`` is public.""" + + kernel_workspace: _WorkspaceSection + lse: _WorkspaceSection + total_bytes: int + + +@dataclass(frozen=True) +class _MLAWorkspaceViews: + kernel_workspace: Optional[torch.Tensor] + lse: torch.Tensor + + +@dataclass(frozen=True) +class _MLADecodePlanState: + """Immutable compile and workspace state for one reusable MLA plan.""" + + device: torch.device + batch_size: int + num_heads: int + max_seq_len_q: int + packed_query: bool + kv_lora_rank: int + qk_rope_head_dim: int + page_size: int + q_dtype: torch.dtype + kv_dtype: torch.dtype + output_dtype: torch.dtype + mask_type: Literal["dense", "causal"] + max_kv_len: int + required_page_columns: int + workspace_buffer: torch.Tensor + workspace_layout: _MLAWorkspaceLayout + workspace_views: _MLAWorkspaceViews + compiled: Callable[..., object] + policy: tuple[tuple[str, object], ...] + split_kv: int + + +@dataclass(frozen=True) +class _MLARuntime: + query: torch.Tensor + normalized_cache: torch.Tensor + out: torch.Tensor + num_physical_pages: int + bmm1_scale: float + bmm2_scale: float + + +def _make_mla_workspace_layout( + kernel_workspace_bytes: int, + batch_size: int, + num_heads: int, + max_seq_len_q: int = 1, +) -> _MLAWorkspaceLayout: + kernel_workspace, byte_end = _append_workspace_section( + 0, (kernel_workspace_bytes,), torch.int8 + ) + lse, byte_end = _append_workspace_section( + byte_end, (batch_size, max_seq_len_q, num_heads), torch.float32 + ) + return _MLAWorkspaceLayout( + kernel_workspace=kernel_workspace, + lse=lse, + total_bytes=_align_up(byte_end), + ) + + +def _bind_mla_workspace( + workspace_buffer: torch.Tensor, layout: _MLAWorkspaceLayout +) -> _MLAWorkspaceViews: + kernel_workspace = None + if layout.kernel_workspace.byte_size > 0: + kernel_workspace = _workspace_section_view( + workspace_buffer, layout.kernel_workspace + ) + return _MLAWorkspaceViews( + kernel_workspace=kernel_workspace, + lse=_workspace_section_view(workspace_buffer, layout.lse), + ) + + +def _validate_mla_dims(kv_lora_rank: int, qk_rope_head_dim: int) -> None: + kv_lora_rank = _validate_positive_int(kv_lora_rank, "kv_lora_rank") + qk_rope_head_dim = _validate_positive_int(qk_rope_head_dim, "qk_rope_head_dim") + if (kv_lora_rank, qk_rope_head_dim) != (_MLA_LATENT_DIM, _MLA_ROPE_DIM): + raise NotImplementedError( + "attention-ts MLA decode currently requires " + f"kv_lora_rank={_MLA_LATENT_DIM} and " + f"qk_rope_head_dim={_MLA_ROPE_DIM}; got " + f"{kv_lora_rank} and {qk_rope_head_dim}" + ) + + +def _validate_mla_max_kv_len(value: int, name: str) -> int: + """Reserve the largest padded split-KV coordinate span in signed Int32.""" + value = _validate_positive_int(value, name) + if value > _MLA_MAX_KV_LEN: + raise NotImplementedError( + f"{name} must be <= {_MLA_MAX_KV_LEN} so padded MLA K/V " + "coordinates fit in a signed int32" + ) + return value + + +def _validate_mla_int32_extent(value: int, name: str) -> int: + """Validate a flattened metadata/cache extent used by Int32 coordinates.""" + if value <= 0: + raise ValueError(f"{name} must be positive") + if value > _INT32_MAX: + raise NotImplementedError(f"{name} must fit in a signed int32") + return value + + +def _validate_mla_nonnegative_int32_extent(value: int, name: str) -> int: + """Validate a possibly empty flattened extent used by Int32 coordinates.""" + if value < 0: + raise ValueError(f"{name} must be nonnegative") + if value > _INT32_MAX: + raise NotImplementedError(f"{name} must fit in a signed int32") + return value + + +def _validate_mla_query_head_extent( + *, + batch_size: int, + num_heads: int, + max_seq_len_q: int, + total_q: Optional[int] = None, +) -> None: + """Keep fixed-capacity and packed query-head coordinates in signed Int32.""" + _validate_mla_int32_extent( + batch_size * max_seq_len_q * num_heads, + "batch_size * max_seq_len_q * num_heads", + ) + if total_q is not None: + _validate_mla_nonnegative_int32_extent( + total_q * num_heads, + "total_q * num_heads", + ) + + +def _validate_mla_policy_coordinate_span( + policy: tuple[tuple[str, object], ...], +) -> None: + """Keep the host K/V bound coupled to the automatically selected policy.""" + resolved = dict(policy) + span = ( + int(cast(int, resolved["tile_size_kv"])) + * int(cast(int, resolved["num_insts_kv"])) + * max(int(cast(int, resolved["split_kv"])), 1) + ) + if span > _MLA_MAX_KV_COORDINATE_SPAN: + raise RuntimeError( + "MLA Int32 extent safety assumes a padded K/V coordinate span no " + f"larger than {_MLA_MAX_KV_COORDINATE_SPAN}, got {span}" + ) + + +def _validate_mla_dtype_pair( + q_dtype: torch.dtype, + kv_dtype: torch.dtype, + output_dtype: torch.dtype, +) -> None: + _dtype_key(q_dtype) + _dtype_key(kv_dtype) + _dtype_key(output_dtype) + if q_dtype != kv_dtype: + raise NotImplementedError( + "attention-ts MLA decode requires query and KV cache to use the " + f"same dtype; got {q_dtype} and {kv_dtype}" + ) + if q_dtype not in _SUPPORTED_INPUT_DTYPES: + raise NotImplementedError( + f"attention-ts MLA decode supports BF16 and FP8-E4M3 input; got {q_dtype}" + ) + if output_dtype not in _SUPPORTED_OUTPUT_DTYPES: + raise NotImplementedError( + "attention-ts MLA decode currently supports BF16 output only; " + f"got {output_dtype}" + ) + + +def _validate_int32_cuda_tensor( + tensor: torch.Tensor, + name: str, + *, + ndim: int, + require_contiguous: bool = True, + require_16byte_alignment: bool = True, +) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.ndim != ndim: + raise ValueError(f"{name} must be rank {ndim}, got rank {tensor.ndim}") + if tensor.dtype != torch.int32: + raise TypeError(f"{name} must have dtype torch.int32") + if tensor.device.type != "cuda": + raise ValueError(f"{name} must be a CUDA tensor") + if require_contiguous and not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if require_16byte_alignment: + _validate_16byte_alignment(tensor, name) + + +def _validate_mla_metadata( + block_tables: torch.Tensor, + seq_lens: torch.Tensor, +) -> tuple[torch.device, int, int]: + _validate_int32_cuda_tensor( + block_tables, + "block_tables", + ndim=2, + require_contiguous=False, + require_16byte_alignment=False, + ) + _validate_int32_cuda_tensor( + seq_lens, + "seq_lens", + ndim=1, + require_16byte_alignment=False, + ) + if block_tables.device != seq_lens.device: + raise ValueError("block_tables and seq_lens must be on the same device") + batch_size = int(seq_lens.numel()) + if batch_size <= 0: + raise ValueError("seq_lens must contain at least one request") + if block_tables.shape[0] != batch_size: + raise ValueError( + "block_tables must have one row per request: expected " + f"{batch_size}, got {block_tables.shape[0]}" + ) + max_num_pages = int(block_tables.shape[1]) + if max_num_pages <= 0: + raise ValueError("block_tables must contain at least one page column") + if block_tables.stride(1) != 1: + raise ValueError("block_tables must be contiguous within each row") + if block_tables.stride(0) < max_num_pages: + raise ValueError( + "block_tables rows must not overlap: row stride must be at least " + f"the column count ({max_num_pages}), got {block_tables.stride(0)}" + ) + if block_tables.data_ptr() % 4 != 0: + raise ValueError("block_tables must be 4-byte aligned") + if seq_lens.data_ptr() % 4 != 0: + raise ValueError("seq_lens must be 4-byte aligned") + _validate_mla_int32_extent(batch_size, "batch_size") + _validate_mla_int32_extent(int(block_tables.numel()), "block_tables elements") + return seq_lens.device, batch_size, max_num_pages + + +def _validate_qo_indptr( + qo_indptr: torch.Tensor, + *, + device: torch.device, + batch_size: int, +) -> None: + """Validate the public packed-query metadata without synchronizing.""" + + _validate_int32_cuda_tensor( + qo_indptr, + "qo_indptr", + ndim=1, + require_16byte_alignment=False, + ) + if qo_indptr.device != device: + raise ValueError(f"qo_indptr must be on {device}, got {qo_indptr.device}") + expected_offsets = batch_size + 1 + if qo_indptr.numel() != expected_offsets: + raise ValueError( + "qo_indptr must contain batch_size + 1 cumulative offsets: " + f"expected {expected_offsets}, got {qo_indptr.numel()}" + ) + + +def _derive_max_seq_len_q( + qo_indptr: torch.Tensor, + *, + batch_size: int, +) -> tuple[int, int, tuple[int, ...]]: + """Validate cumulative offsets and derive their maximum nonnegative delta. + + This helper intentionally synchronizes and is therefore used only by + explicit runtime validation and the one-shot convenience path. + """ + + offsets = [int(value) for value in qo_indptr.tolist()] + if len(offsets) != batch_size + 1: + raise ValueError("qo_indptr must contain batch_size + 1 offsets") + if offsets[0] != 0: + raise ValueError("qo_indptr must start at 0") + q_lengths = tuple( + end - start for start, end in zip(offsets[:-1], offsets[1:], strict=True) + ) + if any(length < 0 for length in q_lengths): + raise ValueError("qo_indptr must be nondecreasing") + return max(q_lengths), offsets[-1], q_lengths + + +def _validate_mla_run_metadata( + state: _MLADecodePlanState, + runtime: _MLARuntime, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + qo_indptr: Optional[torch.Tensor], +) -> None: + """Validate per-run request metadata against one static MLA plan. + + Value checks synchronize CUDA metadata with the host. ``run(validate=False)`` + is the synchronization-free path for compilation and graph capture. + """ + + device, batch_size, max_num_pages = _validate_mla_metadata(block_tables, seq_lens) + if device != state.device: + raise ValueError( + f"MLA metadata must be on the planned device {state.device}, got {device}" + ) + if batch_size != state.batch_size: + raise ValueError( + "MLA metadata batch size must match the plan " + f"({state.batch_size}), got {batch_size}" + ) + if max_num_pages < state.required_page_columns: + raise ValueError( + "block_tables must have at least ceil(max_kv_len / page_size) " + f"columns ({state.required_page_columns}), got {max_num_pages}" + ) + + if state.packed_query: + if qo_indptr is None: + raise ValueError("packed-query MLA run requires qo_indptr") + _validate_qo_indptr( + qo_indptr, + device=state.device, + batch_size=state.batch_size, + ) + runtime_max_seq_len_q, total_q, q_lengths = _derive_max_seq_len_q( + qo_indptr, + batch_size=state.batch_size, + ) + if total_q != int(runtime.query.shape[0]): + raise ValueError( + "qo_indptr must end at the packed query row count " + f"({runtime.query.shape[0]}), got {total_q}" + ) + if runtime_max_seq_len_q > state.max_seq_len_q: + raise ValueError( + "qo_indptr contains a per-request Q length larger than " + f"max_seq_len_q ({state.max_seq_len_q}): got " + f"{runtime_max_seq_len_q}" + ) + else: + if qo_indptr is not None: + raise ValueError("fixed-query MLA plan does not accept qo_indptr") + q_lengths = (state.max_seq_len_q,) * state.batch_size + + seq_lens_host = tuple(int(value) for value in seq_lens.tolist()) + if any(seq_len <= 0 for seq_len in seq_lens_host): + raise ValueError("every runtime request must contain at least one KV token") + runtime_max_kv_len = max(seq_lens_host) + if runtime_max_kv_len > state.max_kv_len: + raise ValueError( + "runtime KV metadata contains a request longer than " + f"max_kv_len ({state.max_kv_len}): got {runtime_max_kv_len}" + ) + block_table_rows = block_tables.tolist() + for request_idx, (row, seq_len) in enumerate( + zip(block_table_rows, seq_lens_host, strict=True) + ): + required_pages = _ceil_div(seq_len, state.page_size) + if any( + int(page_id) < 0 or int(page_id) >= runtime.num_physical_pages + for page_id in row[:required_pages] + ): + raise ValueError( + "block_tables values for active pages must index the physical " + f"K/V cache in [0, {runtime.num_physical_pages}); request " + f"{request_idx} contains an invalid page ID" + ) + if state.mask_type == "causal": + for request_idx, (q_len, kv_len) in enumerate( + zip(q_lengths, seq_lens_host, strict=True) + ): + if q_len > kv_len: + raise ValueError( + "causal MLA decode requires every per-request Q length " + "to be no greater than its K/V length; request " + f"{request_idx} has Q={q_len} and K/V={kv_len}" + ) + + +def _resolve_max_seq_len_q_alias( + *, + seq_len_q: Optional[int], + max_seq_len_q: Optional[int], + default: Optional[int], +) -> Optional[int]: + """Resolve the legacy fixed-Q name and the explicit static-bound name.""" + + legacy_bound = ( + _validate_positive_int(seq_len_q, "seq_len_q") + if seq_len_q is not None + else None + ) + explicit_bound = ( + _validate_positive_int(max_seq_len_q, "max_seq_len_q") + if max_seq_len_q is not None + else None + ) + if ( + legacy_bound is not None + and explicit_bound is not None + and legacy_bound != explicit_bound + ): + raise ValueError( + "seq_len_q and max_seq_len_q must agree when both are provided: " + f"got {legacy_bound} and {explicit_bound}" + ) + if explicit_bound is not None: + return explicit_bound + if legacy_bound is not None: + return legacy_bound + return default + + +def _validate_query( + query: torch.Tensor, + *, + packed_query: bool = False, + device: Optional[torch.device] = None, + batch_size: Optional[int] = None, + num_heads: Optional[int] = None, + max_seq_len_q: Optional[int] = None, + q_dtype: Optional[torch.dtype] = None, +) -> None: + if not isinstance(query, torch.Tensor): + raise TypeError("query must be a torch.Tensor") + expected_rank = 3 if packed_query else 4 + if query.ndim != expected_rank: + expected_shape = "[total_q, H, 576]" if packed_query else "[B, SQ, H, 576]" + raise ValueError(f"query must have shape {expected_shape}") + if packed_query: + if int(query.shape[1]) <= 0: + raise ValueError("query head extent must be positive") + elif any(int(extent) <= 0 for extent in query.shape[:-1]): + raise ValueError("query row and head extents must be positive") + if query.shape[-1] != _MLA_QUERY_DIM: + raise ValueError( + f"query last dimension must be {_MLA_QUERY_DIM}, got {query.shape[-1]}" + ) + if query.dtype not in _SUPPORTED_INPUT_DTYPES: + raise NotImplementedError( + f"unsupported attention-ts MLA query dtype {query.dtype}" + ) + if query.device.type != "cuda": + raise ValueError("query must be a CUDA tensor") + if device is not None and query.device != device: + raise ValueError( + f"query must be on the planned device {device}, got {query.device}" + ) + if not packed_query and batch_size is not None and query.shape[0] != batch_size: + raise ValueError( + f"query batch size must match the plan ({batch_size}), got {query.shape[0]}" + ) + head_axis = 1 if packed_query else 2 + if num_heads is not None and query.shape[head_axis] != num_heads: + raise ValueError( + "query head count must match the plan " + f"({num_heads}), got {query.shape[head_axis]}" + ) + if max_seq_len_q is not None: + if packed_query: + if batch_size is None: + raise ValueError("batch_size is required to validate packed query") + total_q = int(query.shape[0]) + if total_q > batch_size * max_seq_len_q: + raise ValueError( + "packed query total rows must be within " + f"[0, {batch_size * max_seq_len_q}], got {total_q}" + ) + elif query.shape[1] != max_seq_len_q: + raise ValueError( + "fixed query length must equal the planned max_seq_len_q " + f"({max_seq_len_q}), got {query.shape[1]}" + ) + if batch_size is not None and num_heads is not None: + _validate_mla_query_head_extent( + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + total_q=int(query.shape[0]) if packed_query else None, + ) + if q_dtype is not None and query.dtype != q_dtype: + raise ValueError( + f"query dtype must match the plan ({q_dtype}), got {query.dtype}" + ) + if not query.is_contiguous(): + layout = "[total_q, H, 576]" if packed_query else "[B, SQ, H, 576]" + raise ValueError(f"query must be compact in {layout} layout") + _validate_16byte_alignment(query, "query") + + +def _normalize_mla_kv_cache( + kv_cache: torch.Tensor, + *, + expected_device: torch.device, +) -> tuple[torch.Tensor, int, int]: + if not isinstance(kv_cache, torch.Tensor): + raise TypeError("kv_cache must be a torch.Tensor") + if kv_cache.ndim == 4: + if kv_cache.shape[1] != 1: + raise ValueError( + "rank-4 kv_cache must have shape [num_pages, 1, page_size, 576]" + ) + if not kv_cache.is_contiguous(): + raise ValueError("rank-4 kv_cache must be compact") + normalized = kv_cache[:, 0] + elif kv_cache.ndim == 3: + if not kv_cache.is_contiguous(): + raise ValueError("rank-3 kv_cache must be compact") + normalized = kv_cache + else: + raise ValueError( + "kv_cache must have shape [num_pages, page_size, 576] or " + "[num_pages, 1, page_size, 576]" + ) + if normalized.device != expected_device: + raise ValueError( + f"kv_cache must be on the planned device {expected_device}, " + f"got {normalized.device}" + ) + if normalized.shape[0] <= 0 or normalized.shape[1] <= 0: + raise ValueError("kv_cache page count and page size must be positive") + _validate_mla_int32_extent(int(normalized.shape[0]), "kv_cache physical pages") + if normalized.shape[2] != _MLA_QUERY_DIM: + raise ValueError( + f"kv_cache last dimension must be {_MLA_QUERY_DIM}, " + f"got {normalized.shape[2]}" + ) + _validate_16byte_alignment(normalized, "kv_cache") + return normalized, int(normalized.shape[0]), int(normalized.shape[1]) + + +def _validate_out( + out: torch.Tensor, + *, + device: torch.device, + batch_size: int, + num_heads: int, + max_seq_len_q: int, + packed_query: bool, + total_q: Optional[int] = None, + output_dtype: torch.dtype, +) -> None: + if not isinstance(out, torch.Tensor): + raise TypeError("out must be a torch.Tensor") + if packed_query: + if total_q is None: + raise ValueError("total_q is required to validate packed output") + expected_shape: tuple[int, ...] + expected_shape = (total_q, num_heads, _MLA_LATENT_DIM) + else: + expected_shape = (batch_size, max_seq_len_q, num_heads, _MLA_LATENT_DIM) + if out.shape != expected_shape: + raise ValueError( + f"out must have shape {expected_shape}, got {tuple(out.shape)}" + ) + if out.dtype != output_dtype: + raise ValueError(f"out must have dtype {output_dtype}, got {out.dtype}") + if out.device != device: + raise ValueError(f"out must be on {device}, got {out.device}") + if not out.is_contiguous(): + layout = "[total_q, H, 512]" if packed_query else "[B, SQ, H, 512]" + raise ValueError(f"out must be compact in {layout} layout") + _validate_16byte_alignment(out, "out") + + +def _kernel_dtype_name(dtype_key: str) -> str: + names = { + "bfloat16": "bf16", + "float8_e4m3fn": "e4m3", + } + try: + return names[dtype_key] + except KeyError as error: + raise NotImplementedError( + f"unsupported attention-ts MLA dtype key {dtype_key!r}" + ) from error + + +def _ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def _separate_reducer_provenance( + kernel, + *, + split_kv: int, + use_cluster_reduction: bool, +) -> tuple[str, Optional[int]]: + """Describe the derived standalone reducer without exposing a knob.""" + + if split_kv <= 1 or use_cluster_reduction: + return "none", None + if bool(getattr(kernel, "use_parallel_reduction", False)): + topology = getattr(kernel, "parallel_reduction_topology", None) + if topology is None: + raise RuntimeError("parallel MLA reducer is missing its topology") + return "parallel", int(topology.cluster_size) + return "reference", 1 + + +@functools.cache +def _resolve_mla_decode_launch_spec( + device_index: int, + batch_size: int, + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + page_size: int, + max_kv_len: int, + q_dtype_key: str, + kv_dtype_key: str, + output_dtype_key: str, + mask_type: str, + seq_len_q: int = 1, +): + """Resolve and cache MLA policy/workspace without compiling.""" + + max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") + + import cutlass + import cutlass.utils as cutlass_utils + from cuda.bindings import driver as cuda_drv + + from .kernels.mla_decode.kernel_policy import ( + resolve_mla_kernel_policy, + select_mla_ts_kernel, + ) + from .kernels.mla_decode.helpers.query import FlatQueryTileLayout + from .kernels.mla_decode.throughput_2cta.config import ( + compute_split_kv, + compute_workspace_size as compute_2cta_workspace_size, + ) + from .kernels.mla_decode.throughput_2cta.kernel import MlaDecodeTs + from .kernels.mla_decode.throughput_latency_1cta.config import ( + compute_workspace_size as compute_1cta_workspace_size, + q_tile_work_count, + resolve_auto_flat_query_launch_shape, + resolve_runtime_cluster_reduction_mode, + select_auto_split_kv, + ) + from .kernels.mla_decode.throughput_latency_1cta.kernel import ( + ThroughputLatencyMlaDecodeTs, + ) + + if q_dtype_key != kv_dtype_key: + raise ValueError("the cached TS MLA compiler requires one QKV dtype") + seq_len_q = _validate_positive_int(seq_len_q, "seq_len_q") + _validate_mla_query_head_extent( + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=seq_len_q, + ) + _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) + qkv_dtype_name = _kernel_dtype_name(q_dtype_key) + output_dtype_name = _kernel_dtype_name(output_dtype_key) + + with torch.cuda.device(device_index): + plan_stream = cuda_drv.CUstream( + torch.cuda.current_stream(device_index).cuda_stream + ) + hardware_info = cutlass_utils.HardwareInfo(device_index) + max_active_one_cta_clusters = hardware_info.get_max_active_clusters( + 1, plan_stream + ) + max_active_two_cta_clusters = hardware_info.get_max_active_clusters( + 2, plan_stream + ) + one_cta_launch_shape = resolve_auto_flat_query_launch_shape( + num_heads_q=num_heads, + seq_len_q=seq_len_q, + ) + one_cta_base_work = q_tile_work_count( + batch_size, + one_cta_launch_shape.num_heads_q, + one_cta_launch_shape.seq_len_q, + one_cta_launch_shape.tile_size_q, + ) + one_cta_split_kv = select_auto_split_kv( + seq_len_kv=max_kv_len, + tile_size_q=one_cta_launch_shape.tile_size_q, + base_work=one_cta_base_work, + target_work=max_active_one_cta_clusters, + ) + two_cta_launch_shape = FlatQueryTileLayout.for_tile(num_heads, seq_len_q, 128) + two_cta_split_kv = compute_split_kv( + batch_size=batch_size, + num_q_tiles=two_cta_launch_shape.num_tiles, + seq_len_kv=max_kv_len, + mma_qk_tiler_mn=(128, 128), + max_active_blocks=max_active_two_cta_clusters * 2, + ) + requested_policy, policy_source = resolve_mla_kernel_policy( + None, + num_heads, + seq_len_q, + one_cta_split_kv=one_cta_split_kv, + two_cta_split_kv=two_cta_split_kv, + ) + use_throughput_latency = requested_policy == "throughput_latency_1cta" + + kernel: Any + if use_throughput_latency: + max_active_clusters = max_active_one_cta_clusters + launch_shape = one_cta_launch_shape + decision = select_mla_ts_kernel( + requested_policy=requested_policy, + batch_size=batch_size, + num_heads=launch_shape.num_heads_q, + seq_len_q=launch_shape.seq_len_q, + seq_len_k=max_kv_len, + latent_dim=kv_lora_rank, + rope_dim=qk_rope_head_dim, + page_size=page_size, + dtype=qkv_dtype_name, + out_dtype=output_dtype_name, + throughput_latency_profile=None, + throughput_latency_tile_size_q=launch_shape.tile_size_q, + max_active_clusters=max_active_clusters, + throughput_latency_split_kv=None, + throughput_latency_persistent=None, + ) + if not decision.implementation_ready or decision.config is None: + raise NotImplementedError(decision.reason) + reduction_mode = resolve_runtime_cluster_reduction_mode( + decision.config, + reduction_mode=None, + hardware_info=hardware_info, + stream=plan_stream, + ) + kernel = ThroughputLatencyMlaDecodeTs( + batch_size=batch_size, + num_heads=launch_shape.num_heads_q, + seq_len_q=launch_shape.seq_len_q, + seq_len_k=max_kv_len, + latent_dim=kv_lora_rank, + rope_dim=qk_rope_head_dim, + page_size=page_size, + max_active_clusters=max_active_clusters, + acc_dtype=cutlass.Float32, + lse_dtype=cutlass.Float32, + qkv_dtype=qkv_dtype_name, + out_dtype=output_dtype_name, + profile=decision.profile_name, + reduction_mode=reduction_mode, + logical_num_heads=num_heads, + logical_seq_len_q=seq_len_q, + tile_size_q=launch_shape.tile_size_q, + explicit_split_kv=None, + explicit_persistent=None, + mask_type=mask_type, + ) + final_cfg = kernel._make_config() + split_kv = int(final_cfg.num_ctas_per_seq_kv) + workspace_size = compute_1cta_workspace_size( + cfg=final_cfg, + partial_o_dtype=cutlass.BFloat16, + lse_dtype=cutlass.Float32, + ) + separate_reducer_impl, reducer_cluster_size = _separate_reducer_provenance( + kernel, + split_kv=split_kv, + use_cluster_reduction=bool(final_cfg.use_cluster_reduction), + ) + policy = ( + ("kernel", decision.selected_kernel), + ("source", policy_source), + ("profile", decision.profile_name), + ("tile_size_q", int(final_cfg.tile_size_q)), + ("tile_size_kv", int(final_cfg.tile_size_kv)), + ("num_insts_kv", int(final_cfg.num_insts_kv)), + ("split_kv", split_kv), + ("num_ctas_per_head_dim", int(final_cfg.num_ctas_per_head_dim)), + ("head_dim_per_cta_v", int(final_cfg.head_dim_per_cta_v)), + ("use_cluster_reduction", bool(final_cfg.use_cluster_reduction)), + ( + "use_persistent_scheduler", + bool(final_cfg.use_persistent_scheduler), + ), + ( + "use_clc_dynamic_persistent_scheduler", + bool(final_cfg.use_clc_dynamic_persistent_scheduler), + ), + ("separate_reducer_impl", separate_reducer_impl), + ("reducer_cluster_size", reducer_cluster_size), + ) + else: + max_active_clusters = max_active_two_cta_clusters + launch_shape = two_cta_launch_shape + decision = select_mla_ts_kernel( + requested_policy=requested_policy, + batch_size=batch_size, + num_heads=num_heads, + seq_len_q=seq_len_q, + seq_len_k=max_kv_len, + latent_dim=kv_lora_rank, + rope_dim=qk_rope_head_dim, + page_size=page_size, + dtype=qkv_dtype_name, + out_dtype=output_dtype_name, + throughput_latency_profile=None, + throughput_latency_tile_size_q=None, + max_active_clusters=max_active_clusters, + throughput_latency_split_kv=None, + throughput_latency_persistent=None, + ) + if not decision.implementation_ready: + raise NotImplementedError(decision.reason) + split_kv = two_cta_split_kv + work_clusters = batch_size * launch_shape.num_tiles * max(split_kv, 1) + # Dynamic cluster stealing only helps once logical work exceeds a + # resident wave. Within one wave every cluster already launches, + # so the CLC producer/response pipeline is pure overhead. + is_persistent = work_clusters > max_active_clusters + kernel = MlaDecodeTs( + acc_dtype=cutlass.Float32, + lse_dtype=cutlass.Float32, + mma_qk_tiler_mn=(128, 128), + mma_pv_tiler_mn=(128, 256), + max_active_clusters=max_active_clusters, + page_size=page_size, + is_persistent=is_persistent, + is_var_seq=False, + is_var_split_kv=False, + static_split_kv=split_kv, + static_seq_len_k=None, + qkv_dtype=qkv_dtype_name, + out_dtype=output_dtype_name, + rope_dim=qk_rope_head_dim, + num_heads=num_heads, + seq_len_q=seq_len_q, + batch_size=batch_size, + mask_type=mask_type, + ) + workspace_size = compute_2cta_workspace_size( + tile_size_q=int(launch_shape.tile_size_q), + num_q_tiles=int(launch_shape.num_tiles), + latent_dim=kv_lora_rank, + batch_size=batch_size, + split_kv=split_kv, + partial_o_dtype=cutlass.BFloat16, + lse_dtype=cutlass.Float32, + ) + separate_reducer_impl, reducer_cluster_size = _separate_reducer_provenance( + kernel, + split_kv=split_kv, + use_cluster_reduction=False, + ) + policy = ( + ("kernel", decision.selected_kernel), + ("source", policy_source), + ("profile", None), + ("tile_size_q", 128), + ("tile_size_kv", 128), + ("num_insts_kv", 1), + ("split_kv", int(split_kv)), + ("num_ctas_per_head_dim", 2), + ("head_dim_per_cta_v", 256), + ("use_cluster_reduction", False), + ("use_persistent_scheduler", bool(is_persistent)), + ( + "use_clc_dynamic_persistent_scheduler", + bool(is_persistent and qkv_dtype_name == "bf16"), + ), + ("separate_reducer_impl", separate_reducer_impl), + ("reducer_cluster_size", reducer_cluster_size), + ) + _validate_mla_policy_coordinate_span(policy) + + return _MLADecodeLaunchSpec( + kernel=kernel, + policy=policy, + kernel_workspace_bytes=int(workspace_size), + split_kv=int(split_kv), + ) + + +def _mla_kernel_compile_signature(kernel: Any) -> tuple[object, ...]: + """Return all static kernel state except the batch extent.""" + + make_signature = getattr(kernel, "compile_signature", None) + if not callable(make_signature): + raise TypeError("MLA kernels must define compile_signature()") + signature = ( + type(kernel).__module__, + type(kernel).__qualname__, + make_signature(), + ) + try: + hash(signature) + except TypeError as error: + raise TypeError("MLA kernel compile state must be hashable") from error + return signature + + +def _make_mla_decode_compile_spec( + launch_spec: _MLADecodeLaunchSpec, + *, + device_index: int, + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + page_size: int, + q_dtype_key: str, + output_dtype_key: str, + max_seq_len_q: int, + packed_query: bool, +) -> _MLADecodeCompileSpec: + """Keep policy resolution plan-specific and JIT identity batch-free.""" + + return _MLADecodeCompileSpec( + device_index=device_index, + kernel_signature=_mla_kernel_compile_signature(launch_spec.kernel), + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + page_size=page_size, + q_dtype_key=q_dtype_key, + output_dtype_key=output_dtype_key, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + has_kernel_workspace=launch_spec.kernel_workspace_bytes > 0, + split_kv=launch_spec.split_kv, + kernel=launch_spec.kernel, + ) + + +@functools.cache +def _get_compiled_mla_decode( + compile_spec: _MLADecodeCompileSpec, +): + """Compile and cache one batch-dynamic TS MLA topology.""" + + import cutlass + import cutlass.cute as cute + + device_index = compile_spec.device_index + num_heads = compile_spec.num_heads + kv_lora_rank = compile_spec.kv_lora_rank + qk_rope_head_dim = compile_spec.qk_rope_head_dim + page_size = compile_spec.page_size + max_seq_len_q = compile_spec.max_seq_len_q + packed_query = compile_spec.packed_query + kernel = compile_spec.kernel + dtype_map = { + "bfloat16": cutlass.BFloat16, + "float8_e4m3fn": cutlass.Float8E4M3FN, + } + qkv_dtype = dtype_map[compile_spec.q_dtype_key] + output_dtype = dtype_map[compile_spec.output_dtype_key] + physical_pages = cute.sym_int() + batch_size = cute.sym_int() + runtime_total_q = cute.sym_int() + + # These fake tensors pin the public ABI while allowing runtime page counts, + # table widths/row strides, and batch metadata pointers to vary. + q_stride_h = _MLA_QUERY_DIM + q_stride_q = num_heads * _MLA_QUERY_DIM + q_latent_shape: tuple[int, ...] + q_rope_shape: tuple[int, ...] + q_stride: tuple[int, ...] + if packed_query: + q_latent_shape = (num_heads, kv_lora_rank, runtime_total_q) + q_rope_shape = (num_heads, qk_rope_head_dim, runtime_total_q) + q_stride = (q_stride_h, 1, q_stride_q) + else: + q_stride_batch = max_seq_len_q * q_stride_q + q_latent_shape = (num_heads, kv_lora_rank, max_seq_len_q, batch_size) + q_rope_shape = ( + num_heads, + qk_rope_head_dim, + max_seq_len_q, + batch_size, + ) + q_stride = (q_stride_h, 1, q_stride_q, q_stride_batch) + q_latent_fake = cute.runtime.make_fake_tensor( + qkv_dtype, q_latent_shape, stride=q_stride, assumed_align=16 + ) + q_rope_fake = cute.runtime.make_fake_tensor( + qkv_dtype, q_rope_shape, stride=q_stride, assumed_align=16 + ) + cache_token_stride = _MLA_QUERY_DIM + cache_page_stride = page_size * _MLA_QUERY_DIM + c_latent_fake = cute.runtime.make_fake_tensor( + qkv_dtype, + (page_size, kv_lora_rank, physical_pages), + stride=(cache_token_stride, 1, cache_page_stride), + assumed_align=16, + ) + c_rope_fake = cute.runtime.make_fake_tensor( + qkv_dtype, + (page_size, qk_rope_head_dim, physical_pages), + stride=(cache_token_stride, 1, cache_page_stride), + assumed_align=16, + ) + runtime_page_columns = cute.sym_int() + runtime_page_row_stride = cute.sym_int64(divisibility=1) + page_offsets_fake = cute.runtime.make_fake_tensor( + cutlass.Int32, + (runtime_page_columns, batch_size), + stride=(1, runtime_page_row_stride), + assumed_align=4, + ) + out_stride_row = num_heads * kv_lora_rank + out_shape: tuple[int, ...] + out_stride: tuple[int, ...] + lse_shape: tuple[int, ...] + lse_stride: tuple[int, ...] + if packed_query: + out_shape = (num_heads, kv_lora_rank, runtime_total_q) + out_stride = (kv_lora_rank, 1, out_stride_row) + lse_shape = (num_heads, runtime_total_q) + lse_stride = (1, num_heads) + else: + out_stride_batch = max_seq_len_q * out_stride_row + out_shape = (num_heads, kv_lora_rank, max_seq_len_q, batch_size) + out_stride = (kv_lora_rank, 1, out_stride_row, out_stride_batch) + lse_shape = (num_heads, max_seq_len_q, batch_size) + lse_stride = (1, num_heads, max_seq_len_q * num_heads) + out_fake = cute.runtime.make_fake_tensor( + output_dtype, out_shape, stride=out_stride, assumed_align=16 + ) + lse_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, lse_shape, stride=lse_stride, assumed_align=16 + ) + workspace_fake = None + if compile_spec.has_kernel_workspace: + workspace_bytes = cute.sym_int() + workspace_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int8, + (workspace_bytes,), + stride_order=(0,), + assumed_align=32, + ) + cache_seqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (batch_size,), + stride_order=(0,), + assumed_align=4, + ) + qo_indptr_fake = None + if packed_query: + runtime_num_q_offsets = cute.sym_int() + qo_indptr_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (runtime_num_q_offsets,), + stride_order=(0,), + assumed_align=4, + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + # Task objects carry loop-local state through generated control flow, so + # select the public staged frontend for this compilation. + with torch.cuda.device(device_index): + compiled = cute.compile[cute.FrontendNext]( + kernel, + q_latent_fake, + q_rope_fake, + c_latent_fake, + c_rope_fake, + page_offsets_fake, + out_fake, + lse_fake, + workspace_fake, + cutlass.Int32(compile_spec.split_kv), + cache_seqs_fake, + qo_indptr_fake, + None, + cutlass.Float32(1.0), + cutlass.Float32(1.0), + stream_fake, + options=_COMPILE_OPTIONS, + ) + return compiled + + +def get_prims_ts_batch_mla_decode_workspace_size( + batch_size: int, + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + page_size: int, + max_seq_len: int, + *, + seq_len_q: Optional[int] = None, + max_seq_len_q: Optional[int] = None, + q_dtype: torch.dtype = torch.bfloat16, + kv_dtype: Optional[torch.dtype] = None, + out_dtype: torch.dtype = torch.bfloat16, + mask_type: Literal["dense", "causal"] = "causal", + device=None, +) -> int: + """Return caller-workspace bytes for one automatic MLA policy. + + The arguments define the static geometry used to resolve the same automatic + policy and private scratch layout as + :func:`prims_ts_batch_mla_decode_with_kv_cache`, without compiling a kernel. + ``max_seq_len_q`` is the static per-request Q bound for both fixed and + packed-query launches; + ``seq_len_q`` remains a backward-compatible fixed-Q alias. If neither is + supplied, the bound is one. The returned byte count includes both split-KV + scratch and the internal FP32 LSE tensor. Allocate a contiguous + ``torch.int8`` or ``torch.uint8`` CUDA buffer; MLA does not require its + contents to be initialized before first use. + """ + + batch_size = _validate_positive_int(batch_size, "batch_size") + num_heads = _validate_positive_int(num_heads, "num_heads") + _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) + page_size = _validate_page_size(page_size) + max_seq_len = _validate_mla_max_kv_len(max_seq_len, "max_seq_len") + max_seq_len_q = _resolve_max_seq_len_q_alias( + seq_len_q=seq_len_q, + max_seq_len_q=max_seq_len_q, + default=1, + ) + assert max_seq_len_q is not None + _validate_mla_query_head_extent( + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + ) + _validate_mask(mask_type) + if kv_dtype is None: + kv_dtype = q_dtype + _validate_mla_dtype_pair(q_dtype, kv_dtype, out_dtype) + _, device_index = _resolve_cuda_device(device) + + spec = _resolve_mla_decode_launch_spec( + device_index, + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + page_size, + max_seq_len, + _dtype_key(q_dtype), + _dtype_key(kv_dtype), + _dtype_key(out_dtype), + mask_type, + max_seq_len_q, + ) + return _make_mla_workspace_layout( + spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q + ).total_bytes + + +def _prepare_mla_runtime( + query: torch.Tensor, + kv_cache: torch.Tensor, + *, + device: torch.device, + batch_size: int, + num_heads: int, + max_seq_len_q: int, + packed_query: bool, + qo_indptr: Optional[torch.Tensor], + page_size: int, + q_dtype: torch.dtype, + kv_dtype: torch.dtype, + output_dtype: torch.dtype, + bmm1_scale: float, + bmm2_scale: float, + out: Optional[torch.Tensor], + validate: bool, +) -> _MLARuntime: + """Normalize one launch, optionally validating its public arguments.""" + + if validate: + if packed_query: + if qo_indptr is None: + raise ValueError("packed-query MLA run requires qo_indptr") + _validate_qo_indptr( + qo_indptr, + device=device, + batch_size=batch_size, + ) + elif qo_indptr is not None: + raise ValueError("fixed-query MLA plan does not accept qo_indptr") + _validate_query( + query, + packed_query=packed_query, + device=device, + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + q_dtype=q_dtype, + ) + ( + normalized_cache, + num_physical_pages, + runtime_page_size, + ) = _normalize_mla_kv_cache(kv_cache, expected_device=device) + if runtime_page_size != page_size: + raise ValueError( + "kv_cache page size does not match the launch: expected " + f"{page_size}, got {runtime_page_size}" + ) + if normalized_cache.dtype != kv_dtype: + raise ValueError( + f"kv_cache dtype must match the launch ({kv_dtype}), " + f"got {normalized_cache.dtype}" + ) + effective_bmm1_scale = _validate_scale(bmm1_scale, "bmm1_scale") + effective_bmm2_scale = _validate_scale(bmm2_scale, "bmm2_scale") + else: + normalized_cache = kv_cache[:, 0] if kv_cache.ndim == 4 else kv_cache + num_physical_pages = int(normalized_cache.shape[0]) + effective_bmm1_scale = bmm1_scale + effective_bmm2_scale = bmm2_scale + total_q = int(query.shape[0]) if packed_query else None + if out is None: + out_shape = ( + (total_q, num_heads, _MLA_LATENT_DIM) + if packed_query + else (batch_size, max_seq_len_q, num_heads, _MLA_LATENT_DIM) + ) + out = torch.empty(out_shape, device=device, dtype=output_dtype) + elif validate: + _validate_out( + out, + device=device, + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + total_q=total_q, + output_dtype=output_dtype, + ) + return _MLARuntime( + query=query, + normalized_cache=normalized_cache, + out=out, + num_physical_pages=num_physical_pages, + bmm1_scale=effective_bmm1_scale, + bmm2_scale=effective_bmm2_scale, + ) + + +def _validate_mla_output_aliasing( + runtime: _MLARuntime, + *, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + qo_indptr: Optional[torch.Tensor], + workspace_buffer: torch.Tensor, +) -> None: + """Keep output disjoint from every live MLA decode allocation.""" + + _validate_out_does_not_overlap_inputs( + runtime.out, + ("query", runtime.query), + ("kv_cache", runtime.normalized_cache), + ("block_tables", block_tables), + ("seq_lens", seq_lens), + ("qo_indptr", qo_indptr), + ("workspace_buffer", workspace_buffer), + ) + + +def _launch_mla_decode( + runtime: _MLARuntime, + *, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + qo_indptr: Optional[torch.Tensor], + packed_query: bool, + kv_lora_rank: int, + split_kv: int, + workspace: _MLAWorkspaceViews, + compiled: Callable[..., object], +) -> torch.Tensor: + """Form the dimension-first views and launch one compiled MLA kernel.""" + + if packed_query and int(runtime.query.shape[0]) == 0: + return runtime.out + if packed_query: + q_latent = runtime.query[..., :kv_lora_rank].permute(1, 2, 0) + q_rope = runtime.query[..., kv_lora_rank:].permute(1, 2, 0) + out_kernel = runtime.out.permute(1, 2, 0) + total_q = int(runtime.query.shape[0]) + lse_kernel = workspace.lse.view(-1, workspace.lse.shape[-1])[ + :total_q + ].transpose(0, 1) + else: + q_latent = runtime.query[..., :kv_lora_rank].permute(2, 3, 1, 0) + q_rope = runtime.query[..., kv_lora_rank:].permute(2, 3, 1, 0) + out_kernel = runtime.out.permute(2, 3, 1, 0) + lse_kernel = workspace.lse.permute(2, 1, 0) + c_latent = runtime.normalized_cache[..., :kv_lora_rank].permute(1, 2, 0) + c_rope = runtime.normalized_cache[..., kv_lora_rank:].permute(1, 2, 0) + page_offsets = block_tables.transpose(0, 1) + compiled( + q_latent, + q_rope, + c_latent, + c_rope, + page_offsets, + out_kernel, + lse_kernel, + workspace.kernel_workspace, + split_kv, + seq_lens, + qo_indptr, + None, + runtime.bmm1_scale, + runtime.bmm2_scale, + ) + return runtime.out + + +@flashinfer_api(trace=prims_ts_decode_mla_trace_dispatch) +def prims_ts_batch_mla_decode_with_kv_cache( + query: torch.Tensor, + kv_cache: torch.Tensor, + workspace_buffer: torch.Tensor, + kv_lora_rank: int, + qk_rope_head_dim: int, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + *, + qo_indptr: Optional[torch.Tensor] = None, + max_seq_len_q: Optional[int] = None, + out: Optional[torch.Tensor] = None, + bmm1_scale: float = 1.0, + bmm2_scale: float = 1.0, + mask_type: Literal["dense", "causal"] = "causal", + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Launch fixed or packed-query paged MLA decode with caller-owned scratch. + + With ``qo_indptr=None``, ``query`` has fixed shape ``[B, SQ, H, 576]``. + Otherwise ``query`` has compact shape ``[total_q, H, 576]`` and + ``qo_indptr`` contains the ``B + 1`` cumulative Q offsets. Runtime Q + lengths are exclusively ``qo_indptr[b + 1] - qo_indptr[b]``; + ``max_seq_len_q`` is only the static policy, JIT, and workspace bound and + is required for compact launches. Individual packed requests may be empty, + and an all-empty launch returns its empty output without dispatching a GPU + kernel. The last query dimension concatenates the 512 latent and 64 RoPE + dimensions. ``kv_cache`` accepts compact rank-3 + ``[pages, page_size, 576]`` or rank-4 ``[pages, 1, page_size, 576]`` + storage. ``block_tables`` and ``seq_lens`` follow FlashInfer's native dense + paged-cache ABI. The table is contiguous within each row and may have + padding between rows; ``max_seq_len`` is the exact static policy/JIT + maximum. + Causal masking is bottom-right aligned: query row ``i`` can attend through + KV row ``seq_lens[b] - q_len[b] + i`` for request ``b``. + + The workspace is exclusive to one in-flight launch or captured graph and + must not overlap query, K/V cache, metadata, or output storage. + Runtime K/V lengths must remain positive and no larger than ``max_seq_len``; + this hot path deliberately performs no device-to-host metadata reads. For + packed launches, callers must ensure that offsets start at zero, are + nondecreasing, end at ``query.shape[0]``, and have every delta no + larger than ``max_seq_len_q``. For causal masking, every fixed or packed + per-request Q length must also be no greater than the corresponding live + ``seq_lens`` value. Warm the planned topology before CUDA graph + capture and provide ``out`` to avoid an output allocation. Captured graphs + must retain stable ``block_tables``, ``seq_lens``, and, for packed Q, + ``qo_indptr`` storage. Values may change only between completed replays + while the runtime metadata contracts and captured query/output extents + remain valid. No backend fallback or scheduling knob is exposed. + + Parameters + ---------- + query : torch.Tensor + Fixed or packed query tensor with concatenated latent and RoPE heads. + kv_cache : torch.Tensor + Compact paged latent K/V cache. + workspace_buffer : torch.Tensor + Caller-owned byte workspace for this planned layout. + kv_lora_rank, qk_rope_head_dim : int + Latent and RoPE dimensions. + block_tables : torch.Tensor + Dense physical-page table for each request. Rows must be inner + contiguous and non-overlapping, but may have padding between them. + seq_lens : torch.Tensor + Live K/V sequence lengths. + max_seq_len : int + Static maximum K/V length used for policy selection and JIT caching. + qo_indptr : torch.Tensor, optional + Cumulative query offsets selecting packed-query mode. + max_seq_len_q : int, optional + Static packed-query length bound. + out : torch.Tensor, optional + Caller-owned output tensor. + bmm1_scale, bmm2_scale : float + QK and value/output scaling factors. + mask_type : {"dense", "causal"} + Attention mask mode. + out_dtype : torch.dtype + Output dtype. + """ + + packed_query = qo_indptr is not None + _validate_query(query, packed_query=packed_query) + metadata_device, batch_size, max_num_pages = _validate_mla_metadata( + block_tables, seq_lens + ) + if metadata_device != query.device: + raise ValueError( + f"MLA metadata must be on {query.device}, got {metadata_device}" + ) + normalized_cache, _, page_size = _normalize_mla_kv_cache( + kv_cache, expected_device=query.device + ) + if packed_query: + _validate_qo_indptr( + qo_indptr, + device=query.device, + batch_size=batch_size, + ) + if max_seq_len_q is None: + raise ValueError( + "max_seq_len_q is required when qo_indptr selects packed query" + ) + max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") + num_heads = int(query.shape[1]) + else: + fixed_seq_len_q = int(query.shape[1]) + if max_seq_len_q is None: + max_seq_len_q = fixed_seq_len_q + else: + max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") + if max_seq_len_q != fixed_seq_len_q: + raise ValueError( + "fixed query length must equal max_seq_len_q: " + f"got SQ={fixed_seq_len_q} and max_seq_len_q={max_seq_len_q}" + ) + num_heads = int(query.shape[2]) + _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) + _validate_page_size(page_size) + max_seq_len = _validate_mla_max_kv_len(max_seq_len, "max_seq_len") + required_page_columns = _ceil_div(max_seq_len, page_size) + if max_num_pages < required_page_columns: + raise ValueError( + "block_tables must have at least ceil(max_seq_len / page_size) " + f"columns ({required_page_columns}), got {max_num_pages}" + ) + _validate_mask(mask_type) + _validate_mla_dtype_pair(query.dtype, normalized_cache.dtype, out_dtype) + device_index = _validate_runtime_device(query.device) + spec_key = ( + device_index, + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + page_size, + max_seq_len, + _dtype_key(query.dtype), + _dtype_key(normalized_cache.dtype), + _dtype_key(out_dtype), + mask_type, + max_seq_len_q, + ) + spec = _resolve_mla_decode_launch_spec(*spec_key) + layout = _make_mla_workspace_layout( + spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q + ) + _validate_workspace_buffer( + workspace_buffer, + device=query.device, + required_bytes=layout.total_bytes, + ) + caller_provided_out = out is not None + runtime = _prepare_mla_runtime( + query, + normalized_cache, + device=query.device, + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + qo_indptr=qo_indptr, + page_size=page_size, + q_dtype=query.dtype, + kv_dtype=normalized_cache.dtype, + output_dtype=out_dtype, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + out=out, + validate=True, + ) + _validate_tensor_does_not_overlap_inputs( + workspace_buffer, + "workspace_buffer", + ("query", runtime.query), + ("kv_cache", runtime.normalized_cache), + ("block_tables", block_tables), + ("seq_lens", seq_lens), + ("qo_indptr", qo_indptr), + ("out", runtime.out), + ) + if caller_provided_out: + _validate_mla_output_aliasing( + runtime, + block_tables=block_tables, + seq_lens=seq_lens, + qo_indptr=qo_indptr, + workspace_buffer=workspace_buffer, + ) + compile_spec = _make_mla_decode_compile_spec( + spec, + device_index=device_index, + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + page_size=page_size, + q_dtype_key=_dtype_key(query.dtype), + output_dtype_key=_dtype_key(out_dtype), + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + ) + compiled = _get_compiled_mla_decode(compile_spec) + workspace = _bind_mla_workspace(workspace_buffer, layout) + return _launch_mla_decode( + runtime, + block_tables=block_tables, + seq_lens=seq_lens, + qo_indptr=qo_indptr, + packed_query=packed_query, + kv_lora_rank=kv_lora_rank, + split_kv=spec.split_kv, + workspace=workspace, + compiled=compiled, + ) + + +class BatchMLADecodePagedTSWrapper: + """Compile and reuse task-scheduled paged MLA decode launches.""" + + @flashinfer_api + def __init__(self) -> None: + """Initialize an unplanned task-scheduled paged-MLA wrapper.""" + self._plan_state: Optional[_MLADecodePlanState] = None + + @flashinfer_api + def plan( + self, + device: int | str | torch.device, + batch_size: int, + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + page_size: int, + max_kv_len: int, + *, + max_seq_len_q: int, + packed_query: bool, + q_data_type: torch.dtype, + kv_data_type: torch.dtype, + o_data_type: torch.dtype, + mask_type: Literal["dense", "causal"] = "causal", + workspace_buffer: Optional[torch.Tensor] = None, + ) -> None: + """Compile one static MLA shape and bind its reusable workspace. + + Planning consumes only compile-time geometry, capacity, dtype, and + storage-mode inputs. Request metadata belongs exclusively to + :meth:`run` and is never retained by the wrapper. A successful re-plan + atomically replaces the previous immutable state; a failed re-plan + leaves that state usable. + + ``packed_query=False`` selects fixed ``[B, SQ, H, 576]`` query storage, + where ``SQ`` is exactly ``max_seq_len_q``. ``packed_query=True`` selects + ``[total_q, H, 576]`` storage with per-run cumulative offsets supplied to + every run. ``max_seq_len_q`` is then the per-request capacity. + Individual packed requests may be empty; an all-empty run returns its + empty output without dispatching a GPU kernel. + + If ``workspace_buffer`` is omitted, the plan allocates private scratch. + A workspace is mutable and exclusive to one in-flight launch or graph + replay. Warm the plan before graph capture, and call ``run`` with + ``validate=False`` inside compiled or captured regions. + + Parameters + ---------- + device : int, str, or torch.device + CUDA device on which the plan will execute. + batch_size : int + Exact runtime request count. + num_heads, kv_lora_rank, qk_rope_head_dim, page_size : int + Static MLA head geometry and K/V page size. + max_kv_len, max_seq_len_q : int + Static per-request K/V and Q capacities. + packed_query : bool + Select packed rather than fixed query storage. + q_data_type, kv_data_type, o_data_type : torch.dtype + Query, K/V, and output dtypes used to compile the plan. + mask_type : {"dense", "causal"} + Attention mask mode. + workspace_buffer : torch.Tensor, optional + Caller-owned contiguous int8 or uint8 scratch on ``device``. It + must be 32-byte aligned and large enough for the selected plan. + When omitted, planning allocates the buffer. The retained buffer + is exclusive to one in-flight launch or graph replay. + """ + + if not isinstance(packed_query, bool): + raise TypeError("packed_query must be a bool") + _validate_mask(mask_type) + batch_size = _validate_positive_int(batch_size, "batch_size") + _validate_mla_int32_extent(batch_size, "batch_size") + num_heads = _validate_positive_int(num_heads, "num_heads") + _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) + page_size = _validate_page_size(page_size) + max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") + max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") + _validate_mla_query_head_extent( + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + ) + _validate_mla_dtype_pair(q_data_type, kv_data_type, o_data_type) + device, device_index = _resolve_cuda_device(device) + required_page_columns = _ceil_div(max_kv_len, page_size) + + spec_key = ( + device_index, + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + page_size, + max_kv_len, + _dtype_key(q_data_type), + _dtype_key(kv_data_type), + _dtype_key(o_data_type), + mask_type, + max_seq_len_q, + ) + spec = _resolve_mla_decode_launch_spec(*spec_key) + compile_spec = _make_mla_decode_compile_spec( + spec, + device_index=device_index, + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + page_size=page_size, + q_dtype_key=_dtype_key(q_data_type), + output_dtype_key=_dtype_key(o_data_type), + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + ) + policy = spec.policy + workspace_layout = _make_mla_workspace_layout( + spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q + ) + if workspace_buffer is None: + workspace_buffer = torch.empty( + workspace_layout.total_bytes, device=device, dtype=torch.int8 + ) + else: + _validate_workspace_buffer( + workspace_buffer, + device=device, + required_bytes=workspace_layout.total_bytes, + ) + workspace = _bind_mla_workspace(workspace_buffer, workspace_layout) + compiled = _get_compiled_mla_decode(compile_spec) + + # Publish only after validation, compilation, allocation, and binding + # succeed, so a failed re-plan leaves the previous plan usable. + self._plan_state = _MLADecodePlanState( + device=device, + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + page_size=page_size, + q_dtype=q_data_type, + kv_dtype=kv_data_type, + output_dtype=o_data_type, + mask_type=mask_type, + max_kv_len=max_kv_len, + required_page_columns=required_page_columns, + workspace_buffer=workspace_buffer, + workspace_layout=workspace_layout, + workspace_views=workspace, + compiled=compiled, + policy=policy, + split_kv=int(dict(policy)["split_kv"]), + ) + + @flashinfer_api + def run( + self, + query: torch.Tensor, + kv_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + *, + qo_indptr: Optional[torch.Tensor] = None, + bmm1_scale: float = 1.0, + bmm2_scale: float = 1.0, + out: Optional[torch.Tensor] = None, + validate: bool = True, + ) -> torch.Tensor: + """Launch the most recently planned MLA decode on the current stream. + + ``block_tables`` and ``seq_lens`` are required per-run bindings. + ``qo_indptr`` is required by a packed-query plan and rejected by a + fixed-query plan. With validation enabled, tensor structure, metadata + values, scales, aliases, and every static capacity are checked before + launch. These checks synchronize metadata to the host. Set + ``validate=False`` only after validating representative inputs, and use + it for ``torch.compile`` or CUDA graph capture. + + Parameters + ---------- + query : torch.Tensor + Runtime fixed or packed query tensor matching the plan. + kv_cache : torch.Tensor + Runtime compact paged latent K/V cache. + block_tables : torch.Tensor + Runtime physical-page table with one inner-contiguous, + non-overlapping row per request. Inter-row padding is accepted. + seq_lens : torch.Tensor + Runtime K/V lengths with one element per request. + qo_indptr : torch.Tensor, optional + Runtime cumulative query offsets for a packed-query plan. + bmm1_scale, bmm2_scale : float + QK and value/output scaling factors. + out : torch.Tensor, optional + Caller-owned output tensor. A new tensor is allocated when omitted. + validate : bool + Enable explicit runtime validation. Defaults to ``True``. + + Returns + ------- + torch.Tensor + The fixed or packed MLA attention output. + """ + + state = self._plan_state + if state is None: + raise RuntimeError("plan() must be called before run()") + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") + runtime_qo_indptr = qo_indptr if state.packed_query else None + caller_provided_out = out is not None + runtime = _prepare_mla_runtime( + query, + kv_cache, + device=state.device, + batch_size=state.batch_size, + num_heads=state.num_heads, + max_seq_len_q=state.max_seq_len_q, + packed_query=state.packed_query, + qo_indptr=runtime_qo_indptr, + page_size=state.page_size, + q_dtype=state.q_dtype, + kv_dtype=state.kv_dtype, + output_dtype=state.output_dtype, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + out=out, + validate=validate, + ) + if validate: + _validate_mla_run_metadata( + state, + runtime, + block_tables, + seq_lens, + qo_indptr, + ) + _validate_tensor_does_not_overlap_inputs( + state.workspace_buffer, + "workspace_buffer", + ("query", runtime.query), + ("kv_cache", runtime.normalized_cache), + ("block_tables", block_tables), + ("seq_lens", seq_lens), + ("qo_indptr", runtime_qo_indptr), + ("out", runtime.out), + ) + if caller_provided_out: + _validate_mla_output_aliasing( + runtime, + block_tables=block_tables, + seq_lens=seq_lens, + qo_indptr=runtime_qo_indptr, + workspace_buffer=state.workspace_buffer, + ) + return _launch_mla_decode( + runtime, + block_tables=block_tables, + seq_lens=seq_lens, + qo_indptr=runtime_qo_indptr, + packed_query=state.packed_query, + kv_lora_rank=state.kv_lora_rank, + split_kv=state.split_kv, + workspace=state.workspace_views, + compiled=state.compiled, + ) + + +@flashinfer_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) +def batch_mla_decode_with_paged_kv_cache( + query: torch.Tensor, + kv_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + *, + qo_indptr: Optional[torch.Tensor] = None, + max_seq_len_q: Optional[int] = None, + kv_lora_rank: int = _MLA_LATENT_DIM, + qk_rope_head_dim: int = _MLA_ROPE_DIM, + mask_type: Literal["dense", "causal"] = "causal", + max_kv_len: Optional[int] = None, + bmm1_scale: float = 1.0, + bmm2_scale: float = 1.0, + out: Optional[torch.Tensor] = None, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """One-shot convenience wrapper for fixed or packed-query MLA decode. + + This helper reads ``seq_lens`` and, for packed Q, ``qo_indptr`` on the host + to derive plan bounds, then constructs a temporary wrapper. Invoke it + outside CUDA Graph capture. Capture-sensitive callers should pre-plan + :class:`BatchMLADecodePagedTSWrapper` and use ``run(validate=False)``. + + Parameters + ---------- + query : torch.Tensor + Fixed or packed query tensor with concatenated latent and RoPE heads. + kv_cache : torch.Tensor + Compact paged latent K/V cache. + block_tables : torch.Tensor + Dense physical-page table for each request. Rows must be inner + contiguous and non-overlapping, but may have padding between them. + seq_lens : torch.Tensor + Live K/V sequence lengths. + qo_indptr : torch.Tensor, optional + Cumulative query offsets selecting packed-query mode. + max_seq_len_q : int, optional + Per-request packed-query length capacity. For packed Q it defaults to + the maximum delta in ``qo_indptr``; an explicit value may be larger. + For fixed Q it defaults to the query's sequence extent and, when + provided, must equal that extent. + kv_lora_rank, qk_rope_head_dim : int + Latent and RoPE dimensions. + mask_type : {"dense", "causal"} + Attention mask mode. + max_kv_len : int, optional + Static K/V length bound; defaults to the metadata maximum. + bmm1_scale, bmm2_scale : float + QK and value/output scaling factors. + out : torch.Tensor, optional + Caller-owned output tensor. + out_dtype : torch.dtype + Output dtype. + + Returns + ------- + torch.Tensor + The fixed or packed MLA attention output. + """ + + packed_query = qo_indptr is not None + _validate_query(query, packed_query=packed_query) + metadata_device, batch_size, _ = _validate_mla_metadata(block_tables, seq_lens) + if metadata_device != query.device: + raise ValueError( + f"MLA metadata must be on {query.device}, got {metadata_device}" + ) + normalized_cache, _, page_size = _normalize_mla_kv_cache( + kv_cache, expected_device=query.device + ) + _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) + _validate_page_size(page_size) + _validate_mla_dtype_pair(query.dtype, normalized_cache.dtype, out_dtype) + if packed_query: + _validate_qo_indptr( + qo_indptr, + device=query.device, + batch_size=batch_size, + ) + num_heads = int(query.shape[1]) + derived_max_seq_len_q, total_q, runtime_q_lengths = _derive_max_seq_len_q( + qo_indptr, + batch_size=batch_size, + ) + if total_q != int(query.shape[0]): + raise ValueError( + "qo_indptr must end at the packed query row count " + f"({query.shape[0]}), got {total_q}" + ) + if max_seq_len_q is None: + if derived_max_seq_len_q == 0: + raise ValueError( + "max_seq_len_q is required for an all-empty packed query" + ) + max_seq_len_q = derived_max_seq_len_q + else: + max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") + if derived_max_seq_len_q > max_seq_len_q: + raise ValueError( + "qo_indptr contains a per-request Q length larger than " + f"max_seq_len_q ({max_seq_len_q}): got " + f"{derived_max_seq_len_q}" + ) + _validate_mla_query_head_extent( + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + total_q=int(query.shape[0]), + ) + else: + num_heads = int(query.shape[2]) + fixed_seq_len_q = int(query.shape[1]) + if max_seq_len_q is None: + max_seq_len_q = fixed_seq_len_q + else: + max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") + if max_seq_len_q != fixed_seq_len_q: + raise ValueError( + "fixed query length must equal max_seq_len_q: " + f"got SQ={fixed_seq_len_q} and " + f"max_seq_len_q={max_seq_len_q}" + ) + _validate_mla_query_head_extent( + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + ) + runtime_q_lengths = (max_seq_len_q,) * batch_size + seq_lens_host = tuple(int(value) for value in seq_lens.tolist()) + if any(seq_len <= 0 for seq_len in seq_lens_host): + raise ValueError("every runtime request must contain at least one KV token") + metadata_max_kv_len = max(seq_lens_host) + if max_kv_len is None: + max_kv_len = metadata_max_kv_len + else: + max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") + if metadata_max_kv_len > max_kv_len: + raise ValueError( + "runtime KV metadata contains a request longer than " + f"max_kv_len ({max_kv_len}): got {metadata_max_kv_len}" + ) + if mask_type == "causal": + for request_idx, (q_len, kv_len) in enumerate( + zip(runtime_q_lengths, seq_lens_host, strict=True) + ): + if q_len > kv_len: + raise ValueError( + "causal MLA decode requires every per-request Q length " + "to be no greater than its K/V length; request " + f"{request_idx} has Q={q_len} and K/V={kv_len}" + ) + assert max_seq_len_q is not None + assert max_kv_len is not None + if out is not None: + _validate_out( + out, + device=query.device, + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + total_q=int(query.shape[0]) if packed_query else None, + output_dtype=out_dtype, + ) + + wrapper = BatchMLADecodePagedTSWrapper() + wrapper.plan( + query.device, + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + page_size, + max_kv_len, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + q_data_type=query.dtype, + kv_data_type=normalized_cache.dtype, + o_data_type=out_dtype, + mask_type=mask_type, + ) + return wrapper.run( + query, + normalized_cache, + block_tables, + seq_lens, + qo_indptr=qo_indptr, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + out=out, + ) + + +__all__ = [ + "BatchMLADecodePagedTSWrapper", + "batch_mla_decode_with_paged_kv_cache", + "get_prims_ts_batch_mla_decode_workspace_size", + "prims_ts_batch_mla_decode_with_kv_cache", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/split_kv_mode_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/split_kv_mode_policy.py new file mode 100644 index 000000000000..083cc9b68ef7 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/split_kv_mode_policy.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared split-KV reduction-mode policy for FMHA and MLA decode.""" + +from collections.abc import Iterable + + +DIRECT_MODES = frozenset({"direct", "disabled"}) +SEPARATE_MODES = frozenset({"gmem_separate", "gmem_reduction_with_separate_kernel"}) +INLINE_MODES = frozenset({"gmem_inline", "gmem_reduction"}) +CLUSTER_MODES = frozenset({"cluster_smem", "cluster", "cluster_smem_reduction"}) + + +def canonical_split_kv_mode(mode: str) -> str: + """Return the canonical name for a production or selector mode spelling.""" + + normalized = str(mode).strip().lower() + if normalized in DIRECT_MODES: + return "direct" + if normalized in SEPARATE_MODES: + return "gmem_separate" + if normalized in INLINE_MODES: + return "gmem_inline" + if normalized in CLUSTER_MODES: + return "cluster_smem" + raise ValueError(f"unsupported reduction mode {mode!r}") + + +def select_split_kv_modes( + *, + family: str, + topology: str, + tile_size_q: int, + head_dim: int, + head_dim_per_cta_v: int | None, + split_kv: int, + available_modes: Iterable[str], +) -> tuple[str, ...]: + """Return available split-KV modes in the order they should be tried. + + This policy only orders modes. Callers remain responsible for production + support, SMEM limits, model coverage, and exact cluster one-wave residency. + """ + + if split_kv < 1: + raise ValueError("split_kv must be positive") + + family = family.strip().lower() + if family not in {"fmha_decode", "mla_decode"}: + raise ValueError(f"unsupported decode family {family!r}") + topology = topology.strip().lower() + modes_by_name = {canonical_split_kv_mode(mode): mode for mode in available_modes} + if split_kv == 1: + direct = modes_by_name.get("direct") + if direct is None: + raise ValueError("split_kv=1 requires direct/disabled mode") + return (direct,) + + mode_order: tuple[str, ...] + if family == "fmha_decode": + # Keep automatic selection structural: use cluster when the caller's + # exact residency/support checks accept it, otherwise prefer the + # standalone reducer and retain inline reduction only as a support + # fallback. This avoids shape-specific measured crossover tables. + mode_order = ( + "cluster_smem", + "gmem_separate", + "gmem_inline", + ) + else: + # Prefer cluster for every structurally capable 1CTA MLA split profile. + # The profile factory rejects Keeps-MMA-AB and incomplete Q tiles, + # applies the static SMEM budget, and the public planner performs the + # exact cluster-size occupancy query before accepting cluster. Keeping + # shape values out of this ordering avoids a second, narrower support + # matrix that can disagree with those authoritative checks. + use_cluster = topology == "1cta" + mode_order = ( + ("cluster_smem", "gmem_separate") if use_cluster else ("gmem_separate",) + ) + + return tuple(modes_by_name[mode] for mode in mode_order if mode in modes_by_name) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_utils.py index 71d7c37f7fa7..e7e499b7d7a8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_utils.py @@ -29,11 +29,13 @@ def _install_msa_cutlass_compatibility() -> None: except ImportError: return - # MSA has not yet migrated these two aliases to their CUTLASS DSL 4.6 + # MSA has not yet migrated these aliases to their newer CUTLASS DSL # names. Keep this shim local to the MSA import path and remove it once the - # packaged sources use cute.ThrMma and cute.make_rmem_tensor directly. - if not hasattr(cute.core, "ThrMma"): - setattr(cute.core, "ThrMma", cute.ThrMma) + # packaged sources use the top-level thread classes and + # cute.make_rmem_tensor directly. + for thread_class_name in ("ThrCopy", "ThrMma"): + if not hasattr(cute.core, thread_class_name): + setattr(cute.core, thread_class_name, getattr(cute, thread_class_name)) if not hasattr(cute, "make_fragment"): setattr(cute, "make_fragment", cute.make_rmem_tensor) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py index d134d3f9a8dc..5024a43b192a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -27,8 +27,21 @@ from ...attention.backends.interface import PredefinedAttentionMask from .interface import AttentionBackend, AttentionTensorLayout + +def _install_cutlass_dsl_compatibility() -> None: + """Restore CuTe aliases required by pinned third-party FA4 and QuACK.""" + import cutlass.cute as cute + + for name in ("ThrCopy", "ThrMma"): + if not hasattr(cute.core, name) and hasattr(cute, name): + setattr(cute.core, name, getattr(cute, name)) + if not hasattr(cute, "make_fragment") and hasattr(cute, "make_rmem_tensor"): + cute.make_fragment = cute.make_rmem_tensor + + _flash_attn_fwd_import_error = None try: + _install_cutlass_dsl_compatibility() from flash_attn.cute.interface import _flash_attn_fwd except (ImportError, OSError) as e: _flash_attn_fwd = None diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index 482dff8f18c7..a2e1a78d2171 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -32,10 +32,12 @@ from tensorrt_llm._torch.distributed import all_to_all_4d, all_to_all_5d from ...attention.backends.interface import PredefinedAttentionMask +from .flash_attn4 import _install_cutlass_dsl_compatibility from .interface import AttentionBackend, AttentionTensorLayout _flash_attn_combine_import_error = None try: + _install_cutlass_dsl_compatibility() from flash_attn.cute.interface import flash_attn_combine as _flash_attn_combine except (ImportError, OSError) as e: _flash_attn_combine = None diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0ead08256971..c8102d37d26c 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -21,6 +21,7 @@ import pytest import torch +import torch._inductor.config as inductor_config from datasets import load_dataset from defs.conftest import get_sm_version, is_sm_100f from mpi4py.futures import MPIPoolExecutor @@ -79,6 +80,31 @@ def patched_start_mpi_pool(self): patched_start_mpi_pool) +def _count_prims_ts_phase_calls(mocker): + """Count PrimTS phase launches while keeping the real kernels installed.""" + from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha + + calls = { + "context": 0, + "generation": 0, + "mla_generation": 0, + } + + def patch_method(method_name, counter_name): + original = getattr(PrimsTSFmha, method_name) + + def counted(self, params): + calls[counter_name] += 1 + return original(self, params) + + mocker.patch.object(PrimsTSFmha, method_name, counted) + + patch_method("run_context", "context") + patch_method("run_generation", "generation") + patch_method("run_mla_generation", "mla_generation") + return calls + + # MPI session reuse cannot safely reset Userbuffers between Engines. Tests # using this helper must keep ``torch_compile=True`` in their node id; tests # with unconditional compile configs must keep ``piecewise_cuda_graph`` in @@ -1573,6 +1599,43 @@ class TestDeepSeekV3Lite(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V3-Lite" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V3-Lite/bf16" + @skip_pre_blackwell + @pytest.mark.skip_less_device_memory(60000) + def test_prims_ts_bfloat16(self, mocker): + if get_sm_version() not in (100, 103): + pytest.skip("PrimTS requires SM100 or SM103") + + calls = _count_prims_ts_phase_calls(mocker) + env = { + "TLLM_FMHA_LIBS": "+prims_ts", + "TLLM_WORKER_USE_SINGLE_PROCESS": "1", + } + kv_cache_config = KvCacheConfig( + free_gpu_memory_fraction=0.75, + tokens_per_block=32, + use_kv_cache_manager_v2=False, + ) + # Keep the TP=1 worker in this process so the call counter observes the + # real PrimTS launch. Compile Inductor kernels synchronously because its + # process-global async reader thread otherwise outlives this test and is + # reported as a leak; persistent compiler caches remain enabled. + with (inductor_config.patch(compile_threads=1), + mock.patch.dict(os.environ, env)): + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + kv_cache_config=kv_cache_config, + disable_overlap_scheduler=True, + enable_chunked_prefill=False, + enable_attention_dp=False, + cuda_graph_config=None, + speculative_config=None, + max_batch_size=1350) as llm: + calls.update({name: 0 for name in calls}) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + assert calls["mla_generation"] > 0 + @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("v2_kv_cache", [True, False]) # Chunked Prefill for MLA can only be enabled on SM100 @@ -4865,6 +4928,34 @@ def test_bf16(self, tp_size, pp_size, ep_size, attention_dp, cuda_graph, task = MMLU(self.MODEL_NAME) task.evaluate(llm) + @skip_pre_blackwell + def test_prims_ts_bfloat16(self, mocker): + if get_sm_version() not in (100, 103): + pytest.skip("PrimTS requires SM100 or SM103") + + calls = _count_prims_ts_phase_calls(mocker) + env = { + "TLLM_FMHA_LIBS": "+prims_ts", + "TLLM_WORKER_USE_SINGLE_PROCESS": "1", + } + kv_cache_config = KvCacheConfig( + free_gpu_memory_fraction=0.75, + tokens_per_block=32, + use_kv_cache_manager_v2=True, + ) + with mock.patch.dict(os.environ, env): + with LLM(f"{llm_models_root()}/Qwen3/Qwen3-8B", + attn_backend="TRTLLM", + kv_cache_config=kv_cache_config, + disable_overlap_scheduler=True, + cuda_graph_config=None) as llm: + calls.update({name: 0 for name in calls}) + task = CnnDailymail(self.MODEL_NAME) + task.evaluate(llm) + + assert calls["context"] > 0 + assert calls["generation"] > 0 + @parametrize_with_ids( "eagle3_one_model,enable_chunked_prefill,enable_max_concurrency,enable_draft_len_schedule", [ diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 1f7c08f508da..08c7968aec65 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -50,6 +50,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_prims_ts_bfloat16 - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=0-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] @@ -69,6 +70,7 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu_suspend_resume - accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dummy_load_format - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] # Cover nvbugs 5461712 and 5505402 + - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_prims_ts_bfloat16 - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dspark[TRTLLM] # SM100+ only; l0_h100 runs [VANILLA] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_trtllm-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8_mxfp4[fp8-latency-TRTLLM] @@ -280,6 +282,7 @@ l0_b200: - unittest/_torch/visual_gen/test_quant_ops.py - unittest/_torch/visual_gen/test_pertoken_adaln.py - unittest/_torch/visual_gen/test_attention_cute_dsl.py + - unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_integration.py diff --git a/tests/unittest/_torch/attention/fmha_test_utils.py b/tests/unittest/_torch/attention/fmha_test_utils.py index aa228e9269ae..01e5609a686d 100644 --- a/tests/unittest/_torch/attention/fmha_test_utils.py +++ b/tests/unittest/_torch/attention/fmha_test_utils.py @@ -23,7 +23,7 @@ class FakeAttention: - def __init__(self) -> None: + def __init__(self, local_layer_idx: int = 0) -> None: self.is_mla_enable = False self.kv_lora_rank = None self.v_head_dim = None @@ -33,6 +33,7 @@ def __init__(self) -> None: self.predicted_tokens_per_seq = 1 self.flashinfer_mla_backend = None self.has_fp8_kv_cache = False + self.local_layer_idx = local_layer_idx class FakePhasedFmha(PhasedFmha): @@ -81,10 +82,28 @@ def prepare_workspace( workspace.resize_(self._workspace_size) def run_context(self, params: FmhaParams) -> None: - self._events.append(("run", self._name, FmhaPhase.CONTEXT, params.num_tokens)) + self._events.append( + ( + "run", + self._name, + FmhaPhase.CONTEXT, + params.num_tokens, + params.batch_size, + params.num_requests, + ) + ) def run_generation(self, params: FmhaParams) -> None: - self._events.append(("run", self._name, FmhaPhase.GENERATION, params.num_tokens)) + self._events.append( + ( + "run", + self._name, + FmhaPhase.GENERATION, + params.num_tokens, + params.batch_size, + params.num_requests, + ) + ) class FakeFmha(Fmha): diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index cec23ea0b17a..89ced8124d89 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -32,6 +32,7 @@ def test_msa_package_availability_installs_cutlass_46_compatibility_aliases(monk cute = ModuleType("cutlass.cute") cute.core = SimpleNamespace() + cute.ThrCopy = object() cute.ThrMma = object() cute.make_rmem_tensor = object() cutlass = ModuleType("cutlass") @@ -43,12 +44,33 @@ def test_msa_package_availability_installs_cutlass_46_compatibility_aliases(monk msa_package_available.cache_clear() try: assert msa_package_available() + assert cute.core.ThrCopy is cute.ThrCopy assert cute.core.ThrMma is cute.ThrMma assert cute.make_fragment is cute.make_rmem_tensor finally: msa_package_available.cache_clear() +def test_msa_import_preserves_cute_compile_option_selection() -> None: + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import ( + msa_package_available, + ) + + if not msa_package_available(): + pytest.skip("fmha_sm100 (MSA) not importable") + + import cutlass.cute as cute + + original_compile = cute.compile + from fmha_sm100.cute import interface as sparse_interface + + del sparse_interface + assert cute.compile is original_compile + assert callable(cute.compile) + selected_compile = cute.compile[cute.FrontendNext] + assert callable(selected_compile) + + def test_resolver_selects_msa_backend_when_available(monkeypatch): import tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_availability as avail diff --git a/tests/unittest/_torch/attention/test_combined_fmha.py b/tests/unittest/_torch/attention/test_combined_fmha.py index e70763b7d6fc..65c20fb60ed2 100644 --- a/tests/unittest/_torch/attention/test_combined_fmha.py +++ b/tests/unittest/_torch/attention/test_combined_fmha.py @@ -29,6 +29,7 @@ AttentionForwardArgs, AttentionInputType, ) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm.bindings import DataType from tensorrt_llm.quantization.mode import QuantMode @@ -55,14 +56,14 @@ def test_combined_fmha_delegates_phases_and_prepares_max_workspace() -> None: metadata = SimpleNamespace( kv_cache_block_offsets=object(), effective_workspace=torch.empty(0, dtype=torch.uint8), - num_contexts=1, - num_ctx_tokens=2, - num_generations=1, - kv_lens_cuda_runtime=torch.tensor([2, 5], dtype=torch.int32), - kv_lens_runtime=torch.tensor([2, 5], dtype=torch.int32), - prompt_lens_cuda_runtime=torch.tensor([2, 1], dtype=torch.int32), - prompt_lens_cpu_runtime=torch.tensor([2, 1], dtype=torch.int32), - beam_width=1, + num_contexts=2, + num_ctx_tokens=3, + num_generations=4, + kv_lens_cuda_runtime=torch.tensor([2, 1, 5, 5, 5, 5], dtype=torch.int32), + kv_lens_runtime=torch.tensor([2, 1, 5, 5, 5, 5], dtype=torch.int32), + prompt_lens_cuda_runtime=torch.tensor([2, 1, 1, 1, 1, 1], dtype=torch.int32), + prompt_lens_cpu_runtime=torch.tensor([2, 1, 1, 1, 1, 1], dtype=torch.int32), + beam_width=2, cache_indirection=None, tokens_per_block=32, kv_cache_manager=None, @@ -70,35 +71,39 @@ def test_combined_fmha_delegates_phases_and_prepares_max_workspace() -> None: is_spec_decoding_enabled=False, ) forward_args = AttentionForwardArgs( - output=torch.empty((3, 4)), + output=torch.empty((7, 4)), attention_input_type=AttentionInputType.mixed, attention_window_size=8, ) - combined_fmha.forward(torch.empty((3, 4)), None, None, metadata, forward_args) + combined_fmha.forward(torch.empty((7, 4)), None, None, metadata, forward_args) assert events == [ ("prepare", "context"), ("prepare", "generation"), - ("run", "context", FmhaPhase.CONTEXT, 2), - ("run", "generation", FmhaPhase.GENERATION, 1), + ("run", "context", FmhaPhase.CONTEXT, 3, 2, 2), + ("run", "generation", FmhaPhase.GENERATION, 4, 4, 2), ] assert metadata.effective_workspace.numel() == 8 def test_combined_fmha_uses_flattened_v2_page_bound() -> None: - attn = FakeAttention() + attn = FakeAttention(local_layer_idx=3) combined_fmha = CombinedFmha(attn) - kv_cache_manager = SimpleNamespace( - impl=SimpleNamespace(get_page_index_upper_bound=lambda: 23), - blocks_in_primary_pool=23, - num_local_layers=4, - ) + calls: list[tuple[int, object]] = [] + + def get_page_index_upper_bound(local_layer_idx: int, role: object) -> int: + calls.append((local_layer_idx, role)) + return 23 + + kv_cache_manager = object.__new__(KVCacheManagerV2) + kv_cache_manager.impl = SimpleNamespace(get_page_index_upper_bound=get_page_index_upper_bound) assert ( combined_fmha._get_total_num_blocks(SimpleNamespace(kv_cache_manager=kv_cache_manager)) == 23 ) + assert calls == [(3, Role.KEY)] def test_flashinfer_fp8_mode_remains_implementation_local() -> None: @@ -145,6 +150,7 @@ def test_flashinfer_context_fallback_scope( ) fmha = object.__new__(FlashInferTrtllmGenFmha) fmha.kv_factor = 2 + monkeypatch.setattr(fmha, "_get_total_num_blocks", lambda _: 0) attn = FakeAttention() attn.sparse_params = None attn.position_embedding_type = 0 @@ -230,6 +236,7 @@ def test_flashinfer_quantized_kv_context_avoids_fp16_bf16_fallback( ) fmha = object.__new__(FlashInferTrtllmGenFmha) fmha.kv_factor = 2 + monkeypatch.setattr(fmha, "_get_total_num_blocks", lambda _: 0) attn = SimpleNamespace( is_mla_enable=False, sparse_params=None, diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 51d97fb4cfe8..763345829fed 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -11,18 +11,24 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.fmha import ( + flashinfer_trtllm_gen as flashinfer_trtllm_gen_module, +) from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import CuteDslMlaFmha from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( FlashInferTrtllmGenFmha, _get_multi_ctas_kv_counter_size, ) from tensorrt_llm._torch.attention.backends.fmha.interface import _CuteDslMlaStagingKey +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, ) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.autotuner import AutoTuner +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager class _AttentionStub: @@ -32,6 +38,7 @@ def __init__( is_mla_enable: bool, has_fp8_kv_cache: bool, flashinfer_mla_backend: str | None = None, + local_layer_idx: int = 0, ) -> None: self.is_mla_enable = is_mla_enable self.has_fp8_kv_cache = has_fp8_kv_cache @@ -39,33 +46,82 @@ def __init__( self.kv_lora_rank = 512 if is_mla_enable else None self.head_dim = 576 self.v_head_dim = 512 if is_mla_enable else None + self.local_layer_idx = local_layer_idx _MlaBackendPolicy: TypeAlias = Callable[[str, SimpleNamespace, int], str] -def _get_total_num_blocks(manager: SimpleNamespace, kv_factor: int = 2) -> int: - fmha = object.__new__(FlashInferTrtllmGenFmha) - fmha.kv_factor = kv_factor - return fmha._get_total_num_blocks(SimpleNamespace(kv_cache_manager=manager)) +def test_flashinfer_uses_v2_page_index_upper_bound_directly() -> None: + calls: list[tuple[int, object]] = [] + bounds = iter((97, 101)) + def get_page_index_upper_bound(local_layer_idx: int, role: object) -> int: + calls.append((local_layer_idx, role)) + return next(bounds) -def test_flashinfer_uses_v2_page_index_upper_bound_directly() -> None: - manager = SimpleNamespace( - blocks_in_primary_pool=50_000_000, - impl=SimpleNamespace(get_page_index_upper_bound=lambda *_: 50_000_000), - num_local_layers=36, + manager = object.__new__(KVCacheManagerV2) + manager.impl = SimpleNamespace(get_page_index_upper_bound=get_page_index_upper_bound) + fmha = object.__new__(FlashInferTrtllmGenFmha) + fmha.kv_factor = 2 + fmha._v1_total_num_blocks_cache = None + attn = SimpleNamespace(local_layer_idx=7) + fmha._attn_ref = lambda: attn + metadata = SimpleNamespace( + kv_cache_manager=manager, + host_kv_cache_pool_mapping=None, ) - assert _get_total_num_blocks(manager) == 50_000_000 + + assert fmha._get_total_num_blocks(metadata) == 97 + assert fmha._get_total_num_blocks(metadata) == 101 + assert calls == [(7, Role.KEY), (7, Role.KEY)] -def test_flashinfer_preserves_legacy_pool_scaling() -> None: - manager = SimpleNamespace( - blocks_in_primary_pool=1024, - impl=SimpleNamespace(), - num_local_layers=36, +@pytest.mark.parametrize("kv_factor", [1, 2]) +def test_flashinfer_uses_remaining_v1_selected_pool_extent(kv_factor: int) -> None: + calls: list[int] = [] + + def get_primary_pool_data(local_layer_idx: int) -> SimpleNamespace: + calls.append(local_layer_idx) + return SimpleNamespace(shape=(1024,)) + + pool_mapping = torch.tensor( + [ + [0, 0], + [1, 0], + [0, 1], + [1, 1], + [0, 2], + ], + dtype=torch.int32, ) - assert _get_total_num_blocks(manager, kv_factor=2) == 1024 * 36 * 2 + manager = object.__new__(KVCacheManager) + manager.impl = SimpleNamespace(get_primary_pool_data=get_primary_pool_data) + fmha = object.__new__(FlashInferTrtllmGenFmha) + fmha.kv_factor = kv_factor + fmha._v1_total_num_blocks_cache = None + attn = SimpleNamespace(local_layer_idx=4) + fmha._attn_ref = lambda: attn + metadata = SimpleNamespace( + kv_cache_manager=manager, + host_kv_cache_pool_mapping=pool_mapping, + ) + + expected = (1024 * 3 - 2) * kv_factor + assert fmha._get_total_num_blocks(metadata) == expected + assert fmha._get_total_num_blocks(metadata) == expected + assert calls == [4] + + +def test_phased_fmha_rejects_unknown_kv_cache_manager_type() -> None: + fmha = object.__new__(FlashInferTrtllmGenFmha) + fmha.kv_factor = 2 + fmha._v1_total_num_blocks_cache = None + fmha._attn_ref = lambda: SimpleNamespace(local_layer_idx=0) + metadata = SimpleNamespace(kv_cache_manager=SimpleNamespace()) + + with pytest.raises(TypeError, match="Unsupported KV cache manager: SimpleNamespace"): + fmha._get_total_num_blocks(metadata) def test_multi_ctas_kv_counter_size_covers_beam_expanded_batch() -> None: @@ -130,6 +186,106 @@ def check_counter_size_args( ) +def test_flashinfer_generation_uses_phase_batch_size_for_padded_cross_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cross-attention padding can make active rows irrecoverable from requests and beam width.""" + batch_size = 3 + preprocess_calls: list[tuple[object, ...]] = [] + + def generation_preprocess(*args: object) -> tuple[object, ...]: + preprocess_calls.append(args) + return ( + torch.empty((batch_size, 2, 4)), + torch.empty(1), + torch.empty((batch_size, 1), dtype=torch.int32), + None, + None, + None, + torch.empty(0, dtype=torch.uint8), + None, + 1, + 1, + -1, + False, + ) + + decode_calls: list[dict[str, object]] = [] + monkeypatch.setattr( + flashinfer_trtllm_gen_module.thop, + "trtllm_gen_generation_preprocess", + generation_preprocess, + ) + monkeypatch.setattr( + flashinfer_trtllm_gen_module, + "flashinfer", + SimpleNamespace( + decode=SimpleNamespace( + trtllm_batch_decode_with_kv_cache=lambda **kwargs: decode_calls.append(kwargs) + ) + ), + raising=False, + ) + + attn = SimpleNamespace( + local_layer_idx=0, + num_heads=2, + num_kv_heads=1, + head_dim=4, + q_scaling=1.0, + quant_mode=0, + predicted_tokens_per_seq=1, + attention_chunk_size=None, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + rope_params=SimpleNamespace(dim=0, theta=1.0, scale_type=0, scale=1.0, max_positions=1), + ) + metadata = SimpleNamespace( + beam_width=2, + kv_cache_block_offsets=torch.empty(0), + host_kv_cache_pool_pointers=torch.empty(0), + host_kv_cache_pool_mapping=torch.empty(0), + num_contexts=2, + ) + output = torch.empty((batch_size, 2, 4)) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.mixed, + ) + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=forward_args, + workspace=torch.empty(0, dtype=torch.uint8), + qkv_input=torch.empty((batch_size, 2, 4)), + context_buf=output, + sequence_lengths=torch.ones(batch_size, dtype=torch.int32), + input_seq_length=1, + num_tokens=batch_size, + seq_offset=2, + tokens_per_block=32, + kv_factor=2, + total_num_blocks=8, + batch_size=batch_size, + num_requests=1, + is_cross=True, + ) + fmha = SimpleNamespace( + _layout="HND", + _enable_pdl=False, + USE_SHARED_PAGED_KV_IDX=False, + _multi_processor_count=1, + _use_fp8_context_fmha=lambda _output, _input_type: False, + _get_multi_ctas_kv_counter_buffer=lambda: None, + ) + + FlashInferTrtllmGenFmha.run_generation(fmha, params) + + assert preprocess_calls[0][24] == batch_size + assert len(decode_calls) == 1 + + def test_flashinfer_cute_dsl_mla_backend_rejects_fp8_kv_cache() -> None: attn = _AttentionStub( is_mla_enable=True, diff --git a/tests/unittest/_torch/attention/test_fmha_registry.py b/tests/unittest/_torch/attention/test_fmha_registry.py new file mode 100644 index 000000000000..cd4cdb4fd7a2 --- /dev/null +++ b/tests/unittest/_torch/attention/test_fmha_registry.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from tensorrt_llm._torch.attention.backends.fmha import registry + +PRIMS_TS = "prims_ts" + + +def _canonical_names() -> tuple[str, ...]: + return tuple(registry.FMHA_LIBS) + + +def _enabled_names() -> tuple[str, ...]: + classes = registry.get_enabled_fmha_lib_classes() + names_by_class = {cls: name for name, cls in registry.FMHA_LIBS.items()} + return tuple(names_by_class[cls] for cls in classes) + + +def test_prims_ts_precedes_trtllm_gen_in_canonical_order() -> None: + names = _canonical_names() + assert names.index(PRIMS_TS) < names.index("flashinfer_trtllm_gen") + + +def test_default_fmha_libs_exclude_prims_ts(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + + assert PRIMS_TS not in registry.DEFAULT_FMHA_LIBS + assert set(registry.DEFAULT_FMHA_LIBS) <= set(registry.FMHA_LIBS) + assert _enabled_names() == registry.DEFAULT_FMHA_LIBS + + +@pytest.mark.parametrize("value", ["", " ", ", ,"]) +def test_empty_fmha_lib_env_uses_default( + monkeypatch: pytest.MonkeyPatch, + value: str, +) -> None: + monkeypatch.setenv("TLLM_FMHA_LIBS", value) + + assert _enabled_names() == registry.DEFAULT_FMHA_LIBS + + +def test_exact_fmha_lib_env_preserves_order_and_deduplicates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + other_name = next(name for name in reversed(_canonical_names()) if name != PRIMS_TS) + monkeypatch.setenv("TLLM_FMHA_LIBS", f" {other_name}, {PRIMS_TS}, {other_name} ") + + assert _enabled_names() == (other_name, PRIMS_TS) + + +def test_delta_fmha_lib_env_adds_prims_ts_in_canonical_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TLLM_FMHA_LIBS", f"+{PRIMS_TS}") + + expected_names = set(registry.DEFAULT_FMHA_LIBS) | {PRIMS_TS} + assert _enabled_names() == tuple(name for name in _canonical_names() if name in expected_names) + + +def test_delta_fmha_lib_env_removes_default_library(monkeypatch: pytest.MonkeyPatch) -> None: + removed_name = registry.DEFAULT_FMHA_LIBS[-1] + monkeypatch.setenv("TLLM_FMHA_LIBS", f"-{removed_name}") + + expected_names = set(registry.DEFAULT_FMHA_LIBS) - {removed_name} + assert _enabled_names() == tuple(name for name in _canonical_names() if name in expected_names) + + +def test_mixed_exact_and_delta_fmha_lib_env_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("TLLM_FMHA_LIBS", f"{PRIMS_TS},-{PRIMS_TS}") + + with pytest.raises(ValueError, match="either an exact comma-separated list"): + registry.get_enabled_fmha_lib_classes() + + +def test_unknown_fmha_lib_env_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None: + unknown_name = "unknown" + while unknown_name in registry.FMHA_LIBS: + unknown_name += "_" + monkeypatch.setenv("TLLM_FMHA_LIBS", f"{PRIMS_TS},{unknown_name}") + + with pytest.raises(ValueError, match=f"Unknown FMHA library '{unknown_name}'"): + registry.get_enabled_fmha_lib_classes() + + +def test_empty_delta_fmha_lib_env_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("TLLM_FMHA_LIBS", "+") + + with pytest.raises(ValueError, match="Invalid empty FMHA library entry"): + registry.get_enabled_fmha_lib_classes() diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py new file mode 100644 index 000000000000..5be953227100 --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -0,0 +1,1115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import functools +import inspect + +import pytest +import torch +from backend_case import BackendCase, generate_inputs, run_backend, run_case +from utils.util import isSM100Family + +pytestmark = pytest.mark.skipif( + not isSM100Family(), + reason="PrimsTS attention kernels require SM100 or SM103", +) + + +_QWEN2_7B = { + "num_heads": 28, + "num_kv_heads": 4, + "head_dim": 128, + "dtype": "bfloat16", + "kv_layout": "HND", + "page_size": 32, + "rope": { + "dim": 128, + "theta": 1_000_000.0, + "max_positions": 8192, + "is_neox": True, + }, + "fused_rope": True, +} + +_DEEPSEEK_V3_LITE_MLA = { + "num_heads": 32, + "num_kv_heads": 1, + "head_dim": 192, + "dtype": "bfloat16", + "kv_layout": "HND", + "page_size": 32, + "is_mla": True, + "kv_lora_rank": 512, + "q_lora_rank": 1536, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, +} + + +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +@pytest.mark.parametrize( + "phase_args", + [ + pytest.param( + { + "seq_lens": [65, 37], + "num_cached_tokens": [0, 0], + "num_contexts": 2, + }, + id="context", + ), + pytest.param( + { + "seq_lens": [1, 1], + "num_cached_tokens": [64, 96], + "num_contexts": 0, + }, + id="generation", + ), + pytest.param( + { + "seq_lens": [41, 1], + "num_cached_tokens": [0, 63], + "num_contexts": 1, + }, + id="mixed", + ), + ], +) +def test_prims_ts_qwen2_gqa( + monkeypatch: pytest.MonkeyPatch, + use_kv_cache_manager_v2: bool, + phase_args: dict, +) -> None: + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + case = BackendCase( + **_QWEN2_7B, + **phase_args, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ) + + run_case(case) + + +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_prims_ts_context_zero_fills_nan_v_tail( + monkeypatch: pytest.MonkeyPatch, + use_kv_cache_manager_v2: bool, +) -> None: + from tensorrt_llm._torch.attention.backends.prims_ts.context import BatchPrefillPagedTSWrapper + + original_run = BatchPrefillPagedTSWrapper.run + poisoned_tails: list[tuple[int, int]] = [] + + def run_with_nan_v_tail( + self: BatchPrefillPagedTSWrapper, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + qo_indptr: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_kv: torch.Tensor, + **kwargs: object, + ) -> torch.Tensor: + # Poison after TRT-LLM writes valid KV and immediately before PrimTS + # launches, so only the kernel's handling of unused V rows is tested. + page_size = v_cache.shape[2] + seq_lens = seq_lens_kv.cpu().tolist() + page_indices = block_tables.cpu() + for batch_idx, seq_len in enumerate(seq_lens): + logical_last_page = (seq_len - 1) // page_size + tail_start = (seq_len - 1) % page_size + 1 + if tail_start == page_size: + continue + physical_page = int(page_indices[batch_idx, logical_last_page].item()) + v_cache[physical_page, :, tail_start:, :].fill_(float("nan")) + poisoned_tails.append((batch_idx, tail_start)) + + return original_run( + self, + q, + k_cache, + v_cache, + qo_indptr, + block_tables, + seq_lens_kv, + **kwargs, + ) + + monkeypatch.setattr(BatchPrefillPagedTSWrapper, "run", run_with_nan_v_tail) + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + case = BackendCase( + **_QWEN2_7B, + seq_lens=[65, 37], + num_cached_tokens=[0, 0], + num_contexts=2, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ) + + results = run_case(case) + + assert "TRTLLM" in results + assert poisoned_tails == [(0, 1), (1, 5)] + + +def test_prims_ts_fp16_dense_context_with_alternate_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + case = BackendCase( + num_heads=8, + num_kv_heads=2, + head_dim=256, + seq_lens=[67, 33], + num_cached_tokens=[0, 0], + num_contexts=2, + dtype="float16", + causal=False, + kv_layout="HND", + page_size=64, + use_kv_cache_manager_v2=True, + ) + + results = run_case(case) + + assert "TRTLLM" in results + + +def test_prims_ts_uses_compact_preprocessing_and_separate_decode_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import tensorrt_llm._torch.attention.backends.fmha.prims_ts as prims_ts_module + from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha + from tensorrt_llm._torch.attention.backends.prims_ts import ( + get_prims_ts_batch_decode_workspace_size, + ) + + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + case = BackendCase( + **_QWEN2_7B, + seq_lens=[1, 1], + num_cached_tokens=[64, 96], + num_contexts=0, + use_kv_cache_manager_v2=True, + ) + inputs = generate_inputs(case, seed=0) + golden = run_backend( + case, + "VANILLA", + inputs, + kv_dtype=case.compute_dtype, + kv_layout="NHD", + ) + + max_kv_len = 128 + prims_workspace_bytes = get_prims_ts_batch_decode_workspace_size( + case.num_seqs, + case.num_heads, + case.num_kv_heads, + case.head_dim, + case.page_size, + max_kv_len, + seq_len_q=1, + q_dtype=case.compute_dtype, + kv_dtype=case.compute_dtype, + out_dtype=case.compute_dtype, + mask_type="causal", + window_left=-1, + device=torch.device("cuda"), + ) + assert prims_workspace_bytes > 0 + + original_generation_layout = prims_ts_module.thop.get_trtllm_gen_generation_workspace_layout + generation_layout_records = [] + + def record_compact_layout(*args, **kwargs): + layout = original_generation_layout(*args, **kwargs) + generation_layout_records.append((args, kwargs, dict(layout))) + return layout + + monkeypatch.setattr( + prims_ts_module.thop, + "get_trtllm_gen_generation_workspace_layout", + record_compact_layout, + ) + + original_get_decode_workspace = PrimsTSFmha._get_decode_workspace + workspace_records = [] + + def record_decode_workspace( + self: PrimsTSFmha, + root_workspace: torch.Tensor, + ) -> torch.Tensor: + decode_workspace = original_get_decode_workspace(self, root_workspace) + workspace_records.append( + ( + self, + root_workspace.data_ptr(), + root_workspace.numel() * root_workspace.element_size(), + decode_workspace, + ) + ) + return decode_workspace + + monkeypatch.setattr(PrimsTSFmha, "_get_decode_workspace", record_decode_workspace) + + eager = run_backend( + case, + "TRTLLM", + inputs, + kv_dtype=case.compute_dtype, + fuse_rope=True, + kv_layout="HND", + ) + actual = run_backend( + case, + "TRTLLM", + inputs, + kv_dtype=case.compute_dtype, + fuse_rope=True, + cuda_graph=True, + kv_layout="HND", + ) + + torch.testing.assert_close(eager, golden, atol=3e-2, rtol=3e-3) + torch.testing.assert_close(actual, golden, atol=3e-2, rtol=3e-3) + + assert generation_layout_records + compact_preprocess_bytes = int(generation_layout_records[0][2]["total_size"]) + assert all( + kwargs["skip_fmha_workspace"] is True + for _args, kwargs, _layout in generation_layout_records + ) + assert all( + int(layout["trtllm_gen_workspace_size"]) == 0 + for _args, _kwargs, layout in generation_layout_records + ) + + records_by_adapter = {} + for adapter, root_ptr, root_bytes, decode_workspace in workspace_records: + byte_offset = adapter._decode_workspace_offset_bytes + required_bytes = adapter._decode_workspace_required_bytes + assert byte_offset is not None + assert byte_offset % 32 == 0 + assert byte_offset >= compact_preprocess_bytes + assert required_bytes == prims_workspace_bytes + assert root_bytes >= byte_offset + required_bytes + assert decode_workspace.data_ptr() == root_ptr + byte_offset + assert decode_workspace.numel() * decode_workspace.element_size() == required_bytes + wrapper = adapter._decode_wrappers[case.num_seqs] + plan_state = wrapper._plan_state + assert plan_state is not None + assert plan_state.workspace_buffer.data_ptr() == decode_workspace.data_ptr() + records_by_adapter.setdefault(adapter, []).append((root_ptr, decode_workspace.data_ptr())) + + assert len(records_by_adapter) == 2 + for records in records_by_adapter.values(): + assert len({root_ptr for root_ptr, _ in records}) == 1 + assert len({decode_ptr for _, decode_ptr in records}) == 1 + + captured_adapter, _, _, captured_workspace = workspace_records[-1] + captured_wrapper = captured_adapter._decode_wrappers[case.num_seqs] + captured_plan_state = captured_wrapper._plan_state + assert captured_plan_state is not None + assert torch.count_nonzero(captured_plan_state.workspace.split_kv_counter) == 0 + + +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_prims_ts_deepseek_v3_lite_mla_generation( + monkeypatch: pytest.MonkeyPatch, + use_kv_cache_manager_v2: bool, +) -> None: + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + case = BackendCase( + **_DEEPSEEK_V3_LITE_MLA, + seq_lens=[1, 1], + num_cached_tokens=[64, 96], + num_contexts=0, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ) + + run_case(case) + + +def test_prims_ts_context_wrapper_cuda_graph_replay_with_updated_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import tensorrt_llm._torch.attention.backends.prims_ts.context as context_module + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BatchPrefillPagedTSWrapper, + batch_prefill_with_paged_kv_cache, + ) + + # Isolate this regression from compile-cache entries created by earlier + # tests while retaining each compiled module through the full A/E/F + # sequence. The former ragged V tensor map aborted on SM100 only after a + # CLC D128 -> nonpersistent D256 -> distinct CLC D128 compile/run order. + uncached_compile = inspect.unwrap(context_module._get_compiled_paged_context) + compile_records = [] + + @functools.cache + def compile_with_record(*args): + result = uncached_compile(*args) + compiled, policy = result + compile_records.append((args, dict(policy), compiled)) + return result + + monkeypatch.setattr(context_module, "_get_compiled_paged_context", compile_with_record) + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + + a_results = run_case( + BackendCase( + **_QWEN2_7B, + seq_lens=[65, 37], + num_cached_tokens=[0, 0], + num_contexts=2, + use_kv_cache_manager_v2=False, + ) + ) + assert "TRTLLM" in a_results + + e_results = run_case( + BackendCase( + num_heads=8, + num_kv_heads=2, + head_dim=256, + seq_lens=[67, 33], + num_cached_tokens=[0, 0], + num_contexts=2, + dtype="float16", + causal=False, + kv_layout="HND", + page_size=64, + use_kv_cache_manager_v2=True, + ) + ) + assert "TRTLLM" in e_results + + batch_size = 2 + num_qo_heads = 8 + num_kv_heads = 2 + head_dim = 128 + page_size = 32 + max_seq_len = 64 + dtype = torch.bfloat16 + device = torch.device("cuda") + + query = torch.randn(5, num_qo_heads, head_dim, device=device, dtype=dtype) + k_cache = torch.randn( + 4, + num_kv_heads, + page_size, + head_dim, + device=device, + dtype=dtype, + ) + v_cache = torch.randn_like(k_cache) + qo_indptr = torch.tensor([0, 3, 5], device=device, dtype=torch.int32) + trt_block_tables = torch.tensor( + [[[0, 1], [2, 3]], [[2, 3], [0, 1]]], + device=device, + dtype=torch.int32, + ) + block_tables = trt_block_tables[:, 0, :] + assert block_tables.stride() == (4, 1) + seq_lens = torch.tensor([33, 64], device=device, dtype=torch.int32) + output = torch.empty_like(query) + wrapper = BatchPrefillPagedTSWrapper(kv_layout="HND") + wrapper.plan( + device=device, + batch_size=batch_size, + max_seq_len_q=3, + max_kv_len=max_seq_len, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, + page_size=page_size, + mask_type="causal", + ) + plan_state = wrapper._plan_state + assert plan_state is not None + compiled = plan_state.compiled + + assert len(compile_records) == 3 + assert [ + (args[0].max_seq_len_q, args[0].max_kv_len, policy["scheduler"]) + for args, policy, _compiled in compile_records + ] == [ + (96, 96, "clc_dynamic_persistent"), + (128, 128, "nonpersistent"), + (3, 64, "clc_dynamic_persistent"), + ] + assert compile_records[0][0] != compile_records[2][0] + assert len({id(recorded) for _args, _policy, recorded in compile_records}) == 3 + + wrapper.run( + query, + k_cache, + v_cache, + qo_indptr, + block_tables, + seq_lens, + out=output, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = wrapper.run( + query, + k_cache, + v_cache, + qo_indptr, + block_tables, + seq_lens, + out=output, + validate=False, + ) + + query.copy_(query.flip(0).clone()) + qo_indptr.copy_(torch.tensor([0, 2, 5], device=device, dtype=torch.int32)) + seq_lens.copy_(torch.tensor([64, 33], device=device, dtype=torch.int32)) + block_tables.copy_(torch.tensor([[2, 3], [0, 1]], device=device, dtype=torch.int32)) + + reference = batch_prefill_with_paged_kv_cache( + query, + k_cache, + v_cache, + qo_indptr, + block_tables, + seq_lens, + page_size=page_size, + mask_type="causal", + out_dtype=dtype, + ) + graph.replay() + torch.cuda.synchronize() + + assert graph_output.data_ptr() == output.data_ptr() + assert wrapper._plan_state is plan_state + assert wrapper._plan_state.compiled is compiled + torch.testing.assert_close(graph_output, reference, atol=3e-2, rtol=3e-3) + + +def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BatchDecodePagedTSWrapper, + get_prims_ts_batch_decode_workspace_size, + prims_ts_batch_decode_with_kv_cache, + ) + + batch_size = 2 + num_qo_heads = 8 + num_kv_heads = 2 + head_dim = 128 + page_size = 32 + max_seq_len = 64 + dtype = torch.bfloat16 + device = torch.device("cuda") + + query = torch.randn( + batch_size, + num_qo_heads, + head_dim, + device=device, + dtype=dtype, + ) + kv_cache = torch.randn( + 4, + 2, + num_kv_heads, + page_size, + head_dim, + device=device, + dtype=dtype, + ) + trt_block_tables = torch.tensor( + [[[0, 1], [2, 3]], [[2, 3], [0, 1]]], + device=device, + dtype=torch.int32, + ) + block_tables = trt_block_tables[:, 0, :] + assert block_tables.stride() == (4, 1) + seq_lens = torch.tensor([33, 64], device=device, dtype=torch.int32) + output = torch.empty_like(query) + workspace_bytes = get_prims_ts_batch_decode_workspace_size( + batch_size, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + max_seq_len, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, + mask_type="causal", + device=device, + ) + external_workspace = torch.zeros( + max(workspace_bytes, query.numel() * query.element_size()), + device=device, + dtype=torch.uint8, + ) + wrapper = BatchDecodePagedTSWrapper(kv_layout="HND") + wrapper.plan( + device, + batch_size, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + max_seq_len, + max_seq_len_q=1, + packed_query=False, + q_data_type=dtype, + kv_data_type=dtype, + o_data_type=dtype, + mask_type="causal", + workspace_buffer=external_workspace, + ) + plan_state = wrapper._plan_state + assert plan_state is not None + compiled_main = plan_state.compiled_main + + aliased_query = ( + external_workspace[: query.numel() * query.element_size()].view(dtype).view_as(query) + ) + with pytest.raises(ValueError, match="workspace_buffer must not overlap query storage"): + wrapper.run( + aliased_query, + kv_cache, + seq_lens, + block_tables, + out=output, + ) + + external_workspace.zero_() + wrapper.run( + query, + kv_cache, + seq_lens, + block_tables, + out=output, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + external_workspace.zero_() + graph_output = wrapper.run( + query, + kv_cache, + seq_lens, + block_tables, + out=output, + validate=False, + ) + + query.copy_(query.flip(0).clone()) + block_tables.copy_(torch.tensor([[2, 3], [0, 1]], device=device, dtype=torch.int32)) + seq_lens.copy_(torch.tensor([32, 64], device=device, dtype=torch.int32)) + reference_workspace = torch.zeros_like(external_workspace) + reference = prims_ts_batch_decode_with_kv_cache( + query, + kv_cache, + reference_workspace, + block_tables, + seq_lens, + max_seq_len, + out_dtype=dtype, + out=torch.empty_like(query), + mask_type="causal", + kv_layout="HND", + ) + graph.replay() + torch.cuda.synchronize() + + assert graph_output.data_ptr() == output.data_ptr() + assert wrapper._plan_state is plan_state + assert plan_state.workspace_buffer is external_workspace + assert plan_state.compiled_main is compiled_main + assert plan_state.kv_prefix_mode == "dynamic" + assert plan_state.kv_lengths_mode == "dynamic" + torch.testing.assert_close(graph_output, reference, atol=3e-2, rtol=3e-3) + + +def test_prims_ts_mla_live_wrapper_cuda_graph_replay() -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BatchMLADecodePagedTSWrapper, + get_prims_ts_batch_mla_decode_workspace_size, + prims_ts_batch_mla_decode_with_kv_cache, + ) + + batch_size = 2 + num_heads = 32 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + page_size = 32 + max_seq_len = 1024 + dtype = torch.bfloat16 + device = torch.device("cuda") + + query = torch.randn( + batch_size, + 1, + num_heads, + kv_lora_rank + qk_rope_head_dim, + device=device, + dtype=dtype, + ) + kv_cache = torch.randn( + 4, + page_size, + kv_lora_rank + qk_rope_head_dim, + device=device, + dtype=dtype, + ) + block_tables = torch.zeros( + batch_size, + max_seq_len // page_size, + device=device, + dtype=torch.int32, + ) + block_tables[0, :2] = torch.tensor([0, 1], device=device, dtype=torch.int32) + block_tables[1, :2] = torch.tensor([2, 3], device=device, dtype=torch.int32) + seq_lens = torch.tensor([33, 64], device=device, dtype=torch.int32) + output = torch.empty( + batch_size, + 1, + num_heads, + kv_lora_rank, + device=device, + dtype=dtype, + ) + workspace_bytes = get_prims_ts_batch_mla_decode_workspace_size( + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + page_size, + max_seq_len, + max_seq_len_q=1, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, + mask_type="causal", + device=device, + ) + external_workspace = torch.empty( + max(workspace_bytes, query.numel() * query.element_size()), + device=device, + dtype=torch.uint8, + ) + wrapper = BatchMLADecodePagedTSWrapper() + wrapper.plan( + device, + batch_size, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + page_size, + max_seq_len, + max_seq_len_q=1, + packed_query=False, + q_data_type=dtype, + kv_data_type=dtype, + o_data_type=dtype, + mask_type="causal", + workspace_buffer=external_workspace, + ) + plan_state = wrapper._plan_state + assert plan_state is not None + compiled = plan_state.compiled + bmm1_scale = (128 + qk_rope_head_dim) ** -0.5 + + aliased_query = ( + external_workspace[: query.numel() * query.element_size()].view(dtype).view_as(query) + ) + with pytest.raises(ValueError, match="workspace_buffer must not overlap query storage"): + wrapper.run( + aliased_query, + kv_cache, + block_tables, + seq_lens, + bmm1_scale=bmm1_scale, + out=output, + ) + + wrapper.run( + query, + kv_cache, + block_tables, + seq_lens, + bmm1_scale=bmm1_scale, + out=output, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = wrapper.run( + query, + kv_cache, + block_tables, + seq_lens, + bmm1_scale=bmm1_scale, + out=output, + validate=False, + ) + + query.copy_(query.flip(0).clone()) + block_tables.copy_(block_tables.flip(0).clone()) + seq_lens.copy_(seq_lens.flip(0).clone()) + reference_workspace = torch.empty_like(external_workspace) + reference = prims_ts_batch_mla_decode_with_kv_cache( + query, + kv_cache, + reference_workspace, + kv_lora_rank, + qk_rope_head_dim, + block_tables, + seq_lens, + max_seq_len, + max_seq_len_q=1, + bmm1_scale=bmm1_scale, + out_dtype=dtype, + out=torch.empty_like(output), + mask_type="causal", + ) + graph.replay() + torch.cuda.synchronize() + + assert graph_output.data_ptr() == output.data_ptr() + assert wrapper._plan_state is plan_state + assert plan_state.workspace_buffer is external_workspace + assert plan_state.compiled is compiled + torch.testing.assert_close(graph_output, reference, atol=3e-2, rtol=3e-3) + + +def test_prims_ts_decode_graph_profiles_reset_shared_workspace_a_b_a() -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BatchDecodePagedTSWrapper, + get_prims_ts_batch_decode_workspace_size, + prims_ts_batch_decode_with_kv_cache, + ) + + num_qo_heads = 8 + num_kv_heads = 2 + head_dim = 128 + page_size = 32 + max_seq_len_a = 64 + max_seq_len_b = 4096 + dtype = torch.bfloat16 + device = torch.device("cuda") + kv_cache = torch.randn( + 256, + 2, + num_kv_heads, + page_size, + head_dim, + device=device, + dtype=dtype, + ) + workspace_bytes = max( + get_prims_ts_batch_decode_workspace_size( + batch_size, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + max_seq_len, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, + mask_type="causal", + device=device, + ) + for batch_size, max_seq_len in ( + (1, max_seq_len_a), + (2, max_seq_len_b), + ) + ) + shared_workspace = torch.zeros( + workspace_bytes, + device=device, + dtype=torch.uint8, + ) + wrapper = BatchDecodePagedTSWrapper(kv_layout="HND") + query_a = torch.randn(1, num_qo_heads, head_dim, device=device, dtype=dtype) + query_b = torch.randn(2, num_qo_heads, head_dim, device=device, dtype=dtype) + block_tables_a = torch.tensor([[0, 1]], device=device, dtype=torch.int32) + block_tables_b = torch.arange(256, device=device, dtype=torch.int32).view(2, 128) + seq_lens_a = torch.tensor([33], device=device, dtype=torch.int32) + seq_lens_b = torch.tensor([2049, 4096], device=device, dtype=torch.int32) + output_a = torch.empty_like(query_a) + output_b = torch.empty_like(query_b) + + def plan(batch_size: int, max_seq_len: int) -> None: + wrapper.plan( + device, + batch_size, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + max_seq_len, + max_seq_len_q=1, + packed_query=False, + q_data_type=dtype, + kv_data_type=dtype, + o_data_type=dtype, + mask_type="causal", + workspace_buffer=shared_workspace, + ) + + def capture( + query: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + output: torch.Tensor, + ) -> torch.cuda.CUDAGraph: + plan_state = wrapper._plan_state + assert plan_state is not None + control_offset = plan_state.workspace_layout.split_kv_counter.byte_offset + control_end = plan_state.workspace_layout.total_bytes + shared_workspace[control_offset:control_end].zero_() + wrapper.run( + query, + kv_cache, + seq_lens, + block_tables, + out=output, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + shared_workspace[control_offset:control_end].zero_() + wrapper.run( + query, + kv_cache, + seq_lens, + block_tables, + out=output, + validate=False, + ) + return graph + + plan(1, max_seq_len_a) + plan_state_a = wrapper._plan_state + assert plan_state_a is not None + compiled_a = plan_state_a.compiled_main + layout_a = plan_state_a.workspace_layout + control_span_a = slice(layout_a.split_kv_counter.byte_offset, layout_a.total_bytes) + graph_a = capture(query_a, block_tables_a, seq_lens_a, output_a) + plan(2, max_seq_len_b) + plan_state_b = wrapper._plan_state + assert plan_state_b is not None + layout_b = plan_state_b.workspace_layout + assert layout_a.split_kv_counter.byte_offset != layout_b.split_kv_counter.byte_offset + control_span_b = slice(layout_b.split_kv_counter.byte_offset, layout_b.total_bytes) + graph_b = capture(query_b, block_tables_b, seq_lens_b, output_b) + plan(1, max_seq_len_a) + + reference_workspace = torch.zeros_like(shared_workspace) + reference_a = prims_ts_batch_decode_with_kv_cache( + query_a, + kv_cache, + reference_workspace, + block_tables_a, + seq_lens_a, + max_seq_len_a, + out_dtype=dtype, + out=torch.empty_like(output_a), + mask_type="causal", + kv_layout="HND", + ).clone() + reference_workspace.zero_() + reference_b = prims_ts_batch_decode_with_kv_cache( + query_b, + kv_cache, + reference_workspace, + block_tables_b, + seq_lens_b, + max_seq_len_b, + out_dtype=dtype, + out=torch.empty_like(output_b), + mask_type="causal", + kv_layout="HND", + ).clone() + + actual_a = [] + shared_workspace[control_span_a].fill_(0xFF) + graph_a.replay() + torch.cuda.synchronize() + assert torch.count_nonzero(shared_workspace[control_span_a]) == 0 + actual_a.append(output_a.clone()) + shared_workspace[control_span_b].fill_(0xFF) + graph_b.replay() + torch.cuda.synchronize() + assert torch.count_nonzero(shared_workspace[control_span_b]) == 0 + actual_b = output_b.clone() + shared_workspace[control_span_a].fill_(0xFF) + graph_a.replay() + torch.cuda.synchronize() + assert torch.count_nonzero(shared_workspace[control_span_a]) == 0 + actual_a.append(output_a.clone()) + + final_plan_state = wrapper._plan_state + assert final_plan_state is not None + assert final_plan_state.workspace_buffer is shared_workspace + assert final_plan_state.compiled_main is compiled_a + for actual in actual_a: + torch.testing.assert_close(actual, reference_a, atol=3e-2, rtol=3e-3) + torch.testing.assert_close(actual_b, reference_b, atol=3e-2, rtol=3e-3) + + +def test_prims_ts_decode_wrappers_share_workspace_across_serialized_layers() -> None: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BatchDecodePagedTSWrapper, + get_prims_ts_batch_decode_workspace_size, + prims_ts_batch_decode_with_kv_cache, + ) + + batch_size = 2 + num_qo_heads = 8 + num_kv_heads = 2 + head_dim = 128 + page_size = 32 + max_seq_len = 64 + dtype = torch.bfloat16 + device = torch.device("cuda") + workspace_bytes = get_prims_ts_batch_decode_workspace_size( + batch_size, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + max_seq_len, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, + mask_type="causal", + device=device, + ) + shared_workspace = torch.zeros( + workspace_bytes, + device=device, + dtype=torch.uint8, + ) + trt_block_tables = torch.tensor( + [[[0, 1], [2, 3]], [[2, 3], [0, 1]]], + device=device, + dtype=torch.int32, + ) + block_tables = trt_block_tables[:, 0, :] + seq_lens = torch.tensor([33, 64], device=device, dtype=torch.int32) + wrappers = [BatchDecodePagedTSWrapper(kv_layout="HND") for _ in range(2)] + for wrapper in wrappers: + wrapper.plan( + device, + batch_size, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + max_seq_len, + max_seq_len_q=1, + packed_query=False, + q_data_type=dtype, + kv_data_type=dtype, + o_data_type=dtype, + mask_type="causal", + workspace_buffer=shared_workspace, + ) + + for layer_index, wrapper in enumerate(wrappers): + query = torch.randn( + batch_size, + num_qo_heads, + head_dim, + device=device, + dtype=dtype, + ) + kv_cache = torch.randn( + 4, + 2, + num_kv_heads, + page_size, + head_dim, + device=device, + dtype=dtype, + ) + output = torch.empty_like(query) + plan_state = wrapper._plan_state + assert plan_state is not None + control_offset = plan_state.workspace_layout.split_kv_counter.byte_offset + shared_workspace[control_offset : plan_state.workspace_layout.total_bytes].zero_() + actual = wrapper.run( + query, + kv_cache, + seq_lens, + block_tables, + out=output, + ) + reference_workspace = torch.zeros_like(shared_workspace) + reference = prims_ts_batch_decode_with_kv_cache( + query, + kv_cache, + reference_workspace, + block_tables, + seq_lens, + max_seq_len, + out_dtype=dtype, + out=torch.empty_like(query), + mask_type="causal", + kv_layout="HND", + ) + + assert plan_state.workspace_buffer.data_ptr() == shared_workspace.data_ptr(), layer_index + torch.testing.assert_close(actual, reference, atol=3e-2, rtol=3e-3) + + +def test_prims_ts_unsupported_context_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: + from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha + from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha + + calls = {"fallback": 0, "prims_context": 0} + fallback_forward = FallbackFmha.forward + prims_context = PrimsTSFmha.run_context + + def counted_fallback(self, *args, **kwargs): + calls["fallback"] += 1 + return fallback_forward(self, *args, **kwargs) + + def counted_prims_context(self, *args, **kwargs): + calls["prims_context"] += 1 + return prims_context(self, *args, **kwargs) + + monkeypatch.setattr(FallbackFmha, "forward", counted_fallback) + monkeypatch.setattr(PrimsTSFmha, "run_context", counted_prims_context) + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts,fallback") + case = BackendCase( + num_heads=14, + num_kv_heads=2, + head_dim=64, + seq_lens=[65, 37], + num_cached_tokens=[0, 0], + num_contexts=2, + dtype="bfloat16", + kv_layout="HND", + page_size=32, + ) + + run_case(case) + + assert calls["fallback"] > 0 + assert calls["prims_context"] == 0 diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py new file mode 100644 index 000000000000..dccc8ec04bde --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -0,0 +1,2098 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from packaging.version import Version + +import tensorrt_llm._torch.attention.backends.fmha.prims_ts as prims_ts_module +import tensorrt_llm._torch.attention.backends.fmha.utils as fmha_utils +import tensorrt_llm._torch.attention.backends.prims_ts as prims_ts_package +import tensorrt_llm._torch.attention.backends.prims_ts.context as prims_context_module +import tensorrt_llm._torch.attention.backends.prims_ts.decode as prims_decode_module +import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as prims_mla_module +from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha +from tensorrt_llm._torch.attention.backends.fmha.registry import get_enabled_fmha_lib_classes +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings import DataType + + +class _TensorSpec: + """Minimal tensor-like object for the pure support predicate.""" + + def __init__( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + *, + device: str = "cuda", + contiguous: bool = True, + ) -> None: + self.shape = shape + self.dtype = dtype + self.device = torch.device(device) + self.ndim = len(shape) + self._contiguous = contiguous + + def is_contiguous(self) -> bool: + return self._contiguous + + def numel(self) -> int: + return math.prod(self.shape) + + +class _Attention: + def __init__( + self, + *, + head_dim: int = 128, + is_mla: bool = False, + num_heads: int = 8, + num_kv_heads: int | None = None, + ) -> None: + self.num_heads = num_heads + self.num_kv_heads = (1 if is_mla else 2) if num_kv_heads is None else num_kv_heads + self.head_dim = head_dim + self.is_mla_enable = is_mla + self.kv_lora_rank = 512 if is_mla else None + self.qk_rope_head_dim = 64 if is_mla else None + self.qk_nope_head_dim = 128 if is_mla else None + self.v_head_dim = 128 if is_mla else None + self.predicted_tokens_per_seq = 1 + self.sparse_params = None + self.position_embedding_type = 0 + self.quant_mode = 0 + self.q_scaling = 1.0 + self.attention_chunk_size = 0 + self.rope_dim = head_dim + self.local_layer_idx = 0 + self.rope_params = SimpleNamespace( + dim=head_dim, + theta=10000.0, + scale_type=0, + scale=1.0, + max_positions=4096, + ) + self.rotary_inv_freq = None + self.rotary_cos_sin = None + + +def _make_v1_manager(**attributes: object) -> KVCacheManager: + manager = object.__new__(KVCacheManager) + for name, value in attributes.items(): + setattr(manager, name, value) + return manager + + +def _make_v2_manager(**attributes: object) -> KVCacheManagerV2: + manager = object.__new__(KVCacheManagerV2) + for name, value in attributes.items(): + setattr(manager, name, value) + return manager + + +def _support_result( + *, + attention_input_type: AttentionInputType, + head_dim: int = 128, + num_heads: int = 8, + num_kv_heads: int | None = None, + dtype: torch.dtype = torch.bfloat16, + output_dtype: torch.dtype | None = None, + kv_dtype: DataType | None = None, + tokens_per_block: int = 32, + is_mla: bool = False, + is_fused_qkv: bool = True, + has_separate_kv: bool = False, + has_paged_cache: bool = True, + is_cross: bool = False, + beam_width: int = 1, + use_spec_decoding: bool = False, + is_spec_dec_tree: bool = False, + has_attention_sinks: bool = False, + has_relative_attention_bias: bool = False, + has_sparse_attention: bool = False, + has_sparse_runtime_metadata: bool = False, + position_embedding_type: int = 0, + kv_lora_rank: int | None = None, + qk_rope_head_dim: int | None = None, + has_output: bool = True, + attention_window_size: int = 128, + attention_chunk_size: int = 0, + max_seq_len: int = 128, + kv_layout: str = "HND", + num_kv_cache_pools: int = 1, + use_kv_cache_v2: bool = False, + enable_swa_scratch_reuse: bool = False, + phase: FmhaPhase | None = None, +) -> tuple[bool, str]: + attn = _Attention( + head_dim=head_dim, + is_mla=is_mla, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + ) + attn.position_embedding_type = position_embedding_type + attn.attention_chunk_size = attention_chunk_size + if has_sparse_attention: + attn.sparse_params = SimpleNamespace(algorithm="mqa_gqa") + if kv_lora_rank is not None: + attn.kv_lora_rank = kv_lora_rank + if qk_rope_head_dim is not None: + attn.qk_rope_head_dim = qk_rope_head_dim + + output_dtype = dtype if output_dtype is None else output_dtype + if kv_dtype is None: + kv_dtype = DataType.BF16 if dtype == torch.bfloat16 else DataType.HALF + q_width = attn.num_heads * head_dim + if is_fused_qkv and not is_mla: + q_width += 2 * attn.num_kv_heads * head_dim + q = _TensorSpec((4, q_width), dtype) + output_width = attn.num_heads * (512 if is_mla else head_dim) + forward_args = AttentionForwardArgs( + output=_TensorSpec((4, output_width), output_dtype) if has_output else None, + attention_input_type=attention_input_type, + attention_mask=PredefinedAttentionMask.CAUSAL, + attention_window_size=attention_window_size, + attention_sinks=torch.empty(1) if has_attention_sinks else None, + relative_attention_bias=torch.empty(1) if has_relative_attention_bias else None, + is_fused_qkv=is_fused_qkv, + ) + if has_sparse_runtime_metadata: + forward_args.sparse_runtime_params.sparse_kv_indices = torch.empty(1) + if attention_input_type == AttentionInputType.context_only: + num_contexts, num_generations, num_ctx_tokens = 1, 0, 4 + kv_lens = [4] + elif attention_input_type == AttentionInputType.generation_only: + num_contexts, num_generations, num_ctx_tokens = 0, 4, 0 + kv_lens = [128, 96, 64, 32] + else: + num_contexts, num_generations, num_ctx_tokens = 1, 1, 3 + kv_lens = [3, 128] + if use_kv_cache_v2: + kv_cache_manager = _make_v2_manager( + dtype=kv_dtype, + impl=SimpleNamespace(get_page_index_upper_bound=lambda *args: 128), + enable_swa_scratch_reuse=enable_swa_scratch_reuse, + num_local_layers=1, + num_pools=num_kv_cache_pools, + kv_offset=torch.full((num_kv_cache_pools,), 128, dtype=torch.int32), + ) + else: + kv_cache_manager = _make_v1_manager( + dtype=kv_dtype, + impl=SimpleNamespace(), + num_local_layers=1, + num_pools=num_kv_cache_pools, + host_kv_cache_block_offsets=torch.tensor( + [[[[0], [128]]]], + dtype=torch.int32, + ), + ) + metadata = SimpleNamespace( + helix_position_offsets=None, + num_sparse_topk=0, + use_spec_decoding=use_spec_decoding, + is_spec_dec_tree=is_spec_dec_tree, + is_spec_dec_dynamic_tree=False, + is_spec_decoding_enabled=use_spec_decoding, + kv_cache_block_offsets=torch.empty(1) if has_paged_cache else None, + host_kv_cache_pool_pointers=torch.empty(1), + host_kv_cache_pool_mapping=torch.zeros((1, 2), dtype=torch.int32), + kv_cache_manager=kv_cache_manager, + is_cross=is_cross, + beam_width=beam_width, + tokens_per_block=tokens_per_block, + kv_layout=kv_layout, + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, + kv_lens_runtime=torch.tensor(kv_lens, dtype=torch.int32), + max_seq_len=max_seq_len, + ) + fmha = PrimsTSFmha(attn) + k = _TensorSpec((4, attn.num_kv_heads * head_dim), dtype) if has_separate_kv else None + v = _TensorSpec((4, attn.num_kv_heads * head_dim), dtype) if has_separate_kv else None + return fmha._is_supported_with_reason( + q, + k, + v, + attn, + metadata, + forward_args, + phase=phase, + ) + + +@pytest.mark.parametrize( + "case", + [ + { + "attention_input_type": AttentionInputType.context_only, + "head_dim": 128, + }, + { + "attention_input_type": AttentionInputType.mixed, + "head_dim": 256, + "dtype": torch.float16, + }, + { + "attention_input_type": AttentionInputType.generation_only, + "head_dim": 64, + "dtype": torch.float16, + }, + { + "attention_input_type": AttentionInputType.generation_only, + "head_dim": 576, + "is_mla": True, + "num_heads": 128, + }, + { + "attention_input_type": AttentionInputType.context_only, + "num_kv_cache_pools": 2, + "use_kv_cache_v2": True, + }, + { + "attention_input_type": AttentionInputType.context_only, + "num_heads": 64, + "num_kv_heads": 1, + }, + ], + ids=[ + "context", + "mixed", + "generation", + "mla-generation", + "v2-multi-pool", + "context-gqa-ratio-over-32", + ], +) +def test_supported_matrix(case: dict) -> None: + supported, reason = _support_result(**case) + + assert supported, reason + + +@pytest.mark.parametrize("phase", [FmhaPhase.CONTEXT, FmhaPhase.GENERATION]) +def test_phase_support_check_preserves_whole_request_semantics(phase: FmhaPhase) -> None: + supported, reason = _support_result( + attention_input_type=AttentionInputType.mixed, + head_dim=64, + phase=phase, + ) + + assert not supported + assert "context head dimension" in reason + + +def test_is_supported_accepts_and_forwards_phase_keyword( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + support_check = Mock(return_value=(True, "")) + monkeypatch.setattr(fmha, "_is_supported_with_reason", support_check) + q = Mock(spec=torch.Tensor) + metadata = SimpleNamespace() + forward_args = AttentionForwardArgs() + + assert fmha.is_supported( + q, + None, + None, + metadata, + forward_args, + phase=FmhaPhase.GENERATION, + ) + support_check.assert_called_once_with( + q, + None, + None, + attn, + metadata, + forward_args, + phase=FmhaPhase.GENERATION, + ) + + +@pytest.mark.parametrize( + ("case", "expected_reason"), + [ + ( + {"attention_input_type": AttentionInputType.context_only, "head_dim": 64}, + "context head dimension", + ), + ( + {"attention_input_type": AttentionInputType.generation_only, "head_dim": 96}, + "decode head dimension", + ), + ( + {"attention_input_type": AttentionInputType.context_only, "tokens_per_block": 8}, + "page size", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "has_paged_cache": False, + }, + "paged KV-cache block offsets", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "is_fused_qkv": False, + "has_separate_kv": True, + }, + "only fused QKV", + ), + ( + {"attention_input_type": AttentionInputType.context_only, "is_cross": True}, + "cross attention", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "has_sparse_attention": True, + }, + "sparse attention", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "has_sparse_runtime_metadata": True, + }, + "sparse attention metadata", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "use_spec_decoding": True, + }, + "speculative decoding", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "is_spec_dec_tree": True, + "use_spec_decoding": True, + }, + "speculative decoding", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "has_attention_sinks": True, + }, + "attention sinks", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "has_relative_attention_bias": True, + }, + "relative attention bias", + ), + ( + {"attention_input_type": AttentionInputType.generation_only, "beam_width": 2}, + "beam search", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "position_embedding_type": 4, + }, + "position embedding type", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "kv_dtype": DataType.HALF, + }, + "query and KV-cache dtypes", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "output_dtype": torch.float16, + }, + "output dtype must match", + ), + ( + {"attention_input_type": AttentionInputType.context_only, "has_output": False}, + "output tensor", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "num_heads": 7, + "num_kv_heads": 2, + }, + "divisible", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "num_heads": 64, + "num_kv_heads": 1, + }, + "GQA ratio", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "dtype": torch.float8_e4m3fn, + "kv_dtype": DataType.FP8, + }, + "query dtype", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "head_dim": 576, + "is_mla": True, + }, + "generation-only", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "head_dim": 576, + "is_mla": True, + "kv_lora_rank": 256, + }, + "kv_lora_rank=512", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "head_dim": 576, + "is_mla": True, + "qk_rope_head_dim": 32, + }, + "qk_rope_head_dim=64", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "head_dim": 576, + "is_mla": True, + "num_heads": 129, + }, + "at most 128 local query heads", + ), + ( + { + "attention_input_type": AttentionInputType.generation_only, + "head_dim": 640, + "is_mla": True, + }, + "latent plus RoPE dimensions", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "attention_window_size": 64, + }, + "cyclic TRT-LLM page tables", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "attention_window_size": 1, + }, + "cyclic TRT-LLM page tables", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "attention_chunk_size": 64, + }, + "chunked context attention", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "kv_layout": "NHD", + }, + "HND KV-cache layout", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "num_kv_cache_pools": 2, + }, + "V1 with multiple memory pools", + ), + ( + { + "attention_input_type": AttentionInputType.context_only, + "use_kv_cache_v2": True, + "enable_swa_scratch_reuse": True, + }, + "V2 SWA scratch reuse", + ), + ], + ids=[ + "context-head-dim", + "generation-head-dim", + "page-size", + "no-paged-cache", + "separate-qkv", + "cross", + "sparse", + "sparse-runtime-metadata", + "spec-decode", + "tree-mask", + "sinks", + "relative-bias", + "beam-search", + "alibi", + "kv-dtype-mismatch", + "output-dtype-mismatch", + "missing-output", + "heads-not-divisible", + "generation-head-ratio", + "fp8", + "mla-context", + "mla-kv-rank", + "mla-rope-dim", + "mla-too-many-heads", + "mla-head-dim", + "sliding-window", + "one-token-window", + "chunked-context", + "nhd-cache", + "v1-multi-pool", + "v2-swa-scratch-reuse", + ], +) +def test_unsupported_matrix_falls_through(case: dict, expected_reason: str) -> None: + supported, reason = _support_result(**case) + + assert not supported + assert expected_reason in reason + + +@pytest.mark.parametrize( + ("sm", "cutlass_version", "compiler_version", "expected"), + [ + (100, "4.7.0", "13.3", True), + (103, "4.7.0", "13.3", True), + (100, "4.7.0", "13.4", True), + (100, "4.7.0", "13.2", False), + (107, "4.7.0", "13.3", False), + (120, "4.7.0", "13.3", False), + (100, "4.6.2", "13.3", False), + ], +) +def test_static_availability_gate( + monkeypatch: pytest.MonkeyPatch, + sm: int, + cutlass_version: str, + compiler_version: str, + expected: bool, +) -> None: + target_version = Mock( + side_effect=lambda *, min_version: Version(compiler_version) >= Version(min_version) + ) + cutlass = SimpleNamespace(target_version=target_version) + + def import_cutlass_module(module_name: str) -> object: + if module_name == "cutlass": + return cutlass + assert module_name == "cutlass.experimental.task_scheduling" + return object() + + monkeypatch.setattr(prims_ts_module, "get_sm_version", lambda: sm) + monkeypatch.setattr(prims_ts_module, "version", lambda _: cutlass_version) + monkeypatch.setattr(prims_ts_module, "import_module", import_cutlass_module) + monkeypatch.setattr(PrimsTSFmha, "_missing_fused_nanobind_ops", staticmethod(lambda: [])) + + assert PrimsTSFmha.is_available(_Attention()) is expected + if sm in (100, 103) and Version(cutlass_version) >= Version("4.7.0"): + target_version.assert_called_once_with(min_version="13.3") + else: + target_version.assert_not_called() + + +def test_static_availability_gate_fails_closed_when_compiler_query_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cutlass = SimpleNamespace(target_version=Mock(side_effect=RuntimeError("query failed"))) + monkeypatch.setattr(prims_ts_module, "get_sm_version", lambda: 100) + monkeypatch.setattr(prims_ts_module, "version", lambda _: "4.7.0") + monkeypatch.setattr( + prims_ts_module, + "import_module", + lambda module_name: cutlass if module_name == "cutlass" else object(), + ) + monkeypatch.setattr(PrimsTSFmha, "_missing_fused_nanobind_ops", staticmethod(lambda: [])) + + assert not PrimsTSFmha.is_available(_Attention()) + cutlass.target_version.assert_called_once_with(min_version="13.3") + + +def test_unsupported_cutlass_compiler_excludes_prims_ts_from_fmha_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + target_version = Mock( + side_effect=lambda *, min_version: Version("13.2") >= Version(min_version) + ) + cutlass = SimpleNamespace(target_version=target_version) + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts,fallback") + monkeypatch.setattr(prims_ts_module, "get_sm_version", lambda: 100) + monkeypatch.setattr(prims_ts_module, "version", lambda _: "4.7.0") + monkeypatch.setattr( + prims_ts_module, + "import_module", + lambda module_name: cutlass if module_name == "cutlass" else object(), + ) + monkeypatch.setattr(PrimsTSFmha, "_missing_fused_nanobind_ops", staticmethod(lambda: [])) + + available_classes = [ + fmha_cls + for fmha_cls in get_enabled_fmha_lib_classes() + if fmha_cls.is_available(_Attention()) + ] + + assert available_classes == [FallbackFmha] + target_version.assert_called_once_with(min_version="13.3") + + +def test_v2_total_page_bound_is_not_expanded() -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + get_page_index_upper_bound = Mock(return_value=4096) + metadata = SimpleNamespace( + kv_cache_manager=_make_v2_manager( + impl=SimpleNamespace(get_page_index_upper_bound=get_page_index_upper_bound) + ) + ) + + assert fmha._get_total_num_blocks(metadata) == 4096 + assert get_page_index_upper_bound.call_args.args[0] == 0 + + +@pytest.mark.parametrize("is_mla", [False, True], ids=["standard", "mla"]) +def test_v1_total_page_bound_excludes_slots_before_selected_layer(is_mla: bool) -> None: + attn = _Attention(is_mla=is_mla) + attn.local_layer_idx = 3 + fmha = PrimsTSFmha(attn) + metadata = SimpleNamespace( + kv_cache_manager=_make_v1_manager( + impl=SimpleNamespace( + get_primary_pool_data=lambda _: torch.empty(64, dtype=torch.uint8) + ), + ), + host_kv_cache_pool_mapping=torch.tensor( + [[0, 0], [0, 1], [0, 2], [0, 3]], dtype=torch.int32 + ), + ) + + kv_factor = 1 if is_mla else 2 + assert fmha._get_total_num_blocks(metadata) == 64 * 4 * kv_factor - 3 * kv_factor + + +def test_kv_page_offset_uses_v2_manager_displacement() -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + metadata = SimpleNamespace( + kv_cache_manager=_make_v2_manager( + impl=SimpleNamespace(get_page_index_upper_bound=lambda *args: 128), + kv_offset=torch.tensor([0, 128]), + ), + host_kv_cache_pool_mapping=torch.tensor([[1, 0]], dtype=torch.int32), + ) + + assert ( + fmha_utils.get_kv_page_offset( + fmha.attn, + metadata, + 0, + cache=fmha._kv_page_offset_cache, + ) + == 128 + ) + + +def test_kv_page_offset_is_inferred_from_v1_host_tables() -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + host_offsets = torch.tensor( + [[[[0, 1, 2], [64, 65, 66]], [[3, 4, 5], [67, 68, 69]]]], + dtype=torch.int32, + ) + metadata = SimpleNamespace( + kv_cache_manager=_make_v1_manager( + impl=SimpleNamespace(), + host_kv_cache_block_offsets=host_offsets, + ), + host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), + ) + + assert ( + fmha_utils.get_kv_page_offset( + fmha.attn, + metadata, + 1, + cache=fmha._kv_page_offset_cache, + ) + == 64 + ) + + +def test_fixed_block_tables_are_zero_copy_plane_zero_view() -> None: + block_tables = torch.tensor( + [ + [[10, 11, 12], [110, 111, 112]], + [[20, 21, 22], [120, 121, 122]], + [[30, 31, 32], [130, 131, 132]], + ], + dtype=torch.int32, + ) + attn = _Attention() + fmha = PrimsTSFmha(attn) + + actual = fmha._get_fixed_block_tables(block_tables, 2) + + assert actual.shape == (2, 3) + assert actual.stride() == (6, 1) + assert actual.data_ptr() == block_tables.data_ptr() + torch.testing.assert_close( + actual, + torch.tensor([[10, 11, 12], [20, 21, 22]], dtype=torch.int32), + ) + + block_tables[1, 0, 2] = 99 + block_tables[:, 1].add_(1000) + + torch.testing.assert_close( + actual, + torch.tensor([[10, 11, 12], [20, 21, 99]], dtype=torch.int32), + ) + + +def test_sequence_lengths_are_live_source_view() -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + sequence_lengths = torch.tensor([0, 33, 64, 96], dtype=torch.int32)[1:] + + actual = fmha._get_sequence_lengths(sequence_lengths, 2) + + assert actual.shape == (2,) + assert actual.stride() == (1,) + assert actual.data_ptr() == sequence_lengths.data_ptr() + sequence_lengths[1] = 65 + torch.testing.assert_close(actual, torch.tensor([33, 65], dtype=torch.int32)) + + +def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + attn = _Attention() + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 120 + q_processed = torch.empty((3, attn.num_heads, attn.head_dim), dtype=torch.bfloat16) + kv_pool = torch.empty((12, attn.num_kv_heads, 32, attn.head_dim), dtype=torch.bfloat16) + block_tables = torch.tensor( + [ + [[0, 1, 2, 3], [6, 7, 8, 9]], + [[2, 3, 4, 5], [8, 9, 10, 11]], + ], + dtype=torch.int32, + ) + cu_q_seqlens = torch.tensor([0, 1, 3], dtype=torch.int32) + cu_kv_seqlens = torch.tensor([0, 7, 18], dtype=torch.int32) + fmha_workspace = torch.empty(0, dtype=torch.uint8) + context_preprocess = Mock( + return_value=( + q_processed, + kv_pool, + block_tables, + None, + 1.0, + 1.0, + fmha_workspace, + cu_q_seqlens, + cu_kv_seqlens, + 2, + 64, + -1, + ) + ) + context_postprocess = Mock() + wrapper = Mock() + wrapper_factory = Mock(return_value=wrapper) + monkeypatch.setattr( + prims_ts_module.thop, + "trtllm_gen_context_preprocess", + context_preprocess, + ) + monkeypatch.setattr( + prims_ts_module.thop, + "trtllm_gen_context_postprocess", + context_postprocess, + ) + monkeypatch.setattr( + prims_context_module, + "BatchPrefillPagedTSWrapper", + wrapper_factory, + ) + + host_block_offsets = torch.tensor( + [ + [ + [[0, 1, 2, 3], [6, 7, 8, 9]], + [[2, 3, 4, 5], [8, 9, 10, 11]], + [[4, 5, 0, 0], [10, 11, 0, 0]], + ] + ], + dtype=torch.int32, + ) + metadata = SimpleNamespace( + kv_cache_block_offsets=torch.empty((3, 2, 4), dtype=torch.int32), + host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), + kv_cache_manager=_make_v1_manager( + impl=SimpleNamespace(), + host_kv_cache_block_offsets=host_block_offsets, + ), + kv_lens_runtime=torch.tensor([7, 33, 64], dtype=torch.int32), + max_context_length=8, + max_seq_len=128, + ) + output = torch.empty((3, attn.num_heads, attn.head_dim), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.context_only, + attention_window_size=64, + is_fused_qkv=True, + ) + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=forward_args, + workspace=torch.empty(32, dtype=torch.uint8), + qkv_input=torch.empty( + (3, (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim), + dtype=torch.bfloat16, + ), + context_buf=output, + sequence_lengths=torch.tensor([33, 64], dtype=torch.int32), + context_lengths=torch.tensor([1, 2], dtype=torch.int32), + input_seq_length=2, + max_past_kv_length=64, + max_attention_window_size=64, + cyclic_attention_window_size=64, + num_tokens=3, + seq_offset=1, + tokens_per_block=32, + kv_factor=2, + total_num_blocks=24, + batch_size=2, + ) + + fmha.run_context(params) + + wrapper_factory.assert_called_once_with(kv_layout="HND") + wrapper.plan.assert_called_once() + plan_args = wrapper.plan.call_args.args + plan_kwargs = wrapper.plan.call_args.kwargs + assert plan_args == () + assert plan_kwargs == { + "device": q_processed.device, + "batch_size": 2, + "max_seq_len_q": 8, + "max_kv_len": 128, + "num_qo_heads": attn.num_heads, + "num_kv_heads": attn.num_kv_heads, + "head_dim": attn.head_dim, + "q_dtype": torch.bfloat16, + "kv_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, + "page_size": 32, + "mask_type": "causal", + "window_left": -1, + "sm_scale": pytest.approx(1.0 / math.sqrt(attn.head_dim)), + "output_scale": 1.0, + } + wrapper.run.assert_called_once() + run_args = wrapper.run.call_args.args + run_kwargs = wrapper.run.call_args.kwargs + assert run_args[0] is q_processed + k_cache, v_cache = run_args[1], run_args[2] + assert k_cache.shape == v_cache.shape == (6, attn.num_kv_heads, 32, attn.head_dim) + assert v_cache.storage_offset() - k_cache.storage_offset() == 6 * math.prod(kv_pool.shape[1:]) + fixed_block_tables = run_kwargs["block_tables"] + seq_lens_kv = run_kwargs["seq_lens_kv"] + first_metadata_ptrs = ( + run_args[3].data_ptr(), + fixed_block_tables.data_ptr(), + seq_lens_kv.data_ptr(), + ) + assert run_kwargs["out"] is output + assert run_kwargs["validate"] is False + assert run_args[3] is cu_q_seqlens + torch.testing.assert_close(seq_lens_kv, torch.tensor([33, 64], dtype=torch.int32)) + assert seq_lens_kv.data_ptr() == params.sequence_lengths.data_ptr() + assert fixed_block_tables.shape == (2, 4) + assert fixed_block_tables.stride() == (8, 1) + assert fixed_block_tables.data_ptr() == block_tables.data_ptr() + torch.testing.assert_close( + fixed_block_tables, + torch.tensor([[0, 1, 2, 3], [2, 3, 4, 5]], dtype=torch.int32), + ) + context_preprocess.assert_called_once() + context_postprocess.assert_called_once() + assert context_preprocess.call_args.kwargs["skip_fmha_workspace"] is True + assert context_postprocess.call_args.kwargs["skip_fmha_workspace"] is True + + block_tables[:, 0].add_(1) + block_tables[:, 1].add_(100) + params.sequence_lengths.copy_(torch.tensor([64, 33], dtype=torch.int32)) + cu_kv_seqlens.copy_(torch.tensor([0, 4, 9], dtype=torch.int32)) + fmha.run_context(params) + + wrapper_factory.assert_called_once_with(kv_layout="HND") + wrapper.plan.assert_called_once() + assert wrapper.run.call_count == 2 + second_run_args = wrapper.run.call_args.args + second_run_kwargs = wrapper.run.call_args.kwargs + assert ( + second_run_args[3].data_ptr(), + second_run_kwargs["block_tables"].data_ptr(), + second_run_kwargs["seq_lens_kv"].data_ptr(), + ) == first_metadata_ptrs + torch.testing.assert_close( + second_run_kwargs["seq_lens_kv"], + torch.tensor([64, 33], dtype=torch.int32), + ) + torch.testing.assert_close( + second_run_kwargs["block_tables"], + torch.tensor([[1, 2, 3, 4], [3, 4, 5, 6]], dtype=torch.int32), + ) + + +@pytest.mark.parametrize( + ( + "use_split_kv", + "use_separate_reduction_kernel", + "use_cluster_smem_reduction", + "requires_control_reset", + ), + ( + pytest.param(False, False, False, False, id="direct"), + pytest.param(True, False, False, True, id="fused-global-reduction"), + pytest.param(True, True, False, False, id="separate-reduction"), + pytest.param(True, False, True, False, id="cluster-smem-reduction"), + ), +) +def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( + monkeypatch: pytest.MonkeyPatch, + use_split_kv: bool, + use_separate_reduction_kernel: bool, + use_cluster_smem_reduction: bool, + requires_control_reset: bool, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + attn = _Attention() + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 120 + fmha._decode_workspace_offset_bytes = 0 + fmha._decode_workspace_required_bytes = 64 + fmha_workspace = torch.empty(0, dtype=torch.uint8) + q_processed = torch.empty((2, attn.num_heads, attn.head_dim), dtype=torch.bfloat16) + kv_pool = torch.empty((12, attn.num_kv_heads, 32, attn.head_dim), dtype=torch.bfloat16) + block_tables = torch.tensor( + [ + [[0, 1, 2, 3], [6, 7, 8, 9]], + [[2, 3, 4, 5], [8, 9, 10, 11]], + ], + dtype=torch.int32, + ) + workspace = torch.full((64,), 7, dtype=torch.uint8) + split_kv_counter = workspace[32:40].view(torch.int32) + generation_preprocess = Mock( + return_value=( + q_processed, + kv_pool, + block_tables, + None, + 1.0, + 1.0, + fmha_workspace, + None, + 1, + 64, + 15, + False, + ) + ) + wrapper = Mock() + wrapper._plan_state = SimpleNamespace( + policy=( + ("use_split_kv", use_split_kv), + ("use_separate_reduction_kernel", use_separate_reduction_kernel), + ("use_cluster_smem_reduction", use_cluster_smem_reduction), + ), + workspace=SimpleNamespace(split_kv_counter=split_kv_counter), + ) + wrapper_factory = Mock(return_value=wrapper) + monkeypatch.setattr( + prims_ts_module.thop, + "trtllm_gen_generation_preprocess", + generation_preprocess, + ) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_decode_workspace_size", + Mock(return_value=64), + ) + monkeypatch.setattr( + prims_decode_module, + "BatchDecodePagedTSWrapper", + wrapper_factory, + ) + + get_page_index_upper_bound = Mock(return_value=12) + metadata = SimpleNamespace( + beam_width=1, + kv_cache_block_offsets=torch.empty((2, 2, 4), dtype=torch.int32), + host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), + kv_cache_manager=_make_v2_manager( + impl=SimpleNamespace(get_page_index_upper_bound=get_page_index_upper_bound), + kv_offset=torch.tensor([6], dtype=torch.int32), + ), + ) + total_num_blocks = fmha._get_total_num_blocks(metadata) + output = torch.empty((2, attn.num_heads, attn.head_dim), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.generation_only, + attention_window_size=64, + is_fused_qkv=True, + ) + sequence_lengths = torch.tensor([0, 33, 64], dtype=torch.int32)[1:] + assert sequence_lengths.data_ptr() % 16 != 0 + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=forward_args, + workspace=workspace, + qkv_input=torch.empty( + (2, (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim), + dtype=torch.bfloat16, + ), + context_buf=output, + sequence_lengths=sequence_lengths, + input_seq_length=1, + max_past_kv_length=64, + max_attention_window_size=64, + cyclic_attention_window_size=64, + num_tokens=2, + seq_offset=1, + tokens_per_block=32, + kv_factor=2, + total_num_blocks=total_num_blocks, + batch_size=2, + num_requests=2, + ) + + fmha.run_generation(params) + + wrapper_factory.assert_called_once_with(kv_layout="HND") + wrapper.plan.assert_called_once() + plan_args = wrapper.plan.call_args.args + plan_kwargs = wrapper.plan.call_args.kwargs + assert plan_args == ( + params.workspace.device, + 2, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + 32, + 128, + ) + decode_workspace = plan_kwargs["workspace_buffer"] + assert decode_workspace.data_ptr() == params.workspace.data_ptr() + assert decode_workspace.numel() == 64 + assert {key: value for key, value in plan_kwargs.items() if key != "workspace_buffer"} == { + "max_seq_len_q": 1, + "packed_query": False, + "q_data_type": torch.bfloat16, + "kv_data_type": torch.bfloat16, + "o_data_type": torch.bfloat16, + "mask_type": "causal", + "window_left": 15, + } + assert fmha._decode_wrappers[2] is wrapper + wrapper.run.assert_called_once() + run_args = wrapper.run.call_args.args + run_kwargs = wrapper.run.call_args.kwargs + assert run_args[0].shape == (2, attn.num_heads, attn.head_dim) + assert run_args[0].data_ptr() == q_processed.data_ptr() + k_cache, v_cache = run_args[1] + assert k_cache.shape == v_cache.shape == (6, attn.num_kv_heads, 32, attn.head_dim) + assert v_cache.storage_offset() - k_cache.storage_offset() == 6 * math.prod(kv_pool.shape[1:]) + assert run_args[2].data_ptr() == params.sequence_lengths.data_ptr() + fixed_block_tables = run_kwargs["block_tables"] + assert fixed_block_tables.shape == (2, 4) + assert fixed_block_tables.stride() == (8, 1) + assert fixed_block_tables.data_ptr() == block_tables.data_ptr() + torch.testing.assert_close( + fixed_block_tables, + torch.tensor([[0, 1, 2, 3], [2, 3, 4, 5]], dtype=torch.int32), + ) + assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(attn.head_dim)) + assert run_kwargs["bmm2_scale"] == 1.0 + assert run_kwargs["out"].shape == (2, attn.num_heads, attn.head_dim) + assert run_kwargs["out"].data_ptr() == output.data_ptr() + assert run_kwargs["validate"] is False + if requires_control_reset: + assert torch.count_nonzero(split_kv_counter) == 0 + else: + torch.testing.assert_close( + split_kv_counter, + torch.full_like(split_kv_counter, 0x07070707), + ) + torch.testing.assert_close(params.workspace[:32], torch.full((32,), 7, dtype=torch.uint8)) + torch.testing.assert_close(params.workspace[40:], torch.full((24,), 7, dtype=torch.uint8)) + preprocess_args = generation_preprocess.call_args.args + assert preprocess_args[15] == params.seq_offset + assert preprocess_args[39] == total_num_blocks + assert generation_preprocess.call_args.kwargs["skip_fmha_workspace"] is True + get_page_index_upper_bound.assert_called_once() + + block_tables[:, 0].add_(20) + block_tables[:, 1].add_(200) + sequence_lengths.add_(1) + params.workspace.fill_(9) + fmha.run_generation(params) + + wrapper_factory.assert_called_once_with(kv_layout="HND") + wrapper.plan.assert_called_once() + assert wrapper.run.call_count == 2 + second_run_kwargs = wrapper.run.call_args.kwargs + assert second_run_kwargs["block_tables"].data_ptr() == block_tables.data_ptr() + assert second_run_kwargs["block_tables"].stride() == (8, 1) + torch.testing.assert_close( + second_run_kwargs["block_tables"], + torch.tensor([[20, 21, 22, 23], [22, 23, 24, 25]], dtype=torch.int32), + ) + torch.testing.assert_close(run_args[2], torch.tensor([34, 65], dtype=torch.int32)) + if requires_control_reset: + assert torch.count_nonzero(split_kv_counter) == 0 + else: + torch.testing.assert_close( + split_kv_counter, + torch.full_like(split_kv_counter, 0x09090909), + ) + torch.testing.assert_close(params.workspace[:32], torch.full((32,), 9, dtype=torch.uint8)) + torch.testing.assert_close(params.workspace[40:], torch.full((24,), 9, dtype=torch.uint8)) + + +def test_decode_layer_adapters_bind_the_same_shared_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrappers = [Mock(), Mock()] + wrapper_factory = Mock(side_effect=wrappers) + monkeypatch.setattr( + prims_decode_module, + "BatchDecodePagedTSWrapper", + wrapper_factory, + ) + attentions = [_Attention(), _Attention()] + layers = [PrimsTSFmha(attn) for attn in attentions] + shared_workspace = torch.empty(64, dtype=torch.uint8) + + def get_wrapper(layer: PrimsTSFmha) -> object: + return layer._get_or_plan_decode_wrapper( + shared_workspace, + batch_size=2, + num_qo_heads=8, + num_kv_heads=2, + head_dim=128, + page_size=32, + seq_len_q=1, + max_kv_len=64, + q_dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + mask_type="causal", + window_left=-1, + ) + + first_results = [get_wrapper(layer) for layer in layers] + second_results = [get_wrapper(layer) for layer in layers] + + assert first_results == second_results == wrappers + assert wrapper_factory.call_count == 2 + for wrapper in wrappers: + wrapper.plan.assert_called_once() + assert wrapper.plan.call_args.kwargs["workspace_buffer"] is shared_workspace + + +def test_context_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrappers = [Mock(), Mock()] + wrapper_factory = Mock(side_effect=wrappers) + monkeypatch.setattr( + prims_context_module, + "BatchPrefillPagedTSWrapper", + wrapper_factory, + ) + attn = _Attention() + fmha = PrimsTSFmha(attn) + k_cache = torch.empty((8, 2, 32, 128), dtype=torch.bfloat16) + v_cache = torch.empty_like(k_cache) + + def get_wrapper(batch_size: int) -> object: + q = torch.empty((batch_size, 8, 128), dtype=torch.bfloat16) + return fmha._get_or_plan_context_wrapper( + q, + k_cache, + v_cache, + batch_size=batch_size, + max_seq_len_q=128, + max_kv_len=256, + page_size=32, + mask_type="causal", + window_left=-1, + sm_scale=1.0 / math.sqrt(128), + output_dtype=torch.bfloat16, + ) + + first_a = get_wrapper(1) + profile_b = get_wrapper(2) + second_a = get_wrapper(1) + + assert first_a is second_a is wrappers[0] + assert profile_b is wrappers[1] + assert wrapper_factory.call_count == 2 + for wrapper in wrappers: + wrapper.plan.assert_called_once() + assert set(fmha._context_wrappers) == {1, 2} + + +def test_decode_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrappers = [Mock(), Mock()] + wrapper_factory = Mock(side_effect=wrappers) + monkeypatch.setattr( + prims_decode_module, + "BatchDecodePagedTSWrapper", + wrapper_factory, + ) + attn = _Attention() + fmha = PrimsTSFmha(attn) + workspace = torch.empty(64, dtype=torch.uint8) + + def get_wrapper(batch_size: int) -> object: + return fmha._get_or_plan_decode_wrapper( + workspace, + batch_size=batch_size, + num_qo_heads=8, + num_kv_heads=2, + head_dim=128, + page_size=32, + seq_len_q=1, + max_kv_len=256, + q_dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + mask_type="causal", + window_left=-1, + ) + + first_a = get_wrapper(1) + profile_b = get_wrapper(2) + second_a = get_wrapper(1) + + assert first_a is second_a is wrappers[0] + assert profile_b is wrappers[1] + assert wrapper_factory.call_count == 2 + for wrapper in wrappers: + wrapper.plan.assert_called_once() + assert set(fmha._decode_wrappers) == {1, 2} + + +def test_workspace_allocation_change_invalidates_only_workspace_bound_wrappers() -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + first_workspace = torch.empty(32, dtype=torch.uint8) + fmha._update_workspace_allocation(first_workspace) + fmha._context_wrappers[1] = Mock() + fmha._decode_wrappers[1] = Mock() + fmha._mla_decode_wrappers[1] = Mock() + + second_workspace = torch.empty(64, dtype=torch.uint8) + fmha._update_workspace_allocation(second_workspace) + + assert set(fmha._context_wrappers) == {1} + assert fmha._decode_wrappers == {} + assert fmha._mla_decode_wrappers == {} + + +def _get_test_mla_wrapper( + fmha: PrimsTSFmha, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + workspace_buffer: torch.Tensor, + *, + mask_type: str = "causal", +) -> object: + return fmha._get_or_plan_mla_decode_wrapper( + workspace_buffer, + batch_size=int(block_tables.shape[0]), + num_heads=4, + kv_lora_rank=512, + qk_rope_head_dim=64, + page_size=32, + max_seq_len_q=1, + max_kv_len=96, + q_dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + mask_type=mask_type, + ) + + +def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + kv_cache = torch.empty((20, 1, 32, 576), dtype=torch.bfloat16) + block_tables = torch.tensor( + [ + [[0, 1, 2], [10, 11, 12]], + [[3, 4, 5], [13, 14, 15]], + ], + dtype=torch.int32, + ) + build_metadata = Mock(return_value=(kv_cache, block_tables, None)) + wrapper = Mock() + wrapper_factory = Mock(return_value=wrapper) + monkeypatch.setattr( + prims_ts_module.thop, + "build_trtllm_gen_kv_cache_metadata", + build_metadata, + ) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_mla_decode_workspace_size", + Mock(return_value=64), + ) + monkeypatch.setattr( + prims_mla_module, + "BatchMLADecodePagedTSWrapper", + wrapper_factory, + ) + + metadata = SimpleNamespace( + is_cuda_graph=False, + beam_width=1, + kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), + host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), + ) + output = torch.empty((2, attn.num_heads, 512), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.generation_only, + attention_window_size=64, + is_fused_qkv=True, + ) + sequence_lengths = torch.tensor([33, 64], dtype=torch.int32) + assert sequence_lengths.data_ptr() % 16 == 0 + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=forward_args, + workspace=torch.empty(64, dtype=torch.uint8), + qkv_input=torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16), + context_buf=output, + sequence_lengths=sequence_lengths, + input_seq_length=1, + max_past_kv_length=64, + max_attention_window_size=64, + cyclic_attention_window_size=64, + num_tokens=2, + seq_offset=2, + tokens_per_block=32, + kv_factor=1, + total_num_blocks=20, + batch_size=2, + num_requests=2, + ) + + fmha.run_mla_generation(params) + + wrapper_factory.assert_called_once_with() + wrapper.plan.assert_called_once() + plan_args = wrapper.plan.call_args.args + plan_kwargs = wrapper.plan.call_args.kwargs + assert plan_args == ( + params.workspace.device, + 2, + attn.num_heads, + 512, + 64, + 32, + 96, + ) + workspace_buffer = plan_kwargs["workspace_buffer"] + assert workspace_buffer.data_ptr() == params.workspace.data_ptr() + assert workspace_buffer.numel() == 64 + assert {key: value for key, value in plan_kwargs.items() if key != "workspace_buffer"} == { + "max_seq_len_q": 1, + "packed_query": False, + "q_data_type": torch.bfloat16, + "kv_data_type": torch.bfloat16, + "o_data_type": torch.bfloat16, + "mask_type": "causal", + } + wrapper.run.assert_called_once() + run_args = wrapper.run.call_args.args + run_kwargs = wrapper.run.call_args.kwargs + assert run_args[0].shape == (2, 1, attn.num_heads, 576) + assert run_args[0].data_ptr() == params.qkv_input.data_ptr() + assert run_args[1] is kv_cache + assert run_kwargs["block_tables"].shape == (2, 3) + assert run_kwargs["block_tables"].stride() == (6, 1) + assert run_kwargs["block_tables"].data_ptr() == block_tables.data_ptr() + torch.testing.assert_close( + run_kwargs["block_tables"], + torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32), + ) + assert run_kwargs["seq_lens"].data_ptr() == sequence_lengths.data_ptr() + torch.testing.assert_close(run_kwargs["seq_lens"], sequence_lengths) + assert run_kwargs["out"].shape == (2, 1, attn.num_heads, 512) + assert run_kwargs["out"].data_ptr() == output.data_ptr() + assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(128 + 64)) + assert run_kwargs["bmm2_scale"] == 1.0 + assert run_kwargs["validate"] is False + + block_tables[:, 0].add_(20) + block_tables[:, 1].add_(200) + sequence_lengths.add_(1) + fmha.run_mla_generation(params) + + wrapper_factory.assert_called_once_with() + wrapper.plan.assert_called_once() + assert wrapper.run.call_count == 2 + assert run_kwargs["block_tables"].data_ptr() == block_tables.data_ptr() + assert run_kwargs["block_tables"].stride() == (6, 1) + torch.testing.assert_close( + run_kwargs["block_tables"], + torch.tensor([[20, 21, 22], [23, 24, 25]], dtype=torch.int32), + ) + torch.testing.assert_close(run_kwargs["seq_lens"], torch.tensor([34, 65], dtype=torch.int32)) + assert fmha._mla_decode_wrappers[2] is wrapper + + +def test_mla_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + wrappers = [Mock(), Mock()] + wrapper_factory = Mock(side_effect=wrappers) + monkeypatch.setattr( + prims_mla_module, + "BatchMLADecodePagedTSWrapper", + wrapper_factory, + ) + block_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32) + seq_lens = torch.tensor([33, 64], dtype=torch.int32) + workspace = torch.empty(64, dtype=torch.uint8) + + first = _get_test_mla_wrapper(fmha, block_tables, seq_lens, workspace) + cached_first = _get_test_mla_wrapper(fmha, block_tables, seq_lens, workspace) + block_pointer_hit = _get_test_mla_wrapper( + fmha, + block_tables.clone(), + seq_lens, + workspace, + ) + seq_pointer_hit = _get_test_mla_wrapper( + fmha, + block_tables, + seq_lens.clone(), + workspace, + ) + profile_b = _get_test_mla_wrapper( + fmha, + block_tables[:1], + seq_lens[:1], + workspace, + ) + profile_a_again = _get_test_mla_wrapper( + fmha, + block_tables, + seq_lens, + workspace, + ) + + assert all( + result is wrappers[0] + for result in ( + first, + cached_first, + block_pointer_hit, + seq_pointer_hit, + profile_a_again, + ) + ) + assert profile_b is wrappers[1] + assert wrapper_factory.call_count == 2 + for wrapper in wrappers: + wrapper.plan.assert_called_once() + assert wrapper.plan.call_args.kwargs["workspace_buffer"] is workspace + assert fmha._mla_decode_wrappers[2] is wrappers[0] + assert fmha._mla_decode_wrappers[1] is wrappers[1] + + +def test_mla_wrapper_capture_uses_cached_plan_and_rejects_plan_miss( + monkeypatch: pytest.MonkeyPatch, +) -> None: + capturing = False + monkeypatch.setattr( + torch.cuda, + "is_current_stream_capturing", + lambda: capturing, + ) + wrapper = Mock() + wrapper_factory = Mock(return_value=wrapper) + monkeypatch.setattr( + prims_mla_module, + "BatchMLADecodePagedTSWrapper", + wrapper_factory, + ) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + block_tables = torch.tensor([[0, 1, 2]], dtype=torch.int32) + seq_lens = torch.tensor([33], dtype=torch.int32) + workspace = torch.empty(64, dtype=torch.uint8) + planned = _get_test_mla_wrapper(fmha, block_tables, seq_lens, workspace) + capturing = True + + cached = _get_test_mla_wrapper(fmha, block_tables, seq_lens, workspace) + with pytest.raises(RuntimeError, match="must be planned before CUDA graph capture"): + _get_test_mla_wrapper( + fmha, + torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32), + torch.tensor([33, 64], dtype=torch.int32), + workspace, + ) + + assert cached is planned is wrapper + wrapper_factory.assert_called_once_with() + wrapper.plan.assert_called_once() + + +@pytest.mark.parametrize("is_cuda_graph", [False, True], ids=["eager", "cuda-graph"]) +def test_mla_wrapper_receives_v2_bound_and_shared_workspace( + monkeypatch: pytest.MonkeyPatch, + is_cuda_graph: bool, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + kv_cache = torch.empty((20, 1, 32, 576), dtype=torch.bfloat16) + block_tables = torch.tensor( + [ + [[0, 1, 2], [10, 11, 12]], + [[3, 4, 5], [13, 14, 15]], + ], + dtype=torch.int32, + ) + build_metadata = Mock(return_value=(kv_cache, block_tables, None)) + wrapper = Mock() + wrapper_factory = Mock(return_value=wrapper) + monkeypatch.setattr( + prims_ts_module.thop, + "build_trtllm_gen_kv_cache_metadata", + build_metadata, + ) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_mla_decode_workspace_size", + Mock(return_value=96), + ) + monkeypatch.setattr( + prims_mla_module, + "BatchMLADecodePagedTSWrapper", + wrapper_factory, + ) + + get_page_index_upper_bound = Mock(return_value=20) + metadata = SimpleNamespace( + is_cuda_graph=is_cuda_graph, + beam_width=1, + kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), + host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), + kv_cache_manager=_make_v2_manager( + impl=SimpleNamespace(get_page_index_upper_bound=get_page_index_upper_bound) + ), + ) + total_num_blocks = fmha._get_total_num_blocks(metadata) + output = torch.empty((2, attn.num_heads, 512), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.generation_only, + attention_window_size=64, + is_fused_qkv=True, + ) + sequence_lengths = torch.tensor([0, 33, 64], dtype=torch.int32)[1:] + assert sequence_lengths.data_ptr() % 16 != 0 + params = FmhaParams( + attn=attn, + meta=metadata, + fwd=forward_args, + workspace=torch.empty(96, dtype=torch.uint8), + qkv_input=torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16), + context_buf=output, + sequence_lengths=sequence_lengths, + input_seq_length=1, + max_past_kv_length=64, + max_attention_window_size=64, + cyclic_attention_window_size=64, + num_tokens=2, + seq_offset=2, + tokens_per_block=32, + kv_factor=1, + total_num_blocks=total_num_blocks, + batch_size=2, + num_requests=2, + ) + + fmha.run_mla_generation(params) + + wrapper_factory.assert_called_once_with() + wrapper.plan.assert_called_once() + plan_args = wrapper.plan.call_args.args + plan_kwargs = wrapper.plan.call_args.kwargs + assert plan_args == ( + params.workspace.device, + 2, + attn.num_heads, + 512, + 64, + 32, + 96, + ) + assert plan_kwargs["packed_query"] is False + assert plan_kwargs["workspace_buffer"].data_ptr() == params.workspace.data_ptr() + assert plan_kwargs["workspace_buffer"].numel() == 96 + wrapper.run.assert_called_once() + run_args = wrapper.run.call_args.args + run_kwargs = wrapper.run.call_args.kwargs + assert run_args[0].shape == (2, 1, attn.num_heads, 576) + assert run_args[0].data_ptr() == params.qkv_input.data_ptr() + assert run_args[1] is kv_cache + torch.testing.assert_close( + run_kwargs["block_tables"], + torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32), + ) + assert run_kwargs["block_tables"].shape == (2, 3) + assert run_kwargs["block_tables"].stride() == (6, 1) + assert run_kwargs["block_tables"].data_ptr() == block_tables.data_ptr() + torch.testing.assert_close(run_kwargs["seq_lens"], params.sequence_lengths) + assert run_kwargs["seq_lens"].data_ptr() == params.sequence_lengths.data_ptr() + assert run_kwargs["out"].shape == (2, 1, attn.num_heads, 512) + assert run_kwargs["out"].data_ptr() == output.data_ptr() + assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(128 + 64)) + assert run_kwargs["bmm2_scale"] == 1.0 + assert run_kwargs["validate"] is False + builder_args = build_metadata.call_args.args + assert builder_args[8] == total_num_blocks + assert builder_args[10] == params.seq_offset + assert builder_args[11] == 2 + assert builder_args[12] == torch.bfloat16 + get_page_index_upper_bound.assert_called_once() + + +@pytest.mark.parametrize("is_cuda_graph", [False, True], ids=["eager", "cuda-graph"]) +def test_mla_prepare_workspace_sizes_caller_owned_workspace( + monkeypatch: pytest.MonkeyPatch, + is_cuda_graph: bool, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + workspace_size = Mock(return_value=48) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_mla_decode_workspace_size", + workspace_size, + ) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) + metadata = SimpleNamespace( + is_cuda_graph=is_cuda_graph, + kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), + max_num_requests=2, + num_contexts=0, + num_generations=2, + num_ctx_tokens=0, + tokens_per_block=32, + kv_lens_runtime=torch.tensor([33, 64], dtype=torch.int32), + ) + output = torch.empty((2, attn.num_heads * 512), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.generation_only, + attention_window_size=96, + ) + + workspace = torch.empty(0, dtype=torch.uint8) + fmha.prepare_workspace( + q, + None, + None, + metadata, + forward_args, + workspace, + ) + + workspace_size.assert_called_once() + assert workspace.numel() == 48 + + +def test_mla_prepare_workspace_preserves_cached_wrappers_with_stable_allocation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + stream = Mock() + monkeypatch.setattr(torch.cuda, "current_stream", Mock(return_value=stream)) + workspace_size = Mock(return_value=48) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_mla_decode_workspace_size", + workspace_size, + ) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + workspace = torch.empty(48, dtype=torch.uint8) + fmha._update_workspace_allocation(workspace) + wrapper = Mock() + fmha._mla_decode_wrappers[2] = wrapper + q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) + metadata = SimpleNamespace( + is_cuda_graph=True, + kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), + max_num_requests=2, + num_contexts=0, + num_generations=2, + num_ctx_tokens=0, + tokens_per_block=32, + kv_lens_runtime=torch.tensor([33, 64], dtype=torch.int32), + ) + output = torch.empty((2, attn.num_heads * 512), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.generation_only, + attention_window_size=96, + ) + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + + stream.synchronize.assert_not_called() + assert fmha._mla_decode_wrappers[2] is wrapper + assert workspace.numel() == 48 + workspace_size.assert_called_once() + + +def test_mla_caller_workspace_grows_across_plan_profiles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + required_bytes = iter((32, 64)) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_mla_decode_workspace_size", + lambda *args, **kwargs: next(required_bytes), + ) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) + metadata = SimpleNamespace( + is_cuda_graph=False, + kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), + max_num_requests=2, + num_contexts=0, + num_generations=2, + num_ctx_tokens=0, + tokens_per_block=32, + kv_lens_runtime=torch.tensor([33, 64], dtype=torch.int32), + ) + forward_args = AttentionForwardArgs( + output=torch.empty((2, attn.num_heads * 512), dtype=torch.bfloat16), + attention_input_type=AttentionInputType.generation_only, + attention_window_size=96, + ) + workspace = torch.empty(0, dtype=torch.uint8) + + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + assert workspace.numel() == 32 + + dense_forward_args = replace( + forward_args, + attention_mask=PredefinedAttentionMask.FULL, + ) + fmha.prepare_workspace(q, None, None, metadata, dense_forward_args, workspace) + + assert workspace.numel() == 64 + + +def test_decode_prepare_workspace_reserves_tail_after_compact_preprocessing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + prims_ts_module.thop, + "get_trtllm_gen_generation_workspace_layout", + lambda *args, **kwargs: {"total_size": 64}, + ) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_decode_workspace_size", + lambda *args, **kwargs: 48, + ) + attn = _Attention() + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + metadata = SimpleNamespace( + kv_cache_block_offsets=torch.empty((2, 2, 4), dtype=torch.int32), + max_num_requests=2, + num_contexts=0, + num_generations=2, + num_ctx_tokens=0, + tokens_per_block=32, + ) + forward_args = AttentionForwardArgs( + output=torch.empty((2, 8 * 128), dtype=torch.bfloat16), + attention_input_type=AttentionInputType.generation_only, + attention_window_size=128, + ) + workspace = torch.empty(0, dtype=torch.uint8) + + fmha.prepare_workspace( + torch.empty((2, 12 * 128), dtype=torch.bfloat16), + None, + None, + metadata, + forward_args, + workspace, + ) + + assert fmha._decode_workspace_offset_bytes == 64 + assert fmha._decode_workspace_required_bytes == 48 + assert workspace.numel() == 112 + decode_workspace = fmha._get_decode_workspace(workspace) + assert decode_workspace.data_ptr() == workspace.data_ptr() + 64 + assert decode_workspace.numel() == 48 + + +def test_decode_workspace_tail_is_stable_across_mixed_context_layouts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + context_layout = Mock( + side_effect=( + {"total_size": 64}, + {"total_size": 320}, + ) + ) + monkeypatch.setattr( + prims_ts_module.thop, + "get_trtllm_gen_context_workspace_layout", + context_layout, + ) + monkeypatch.setattr( + prims_ts_module.thop, + "get_trtllm_gen_generation_workspace_layout", + lambda *args, **kwargs: {"total_size": 64}, + ) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_decode_workspace_size", + lambda *args, **kwargs: 48, + ) + attn = _Attention() + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + metadata = SimpleNamespace( + kv_cache_block_offsets=torch.empty((4, 2, 4), dtype=torch.int32), + max_num_requests=4, + num_contexts=2, + num_generations=2, + num_ctx_tokens=4, + tokens_per_block=32, + ) + forward_args = AttentionForwardArgs( + output=torch.empty((6, 8 * 128), dtype=torch.bfloat16), + attention_input_type=AttentionInputType.mixed, + attention_window_size=128, + ) + q = torch.empty((6, 12 * 128), dtype=torch.bfloat16) + workspace = torch.empty(1024, dtype=torch.uint8) + + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + first_workspace = fmha._get_decode_workspace(workspace) + cached_wrapper = Mock() + fmha._decode_wrappers[2] = cached_wrapper + + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + second_workspace = fmha._get_decode_workspace(workspace) + + assert context_layout.call_count == 2 + assert fmha._decode_wrappers[2] is cached_wrapper + assert fmha._decode_workspace_offset_bytes == 960 + assert first_workspace.data_ptr() == second_workspace.data_ptr() + assert first_workspace.numel() == second_workspace.numel() == 48 + + +def test_workspace_cannot_grow_during_capture(monkeypatch: pytest.MonkeyPatch) -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + monkeypatch.setattr( + prims_ts_module.thop, + "get_trtllm_gen_generation_workspace_layout", + lambda *args, **kwargs: {"total_size": 32}, + ) + monkeypatch.setattr( + prims_ts_package, + "get_prims_ts_batch_decode_workspace_size", + lambda *args, **kwargs: 32, + ) + metadata = SimpleNamespace( + kv_cache_block_offsets=torch.empty((2, 2, 4), dtype=torch.int32), + max_num_requests=2, + num_contexts=0, + num_generations=2, + num_ctx_tokens=0, + tokens_per_block=32, + kv_lens_runtime=torch.tensor([64, 96], dtype=torch.int32), + ) + forward_args = AttentionForwardArgs( + output=torch.empty((2, 8 * 128), dtype=torch.bfloat16), + attention_input_type=AttentionInputType.generation_only, + attention_window_size=128, + ) + + with pytest.raises( + RuntimeError, + match="PrimTS caller workspace must be sized before CUDA graph capture", + ): + fmha.prepare_workspace( + torch.empty((2, 12 * 128), dtype=torch.bfloat16), + None, + None, + metadata, + forward_args, + torch.empty(16, dtype=torch.uint8), + ) + + +def test_phased_forward_routes_mixed_batch_to_context_and_generation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attn = _Attention() + fmha = PrimsTSFmha(attn) + context_calls = [] + generation_calls = [] + run_context = Mock(side_effect=lambda params: context_calls.append(replace(params))) + run_generation = Mock(side_effect=lambda params: generation_calls.append(replace(params))) + monkeypatch.setattr(fmha, "prepare_workspace", Mock()) + monkeypatch.setattr(fmha, "run_context", run_context) + monkeypatch.setattr(fmha, "run_generation", run_generation) + + q = torch.empty((5, 128), dtype=torch.bfloat16) + output = torch.empty((5, attn.num_heads * attn.head_dim), dtype=torch.bfloat16) + metadata = SimpleNamespace( + kv_cache_block_offsets=torch.empty(1), + effective_workspace=torch.empty(0, dtype=torch.int8), + num_contexts=1, + num_ctx_tokens=3, + num_generations=2, + cache_indirection=None, + beam_width=1, + tokens_per_block=32, + kv_lens_cuda_runtime=torch.tensor([3, 65, 97], dtype=torch.int32), + kv_lens_runtime=torch.tensor([3, 65, 97], dtype=torch.int32), + prompt_lens_cuda_runtime=torch.tensor([3, 1, 1], dtype=torch.int32), + prompt_lens_cpu_runtime=torch.tensor([3, 1, 1], dtype=torch.int32), + is_spec_decoding_enabled=False, + is_cross=False, + kv_cache_manager=None, + ) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.mixed, + attention_window_size=128, + ) + + fmha.forward(q, None, None, metadata, forward_args) + + run_context.assert_called_once() + context_params = context_calls[0] + assert context_params.num_tokens == 3 + assert context_params.seq_offset == 0 + assert context_params.batch_size == 1 + assert context_params.num_requests == 1 + assert context_params.attention_input is not None + assert context_params.attention_input.shape[0] == 3 + assert context_params.context_buf is not None + assert context_params.context_buf.shape == (3, attn.num_heads, attn.head_dim) + + run_generation.assert_called_once() + generation_params = generation_calls[0] + assert generation_params.num_tokens == 2 + assert generation_params.seq_offset == 1 + assert generation_params.batch_size == 2 + assert generation_params.num_requests == 2 + assert generation_params.attention_input is not None + assert generation_params.attention_input.shape[0] == 2 + assert generation_params.context_buf is not None + assert generation_params.context_buf.shape == (2, attn.num_heads, attn.head_dim) diff --git a/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py new file mode 100644 index 000000000000..eb9ea6941792 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from importlib import import_module + +import pytest + +cute = pytest.importorskip("cutlass.cute") + +from tensorrt_llm._torch.visual_gen.attention_backend import flash_attn4, parallel # noqa: E402 +from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( # noqa: E402 + _install_cutlass_dsl_compatibility, +) + + +def test_cutlass_dsl_47_moved_names_are_restored(monkeypatch): + monkeypatch.delattr(cute.core, "ThrCopy", raising=False) + monkeypatch.delattr(cute.core, "ThrMma", raising=False) + monkeypatch.delattr(cute, "make_fragment", raising=False) + + _install_cutlass_dsl_compatibility() + + assert cute.core.ThrCopy is cute.ThrCopy + assert cute.core.ThrMma is cute.ThrMma + assert cute.make_fragment is cute.make_rmem_tensor + + +def test_cutlass_dsl_existing_names_are_preserved(monkeypatch): + existing_thr_copy = object() + existing_thr_mma = object() + existing_make_fragment = object() + monkeypatch.setattr(cute.core, "ThrCopy", existing_thr_copy) + monkeypatch.setattr(cute.core, "ThrMma", existing_thr_mma) + monkeypatch.setattr(cute, "make_fragment", existing_make_fragment) + + _install_cutlass_dsl_compatibility() + + assert cute.core.ThrCopy is existing_thr_copy + assert cute.core.ThrMma is existing_thr_mma + assert cute.make_fragment is existing_make_fragment + + +def test_cutlass_dsl_47_aliases_allow_fa4_interface_import() -> None: + _install_cutlass_dsl_compatibility() + interface = import_module("flash_attn.cute.interface") + + assert callable(interface._flash_attn_fwd) + assert callable(interface.flash_attn_combine) + assert callable(flash_attn4._flash_attn_fwd) + assert callable(parallel._flash_attn_combine) diff --git a/tests/unittest/others/test_vendor_sources.py b/tests/unittest/others/test_vendor_sources.py new file mode 100644 index 000000000000..11526714f574 --- /dev/null +++ b/tests/unittest/others/test_vendor_sources.py @@ -0,0 +1,1898 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Hermetic lifecycle tests for the generic source-vendoring tool.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import os +import shutil +import stat +import subprocess +import sys +import time +from pathlib import Path +from types import ModuleType + +import pytest +import yaml + +pytestmark = pytest.mark.cpu_only + + +_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +_VENDOR_SOURCES = _REPO_ROOT / "scripts" / "vendor_sources.py" +_VENDOR_NAME = "example" +_SOURCE = "python/example" +_DESTINATION = "src/example" +_INCLUDE = "**/*.py" + + +def _command_env() -> dict[str, str]: + env = os.environ.copy() + env.update( + { + "GIT_CONFIG_NOSYSTEM": "1", + "GIT_TERMINAL_PROMPT": "0", + "LC_ALL": "C.UTF-8", + } + ) + return env + + +def _run( + command: list[str | Path], + *, + cwd: Path, + check: bool = True, + env_overrides: dict[str, str] | None = None, +) -> subprocess.CompletedProcess[str]: + env = _command_env() + if env_overrides is not None: + env.update(env_overrides) + result = subprocess.run( + [str(argument) for argument in command], + cwd=cwd, + env=env, + capture_output=True, + text=True, + timeout=30, + ) + if check and result.returncode != 0: + pytest.fail( + f"Command failed ({result.returncode}): {' '.join(map(str, command))}\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + return result + + +def _git(repo: Path, *arguments: str) -> str: + result = _run(["git", "-C", repo, *arguments], cwd=repo) + return result.stdout.strip() + + +def _write_files(root: Path, files: dict[str, str]) -> None: + for relative_path, content in files.items(): + path = root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + +def _commit(repo: Path, message: str) -> str: + _git(repo, "add", "-A") + _git(repo, "commit", "-m", message) + return _git(repo, "rev-parse", "HEAD") + + +def _make_upstream(tmp_path: Path, files: dict[str, str]) -> tuple[Path, str]: + upstream = tmp_path / "upstream" + upstream.mkdir() + _run(["git", "init", "--initial-branch=main", upstream], cwd=tmp_path) + _git(upstream, "config", "user.name", "Vendor Sources Test") + _git(upstream, "config", "user.email", "vendor-sources@example.invalid") + _write_files(upstream / _SOURCE, files) + return upstream, _commit(upstream, "initial upstream source") + + +def _make_consumer(tmp_path: Path) -> tuple[Path, Path]: + consumer = tmp_path / "consumer" + lock = consumer / "3rdparty" / "vendor-sources.yml" + lock.parent.mkdir(parents=True) + return consumer, lock + + +def _vendor( + consumer: Path, + lock: Path, + *arguments: str | Path, + check: bool = True, + env_overrides: dict[str, str] | None = None, +) -> subprocess.CompletedProcess[str]: + return _run( + [sys.executable, _VENDOR_SOURCES, "--lock", lock, *arguments], + cwd=consumer, + check=check, + env_overrides=env_overrides, + ) + + +def _copy_python_sources(upstream: Path, consumer: Path) -> None: + source = upstream / _SOURCE + destination = consumer / _DESTINATION + for source_path in source.rglob("*.py"): + relative_path = source_path.relative_to(source) + target = destination / relative_path + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_path, target) + + +def _create_vendor( + consumer: Path, + lock: Path, + upstream: Path, + commit: str, + *, + mode: str, + url: str | None = None, +) -> None: + _vendor( + consumer, + lock, + "create", + _VENDOR_NAME, + "--url", + url or upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + upstream, + "--adopt", + mode, + ) + + +def _lock_data(lock: Path) -> dict[str, object]: + data = yaml.safe_load(lock.read_text(encoding="utf-8")) + assert isinstance(data, dict) + return data + + +def _vendor_data(lock: Path) -> dict[str, object]: + vendors = _lock_data(lock)["vendors"] + assert isinstance(vendors, dict) + vendor = vendors[_VENDOR_NAME] + assert isinstance(vendor, dict) + return vendor + + +def _assert_failure(result: subprocess.CompletedProcess[str], text: str) -> None: + assert result.returncode != 0 + combined_output = f"{result.stdout}\n{result.stderr}".lower() + assert text.lower() in combined_output + + +def _tree_snapshot(root: Path) -> dict[str, tuple[bytes, int]]: + if not root.exists(): + return {} + return { + path.relative_to(root).as_posix(): ( + path.read_bytes(), + stat.S_IMODE(path.stat().st_mode), + ) + for path in sorted(root.rglob("*")) + if path.is_file() + } + + +def _set_patch_digest(lock: Path, content: bytes) -> None: + data = _lock_data(lock) + vendors = data["vendors"] + assert isinstance(vendors, dict) + vendor = vendors[_VENDOR_NAME] + assert isinstance(vendor, dict) + vendor["patch_digest"] = f"sha256:{hashlib.sha256(content).hexdigest()}" + lock.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + + +def _load_vendor_sources_module() -> ModuleType: + module_name = "_vendor_sources_under_test" + existing = sys.modules.get(module_name) + if existing is not None: + return existing + spec = importlib.util.spec_from_file_location(module_name, _VENDOR_SOURCES) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def test_atomic_write_syncs_file_before_replacement( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_vendor_sources_module() + target = tmp_path / "lock" / "vendor-sources.yml" + target.parent.mkdir() + events: list[str] = [] + real_fsync = os.fsync + real_replace = os.replace + + def record_fsync(file_descriptor: int) -> None: + mode = os.fstat(file_descriptor).st_mode + events.append("directory-fsync" if stat.S_ISDIR(mode) else "file-fsync") + real_fsync(file_descriptor) + + def record_replace(source: str | bytes | Path, destination: str | bytes | Path) -> None: + events.append("replace") + real_replace(source, destination) + + with monkeypatch.context() as recording: + recording.setattr(module.os, "fsync", record_fsync) + recording.setattr(module.os, "replace", record_replace) + module._atomic_write(target, b"updated\n") + + assert target.read_bytes() == b"updated\n" + assert events == ["file-fsync", "replace"] + + +def test_create_and_patch_updates_do_not_sync_replacement_directories( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + exact_root = tmp_path / "exact-case" + exact_root.mkdir() + exact_consumer, exact_lock = _make_consumer(exact_root) + patched_root = tmp_path / "patched-case" + patched_root.mkdir() + patched_consumer, patched_lock = _make_consumer(patched_root) + patched_destination = patched_consumer / _DESTINATION + _write_files(patched_destination, {"kernel.py": "VALUE = 'adopted'\n"}) + module = _load_vendor_sources_module() + real_fsync = os.fsync + directory_syncs = 0 + + def reject_directory_sync(file_descriptor: int) -> None: + nonlocal directory_syncs + if stat.S_ISDIR(os.fstat(file_descriptor).st_mode): + directory_syncs += 1 + raise OSError("unexpected replacement-directory sync") + real_fsync(file_descriptor) + + def create_arguments(lock: Path, *, adopt: str | None = None) -> list[str]: + arguments = [ + "--lock", + str(lock), + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + str(upstream), + ] + if adopt is not None: + arguments.extend(["--adopt", adopt]) + return arguments + + with monkeypatch.context() as failure: + failure.setattr(module.os, "fsync", reject_directory_sync) + assert module.main(create_arguments(exact_lock)) == 0 + assert module.main(create_arguments(patched_lock, adopt="patched")) == 0 + + exact_destination = exact_consumer / _DESTINATION + (exact_destination / "kernel.py").write_text("VALUE = 'patched'\n", encoding="utf-8") + assert ( + module.main( + [ + "--lock", + str(exact_lock), + "patch", + _VENDOR_NAME, + "create", + "--repo", + str(upstream), + ] + ) + == 0 + ) + (exact_destination / "kernel.py").write_text("VALUE = 'refreshed'\n", encoding="utf-8") + assert ( + module.main( + [ + "--lock", + str(exact_lock), + "patch", + _VENDOR_NAME, + "refresh", + "--repo", + str(upstream), + ] + ) + == 0 + ) + + assert directory_syncs == 0 + _vendor(exact_consumer, exact_lock, "check", _VENDOR_NAME, "--repo", upstream) + _vendor(patched_consumer, patched_lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_exact_create_check_digest_and_sync(tmp_path: Path) -> None: + upstream, commit = _make_upstream( + tmp_path, + { + "__init__.py": "VALUE = 1\n", + "nested/kernel.py": "def kernel() -> int:\n return 1\n", + "README.md": "not selected\n", + }, + ) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + + _create_vendor(consumer, lock, upstream, commit, mode="exact") + + vendor = _vendor_data(lock) + assert vendor["commit"] == commit + assert str(vendor["digest"]).startswith("sha256-tree-v1:") + assert not (consumer / _DESTINATION / "README.md").exists() + _vendor(consumer, lock, "check", _VENDOR_NAME) + + fake_bin = tmp_path / "fake-bin" + fake_bin.mkdir() + sentinel = tmp_path / "git-was-invoked" + fake_git = fake_bin / "git" + fake_git.write_text( + '#!/bin/sh\nprintf invoked > "$VENDOR_GIT_SENTINEL"\nexit 97\n', + encoding="utf-8", + ) + fake_git.chmod(0o755) + fake_git_environment = { + "PATH": str(fake_bin), + "VENDOR_GIT_SENTINEL": str(sentinel), + } + for mode in ([], ["--offline"]): + _vendor( + consumer, + lock, + "check", + _VENDOR_NAME, + *mode, + env_overrides=fake_git_environment, + ) + assert not sentinel.exists(), "offline checks must not invoke Git" + invalid_mixed_mode = _vendor( + consumer, + lock, + "check", + _VENDOR_NAME, + "--offline", + "--repo", + upstream, + check=False, + env_overrides=fake_git_environment, + ) + _assert_failure(invalid_mixed_mode, "either") + assert not sentinel.exists(), "rejected offline checks must not invoke Git" + + kernel = consumer / _DESTINATION / "nested" / "kernel.py" + kernel.write_text("def kernel() -> int:\n return 99\n", encoding="utf-8") + _assert_failure(_vendor(consumer, lock, "check", _VENDOR_NAME, check=False), "digest") + + _vendor(consumer, lock, "sync", _VENDOR_NAME, "--repo", upstream) + assert kernel.read_text(encoding="utf-8").endswith("return 1\n") + _vendor(consumer, lock, "check", _VENDOR_NAME) + + +def test_create_materializes_exact_tree_and_remote_mismatch_is_fatal(tmp_path: Path) -> None: + upstream, first_commit = _make_upstream( + tmp_path, + { + "kernel.py": "VALUE = 1\n", + "nested/helper.py": "HELPER = True\n", + "README.md": "not selected\n", + }, + ) + consumer, lock = _make_consumer(tmp_path) + + _vendor( + consumer, + lock, + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--branch", + "main", + "--commit", + first_commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + upstream, + ) + + destination = consumer / _DESTINATION + assert (destination / "kernel.py").read_text(encoding="utf-8") == "VALUE = 1\n" + assert (destination / "nested" / "helper.py").is_file() + assert not (destination / "README.md").exists() + _vendor(consumer, lock, "check", _VENDOR_NAME, "--upstream") + + (upstream / _SOURCE / "kernel.py").write_text("VALUE = 2\n", encoding="utf-8") + second_commit = _commit(upstream, "change accessible upstream source") + original_lock = lock.read_text(encoding="utf-8") + assert first_commit in original_lock + lock.write_text(original_lock.replace(first_commit, second_commit), encoding="utf-8") + + # Offline integrity still holds because only source provenance changed. + _vendor(consumer, lock, "check", _VENDOR_NAME) + remote_check = _vendor( + consumer, + lock, + "check", + _VENDOR_NAME, + "--upstream", + check=False, + ) + _assert_failure(remote_check, "materialization digest mismatch") + assert "skip-remote-unavailable" not in f"{remote_check.stdout}\n{remote_check.stderr}".lower() + + +def test_git_subprocesses_ignore_ambient_repository_environment(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + _create_vendor(consumer, lock, upstream, commit, mode="exact") + destination = consumer / _DESTINATION + (destination / "kernel.py").write_text("VALUE = 'downstream'\n", encoding="utf-8") + + upstream_head = _git(upstream, "rev-parse", "HEAD") + upstream_branch = _git(upstream, "symbolic-ref", "HEAD") + upstream_status = _git(upstream, "status", "--porcelain=v1", "--untracked-files=all") + upstream_index = (upstream / ".git" / "index").read_bytes() + upstream_source = _tree_snapshot(upstream / _SOURCE) + + _vendor( + consumer, + lock, + "patch", + _VENDOR_NAME, + "create", + "--repo", + upstream, + env_overrides={ + "GIT_DIR": str(upstream / ".git"), + "GIT_WORK_TREE": str(upstream), + "GIT_INDEX_FILE": str(upstream / ".git" / "index"), + "GIT_OBJECT_DIRECTORY": str(upstream / ".git" / "objects"), + }, + ) + + assert _git(upstream, "rev-parse", "HEAD") == upstream_head + assert _git(upstream, "symbolic-ref", "HEAD") == upstream_branch + assert (upstream / ".git" / "index").read_bytes() == upstream_index + assert _git(upstream, "status", "--porcelain=v1", "--untracked-files=all") == upstream_status + assert _tree_snapshot(upstream / _SOURCE) == upstream_source + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_exact_materialization_ignores_git_export_attributes(tmp_path: Path) -> None: + marker = 'REVISION = "$Format:%H$"\n' + upstream, _ = _make_upstream( + tmp_path, + { + ".gitattributes": "ignored.py export-ignore\nsubstituted.py export-subst\n", + "ignored.py": "IGNORED_BY_ARCHIVE = True\n", + "substituted.py": marker, + }, + ) + (upstream / _SOURCE / "ignored.py").chmod(0o755) + commit = _commit(upstream, "record executable mode") + consumer, lock = _make_consumer(tmp_path) + + _vendor( + consumer, + lock, + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + upstream, + ) + + destination = consumer / _DESTINATION + assert (destination / "ignored.py").read_text(encoding="utf-8") == ( + "IGNORED_BY_ARCHIVE = True\n" + ) + assert (destination / "ignored.py").stat().st_mode & stat.S_IXUSR + assert (destination / "substituted.py").read_text(encoding="utf-8") == marker + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--upstream", "--require-access") + + +def test_patched_adoption_preserves_raw_crlf_with_selected_text_attribute( + tmp_path: Path, +) -> None: + upstream, _ = _make_upstream(tmp_path, {"value.txt": "seed\n"}) + _git(upstream, "config", "core.autocrlf", "false") + upstream_value = upstream / _SOURCE / "value.txt" + upstream_value.write_bytes(b"upstream\r\n") + _commit(upstream, "record raw CRLF blob") + + attributes = upstream / _SOURCE / ".gitattributes" + attributes.write_text("*.txt text\n", encoding="utf-8") + _git(upstream, "add", "--", f"{_SOURCE}/.gitattributes") + _git(upstream, "commit", "-m", "select text attribute") + commit = _git(upstream, "rev-parse", "HEAD") + committed_blob = _git(upstream, "rev-parse", f"{commit}:{_SOURCE}/value.txt") + raw_worktree_blob = _git(upstream, "hash-object", "--no-filters", f"{_SOURCE}/value.txt") + assert committed_blob == raw_worktree_blob + + consumer, lock = _make_consumer(tmp_path) + destination = consumer / _DESTINATION + shutil.copytree(upstream / _SOURCE, destination) + destination_value = destination / "value.txt" + destination_value.write_bytes(b"downstream\r\n") + destination_value.chmod(0o755) + accepted_tree = _tree_snapshot(destination) + assert accepted_tree["value.txt"] == (b"downstream\r\n", 0o755) + + _vendor( + consumer, + lock, + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--repo", + upstream, + "--adopt", + "patched", + ) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + destination_value.write_bytes(b"broken\n") + destination_value.chmod(0o644) + _vendor(consumer, lock, "sync", _VENDOR_NAME, "--repo", upstream) + assert _tree_snapshot(destination) == accepted_tree + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_local_replacement_refs_do_not_change_locked_content(tmp_path: Path) -> None: + upstream, first_commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 1\n"}) + (upstream / _SOURCE / "kernel.py").write_text("VALUE = 2\n", encoding="utf-8") + second_commit = _commit(upstream, "replacement content") + _git(upstream, "replace", first_commit, second_commit) + consumer, lock = _make_consumer(tmp_path) + + _vendor( + consumer, + lock, + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + first_commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + upstream, + ) + + assert (consumer / _DESTINATION / "kernel.py").read_text(encoding="utf-8") == "VALUE = 1\n" + assert _vendor_data(lock)["commit"] == first_commit + _git(upstream, "replace", "-d", first_commit) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_lock_rejects_duplicate_keys_and_unsafe_paths(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 1\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + _create_vendor(consumer, lock, upstream, commit, mode="exact") + + valid_lock = lock.read_text(encoding="utf-8") + first_data_line = next( + line + for line in valid_lock.splitlines() + if line.strip() and not line.lstrip().startswith("#") + ) + lock.write_text(f"{first_data_line}\n{valid_lock}", encoding="utf-8") + _assert_failure(_vendor(consumer, lock, "check", check=False), "duplicate") + + lock.write_text(valid_lock, encoding="utf-8") + vendor_marker = f" {_VENDOR_NAME}:\n" + assert vendor_marker in valid_lock + lock.write_text( + valid_lock.replace(vendor_marker, f"{vendor_marker} unexpected: true\n", 1), + encoding="utf-8", + ) + _assert_failure(_vendor(consumer, lock, "check", check=False), "unexpected") + + lock.write_text( + valid_lock.replace(vendor_marker, f"{vendor_marker} divergence: legacy\n", 1), + encoding="utf-8", + ) + _assert_failure(_vendor(consumer, lock, "check", check=False), "divergence") + + lock.write_text(valid_lock, encoding="utf-8") + unsafe_create = _vendor( + consumer, + lock, + "create", + "unsafe", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + "../outside", + "--repo", + upstream, + check=False, + ) + _assert_failure(unsafe_create, "destination") + assert not (tmp_path / "outside").exists() + + unsafe_source = _vendor( + consumer, + lock, + "create", + "unsafe-source", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + "../outside", + "--destination", + "src/unsafe-source", + "--repo", + upstream, + check=False, + ) + _assert_failure(unsafe_source, "source") + + existing_destination = consumer / "src" / "existing" + _write_files(existing_destination, {"kernel.py": "DO_NOT_REPLACE = True\n"}) + accidental_overwrite = _vendor( + consumer, + lock, + "create", + "existing", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + "src/existing", + "--include", + _INCLUDE, + "--repo", + upstream, + check=False, + ) + _assert_failure(accidental_overwrite, "already exists") + assert (existing_destination / "kernel.py").read_text(encoding="utf-8") == ( + "DO_NOT_REPLACE = True\n" + ) + + git_destination = _vendor( + consumer, + lock, + "create", + "git-internals", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + ".git/vendor", + "--repo", + upstream, + check=False, + ) + _assert_failure(git_destination, "destination") + + unsafe_tree = consumer / "src" / "unsafe-tree" + _write_files( + unsafe_tree, + { + "kernel.py": "VALUE = 1\n", + ".git/config": "[core]\n\thooksPath = /tmp/unsafe\n", + }, + ) + unsafe_tree_create = _vendor( + consumer, + lock, + "create", + "unsafe-tree", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + "src/unsafe-tree", + "--repo", + upstream, + "--adopt", + "patched", + check=False, + ) + _assert_failure(unsafe_tree_create, "unsafe path") + + outside_destination = tmp_path / "outside-destination" + _write_files(outside_destination, {"kernel.py": "VALUE = 1\n"}) + linked_destination = consumer / "src" / "linked-destination" + linked_destination.symlink_to(outside_destination, target_is_directory=True) + symlink_escape = _vendor( + consumer, + lock, + "create", + "unsafe-symlink", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + "src/linked-destination", + "--include", + _INCLUDE, + "--repo", + upstream, + "--adopt", + "exact", + check=False, + ) + _assert_failure(symlink_escape, "destination") + assert (outside_destination / "kernel.py").read_text(encoding="utf-8") == "VALUE = 1\n" + + +def test_branch_tag_validation_and_overlapping_destinations(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 1\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + _create_vendor(consumer, lock, upstream, commit, mode="exact") + + valid_lock = lock.read_text(encoding="utf-8") + url_line = next( + line for line in valid_lock.splitlines(keepends=True) if line.startswith(" url:") + ) + lock.write_text( + valid_lock.replace(url_line, f"{url_line} branch: main\n tag: v1.0\n", 1), + encoding="utf-8", + ) + _assert_failure(_vendor(consumer, lock, "check", check=False), "both branch and tag") + lock.write_text(valid_lock, encoding="utf-8") + + common_arguments: list[str | Path] = [ + "create", + "second", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + "src/second", + "--repo", + upstream, + ] + conflicting_reference = _vendor( + consumer, + lock, + *common_arguments, + "--branch", + "main", + "--tag", + "v1.0", + check=False, + ) + _assert_failure(conflicting_reference, "not allowed with argument") + + long_reference = _vendor( + consumer, + lock, + *common_arguments, + "--branch", + "refs/heads/main", + check=False, + ) + _assert_failure(long_reference, "short Git name") + + overlapping = _vendor( + consumer, + lock, + "create", + "overlap", + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + f"{_DESTINATION}/nested", + "--repo", + upstream, + check=False, + ) + _assert_failure(overlapping, "overlap") + assert set(_lock_data(lock)["vendors"]) == {_VENDOR_NAME} + + +def test_patched_vendor_reproduces_added_modified_and_deleted_files(tmp_path: Path) -> None: + upstream, commit = _make_upstream( + tmp_path, + { + "keep.py": "KEEP = True\n", + "modify.py": "VALUE = 'upstream'\n", + "delete.py": "DELETE_ME = True\n", + }, + ) + consumer, lock = _make_consumer(tmp_path) + destination = consumer / _DESTINATION + _write_files( + destination, + { + "keep.py": "KEEP = True\n", + "modify.py": "VALUE = 'downstream'\n", + "add.py": "ADDED = True\n", + }, + ) + (destination / "modify.py").chmod(0o755) + + _create_vendor(consumer, lock, upstream, commit, mode="patched") + + vendor = _vendor_data(lock) + patch_path = consumer / str(vendor["patch"]) + assert patch_path.is_file() + original_patch = patch_path.read_bytes() + original_patch_mode = patch_path.stat().st_mode & 0o777 + patch_text = original_patch.decode("utf-8") + assert all(filename in patch_text for filename in ("add.py", "modify.py", "delete.py")) + assert "new mode 100755" in patch_text + + valid_lock = lock.read_text(encoding="utf-8") + lock.write_text( + valid_lock.replace(str(vendor["patch"]), "../outside.patch"), + encoding="utf-8", + ) + _assert_failure(_vendor(consumer, lock, "check", _VENDOR_NAME, check=False), "patch") + lock.write_text(valid_lock, encoding="utf-8") + + external_patch = tmp_path / "external.patch" + external_patch.write_bytes(original_patch) + patch_path.unlink() + patch_path.symlink_to(external_patch) + _assert_failure(_vendor(consumer, lock, "check", _VENDOR_NAME, check=False), "patch") + patch_path.unlink() + patch_path.write_bytes(original_patch) + patch_path.chmod(original_patch_mode) + + assert b"a/modify.py" in original_patch + malicious_patch = original_patch.replace(b"a/modify.py", b"a/../outside.py").replace( + b"b/modify.py", b"b/../outside.py" + ) + patch_path.write_bytes(malicious_patch) + _set_patch_digest(lock, malicious_patch) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--offline") + malicious_check = _vendor( + consumer, + lock, + "check", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + env_overrides={"TMPDIR": str(tmp_path)}, + ) + _assert_failure(malicious_check, "failed to apply vendor patch") + assert not (tmp_path / "outside.py").exists() + + outside_include_patch = b"""diff --git a/outside.txt b/outside.txt +new file mode 100644 +--- /dev/null ++++ b/outside.txt +@@ -0,0 +1 @@ ++outside include +""" + patch_path.write_bytes(outside_include_patch) + _set_patch_digest(lock, outside_include_patch) + outside_include_check = _vendor( + consumer, + lock, + "check", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + ) + _assert_failure(outside_include_check, "outside the vendor include set") + + patch_path.write_bytes(original_patch) + patch_path.chmod(original_patch_mode) + lock.write_text(valid_lock, encoding="utf-8") + + (destination / "add.py").unlink() + (destination / "modify.py").write_text("BROKEN = True\n", encoding="utf-8") + (destination / "modify.py").chmod(0o644) + (destination / "delete.py").write_text("STALE = True\n", encoding="utf-8") + _vendor(consumer, lock, "sync", _VENDOR_NAME, "--repo", upstream) + + assert (destination / "keep.py").read_text(encoding="utf-8") == "KEEP = True\n" + assert (destination / "modify.py").read_text(encoding="utf-8") == "VALUE = 'downstream'\n" + assert destination.joinpath("modify.py").stat().st_mode & 0o100 + assert (destination / "add.py").read_text(encoding="utf-8") == "ADDED = True\n" + assert not (destination / "delete.py").exists() + _vendor(consumer, lock, "check", _VENDOR_NAME) + + # A refresh is also the recovery path if a prior update was interrupted + # after replacing the patch but before updating its recorded digest. + patch_path.write_bytes(original_patch + b"\n") + (destination / "add.py").write_text("ADDED = 'refreshed'\n", encoding="utf-8") + _vendor(consumer, lock, "patch", _VENDOR_NAME, "refresh", "--repo", upstream) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + accepted_tree = _tree_snapshot(destination) + + (destination / "add.py").unlink() + (destination / "modify.py").write_text("BROKEN_AGAIN = True\n", encoding="utf-8") + (destination / "modify.py").chmod(0o644) + (destination / "delete.py").write_text("STALE_AGAIN = True\n", encoding="utf-8") + _vendor(consumer, lock, "sync", _VENDOR_NAME, "--repo", upstream) + assert _tree_snapshot(destination) == accepted_tree + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_patch_application_ignores_enclosing_git_worktree(tmp_path: Path) -> None: + enclosing = tmp_path / "enclosing" + enclosing.mkdir() + _run(["git", "init", "--initial-branch=main", enclosing], cwd=tmp_path) + case_root = enclosing / "case" + case_root.mkdir() + upstream, commit = _make_upstream(case_root, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(case_root) + destination = consumer / _DESTINATION + _write_files(destination, {"kernel.py": "VALUE = 'downstream'\n"}) + _create_vendor(consumer, lock, upstream, commit, mode="patched") + + temporary_root = consumer / "tmp" + temporary_root.mkdir() + assert _git(temporary_root, "rev-parse", "--show-toplevel") == str(enclosing) + + _vendor( + consumer, + lock, + "check", + _VENDOR_NAME, + "--repo", + upstream, + env_overrides={"TMPDIR": str(temporary_root)}, + ) + assert (destination / "kernel.py").read_text(encoding="utf-8") == "VALUE = 'downstream'\n" + + +def test_exact_adoption_rejects_unrepresented_differences(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + destination = consumer / _DESTINATION + _write_files(destination, {"kernel.py": "VALUE = 'different'\n"}) + destination_snapshot = _tree_snapshot(destination) + + adoption = _vendor( + consumer, + lock, + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + upstream, + "--adopt", + "exact", + check=False, + ) + + _assert_failure(adoption, "is not exact") + assert "kernel.py" in f"{adoption.stdout}\n{adoption.stderr}" + assert _tree_snapshot(destination) == destination_snapshot + assert not lock.exists() + assert not (consumer / "3rdparty" / "vendor_patches" / "example.patch").exists() + + +def test_create_rolls_back_destination_and_patch_when_lock_save_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], +) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + exact_root = tmp_path / "exact-case" + exact_root.mkdir() + exact_consumer, exact_lock = _make_consumer(exact_root) + patched_root = tmp_path / "patched-case" + patched_root.mkdir() + patched_consumer, patched_lock = _make_consumer(patched_root) + patched_destination = patched_consumer / _DESTINATION + _write_files(patched_destination, {"kernel.py": "VALUE = 'downstream'\n"}) + patched_snapshot = _tree_snapshot(patched_destination) + module = _load_vendor_sources_module() + + def fail_save_lock(_: object) -> None: + raise OSError("injected lock-save failure") + + with monkeypatch.context() as failure: + failure.setattr(module, "_save_lock", fail_save_lock) + exact_result = module.main( + [ + "--lock", + str(exact_lock), + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + str(upstream), + ] + ) + patched_result = module.main( + [ + "--lock", + str(patched_lock), + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + str(upstream), + "--adopt", + "patched", + ] + ) + + assert exact_result == 1 + assert patched_result == 1 + captured = capfd.readouterr() + assert "error:" in captured.err.lower() + assert "injected lock-save failure" in captured.err + assert not exact_lock.exists() + assert not (exact_consumer / _DESTINATION).exists() + assert not list((exact_consumer / "src").glob(".example.vendor-*")) + assert not patched_lock.exists() + assert _tree_snapshot(patched_destination) == patched_snapshot + assert not (patched_consumer / "3rdparty/vendor_patches/example.patch").exists() + assert not list(patched_destination.parent.glob(".example.vendor-*")) + + _vendor( + exact_consumer, + exact_lock, + "create", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--commit", + commit, + "--source", + _SOURCE, + "--destination", + _DESTINATION, + "--include", + _INCLUDE, + "--repo", + upstream, + ) + _create_vendor(patched_consumer, patched_lock, upstream, commit, mode="patched") + + +def test_stale_backup_does_not_leak_staging_or_modify_trees(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + _create_vendor(consumer, lock, upstream, commit, mode="exact") + destination = consumer / _DESTINATION + backup = destination.parent / ".example.vendor-backup" + _write_files(backup, {"sentinel.py": "DO_NOT_TOUCH = True\n"}) + destination_snapshot = _tree_snapshot(destination) + backup_snapshot = _tree_snapshot(backup) + lock_snapshot = lock.read_bytes() + parent_entries = sorted(path.name for path in destination.parent.iterdir()) + + sync = _vendor( + consumer, + lock, + "sync", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + ) + + _assert_failure(sync, "stale vendor backup") + assert sorted(path.name for path in destination.parent.iterdir()) == parent_entries + assert _tree_snapshot(destination) == destination_snapshot + assert _tree_snapshot(backup) == backup_snapshot + assert lock.read_bytes() == lock_snapshot + + +def test_patch_drop_and_remove_preserve_destination_and_unrelated_patch(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + destination = consumer / _DESTINATION + _write_files(destination, {"kernel.py": "VALUE = 'downstream'\n"}) + _create_vendor(consumer, lock, upstream, commit, mode="patched") + patch_path = consumer / str(_vendor_data(lock)["patch"]) + rejected_lock = lock.read_bytes() + rejected_patch = patch_path.read_bytes() + rejected_destination = _tree_snapshot(destination) + + rejected_drop = _vendor( + consumer, + lock, + "patch", + _VENDOR_NAME, + "drop", + "--repo", + upstream, + check=False, + ) + _assert_failure(rejected_drop, "not exact upstream") + assert lock.read_bytes() == rejected_lock + assert patch_path.read_bytes() == rejected_patch + assert _tree_snapshot(destination) == rejected_destination + + (destination / "kernel.py").write_text("VALUE = 'upstream'\n", encoding="utf-8") + exact_snapshot = _tree_snapshot(destination) + unrelated_patch = patch_path.parent / "unrelated.patch" + unrelated_patch.write_bytes(b"unrelated sentinel\n") + _vendor(consumer, lock, "patch", _VENDOR_NAME, "drop", "--repo", upstream) + vendor = _vendor_data(lock) + assert "patch" not in vendor + assert "patch_digest" not in vendor + assert not patch_path.exists() + assert unrelated_patch.read_bytes() == b"unrelated sentinel\n" + assert _tree_snapshot(destination) == exact_snapshot + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + (destination / "kernel.py").write_text("VALUE = 'downstream-again'\n", encoding="utf-8") + _vendor(consumer, lock, "patch", _VENDOR_NAME, "create", "--repo", upstream) + recreated_patch = consumer / str(_vendor_data(lock)["patch"]) + removal_snapshot = _tree_snapshot(destination) + _vendor(consumer, lock, "remove", _VENDOR_NAME) + assert _lock_data(lock)["vendors"] == {} + assert not recreated_patch.exists() + assert unrelated_patch.read_bytes() == b"unrelated sentinel\n" + assert _tree_snapshot(destination) == removal_snapshot + + +def test_export_rejects_dirty_source_and_mismatched_head_without_changes(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + _create_vendor(consumer, lock, upstream, commit, mode="exact") + destination = consumer / _DESTINATION + (destination / "kernel.py").write_text("VALUE = 'downstream'\n", encoding="utf-8") + _assert_failure( + _vendor(consumer, lock, "check", _VENDOR_NAME, "--offline", check=False), + "digest mismatch", + ) + status = _vendor(consumer, lock, "status", _VENDOR_NAME, check=False) + _assert_failure(status, "fail-local-integrity") + lock_snapshot = lock.read_bytes() + destination_snapshot = _tree_snapshot(destination) + + upstream_kernel = upstream / _SOURCE / "kernel.py" + upstream_kernel.write_text("VALUE = 'dirty'\n", encoding="utf-8") + dirty_source_snapshot = _tree_snapshot(upstream / _SOURCE) + dirty_export = _vendor( + consumer, + lock, + "export", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + ) + _assert_failure(dirty_export, "uncommitted changes") + assert _tree_snapshot(upstream / _SOURCE) == dirty_source_snapshot + assert lock.read_bytes() == lock_snapshot + assert _tree_snapshot(destination) == destination_snapshot + + _git(upstream, "restore", "--", _SOURCE) + _write_files(upstream, {"unrelated.txt": "new head\n"}) + _commit(upstream, "unrelated upstream change") + mismatched_source_snapshot = _tree_snapshot(upstream / _SOURCE) + mismatched_export = _vendor( + consumer, + lock, + "export", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + ) + _assert_failure(mismatched_export, "head must equal locked commit") + assert _tree_snapshot(upstream / _SOURCE) == mismatched_source_snapshot + assert lock.read_bytes() == lock_snapshot + assert _tree_snapshot(destination) == destination_snapshot + + +def test_export_rejects_ignored_selected_source_files_without_changes(tmp_path: Path) -> None: + upstream, commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + _create_vendor(consumer, lock, upstream, commit, mode="exact") + destination = consumer / _DESTINATION + (destination / "kernel.py").write_text("VALUE = 'downstream'\n", encoding="utf-8") + ignored_path = upstream / _SOURCE / "generated.py" + ignored_path.write_text("GENERATED = True\n", encoding="utf-8") + (upstream / ".git" / "info" / "exclude").write_text( + f"/{_SOURCE}/generated.py\n", encoding="utf-8" + ) + assert _git(upstream, "status", "--porcelain", "--", _SOURCE) == "" + upstream_snapshot = _tree_snapshot(upstream / _SOURCE) + lock_snapshot = lock.read_bytes() + destination_snapshot = _tree_snapshot(destination) + + result = _vendor( + consumer, + lock, + "export", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + ) + + _assert_failure(result, "ignored untracked files") + assert _tree_snapshot(upstream / _SOURCE) == upstream_snapshot + assert ignored_path.read_text(encoding="utf-8") == "GENERATED = True\n" + assert lock.read_bytes() == lock_snapshot + assert _tree_snapshot(destination) == destination_snapshot + + +def test_upstream_access_policy_export_and_pin(tmp_path: Path) -> None: + upstream, first_commit = _make_upstream(tmp_path, {"kernel.py": "VALUE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + inaccessible_url = (tmp_path / "inaccessible-upstream.git").as_uri() + _create_vendor( + consumer, + lock, + upstream, + first_commit, + mode="exact", + url=inaccessible_url, + ) + + best_effort = _vendor(consumer, lock, "check", _VENDOR_NAME, "--upstream") + assert "unavailable" in f"{best_effort.stdout}\n{best_effort.stderr}".lower() + strict = _vendor( + consumer, + lock, + "check", + _VENDOR_NAME, + "--upstream", + "--require-access", + check=False, + ) + _assert_failure(strict, "unavailable") + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + no_change = _vendor( + consumer, + lock, + "export", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + ) + _assert_failure(no_change, "no downstream change") + + kernel = consumer / _DESTINATION / "kernel.py" + locked_state = lock.read_bytes() + kernel.write_text("VALUE = 'exported-fix'\n", encoding="utf-8") + pending_tree = _tree_snapshot(consumer / _DESTINATION) + _assert_failure( + _vendor(consumer, lock, "check", _VENDOR_NAME, "--offline", check=False), + "digest mismatch", + ) + + _git(upstream, "switch", "-c", "vendor-fix") + _vendor(consumer, lock, "export", _VENDOR_NAME, "--repo", upstream) + assert lock.read_bytes() == locked_state + assert _tree_snapshot(consumer / _DESTINATION) == pending_tree + _assert_failure( + _vendor(consumer, lock, "check", _VENDOR_NAME, "--offline", check=False), + "digest mismatch", + ) + assert (upstream / _SOURCE / "kernel.py").read_text( + encoding="utf-8" + ) == "VALUE = 'exported-fix'\n" + second_commit = _commit(upstream, "apply exported downstream fix") + + _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--url", + upstream.as_uri(), + "--branch", + "vendor-fix", + "--repo", + upstream, + ) + vendor = _vendor_data(lock) + assert vendor["url"] == upstream.as_uri() + assert vendor["branch"] == "vendor-fix" + assert vendor["commit"] == second_commit + assert "patch" not in vendor + assert _tree_snapshot(consumer / _DESTINATION) == pending_tree + _vendor(consumer, lock, "check", _VENDOR_NAME) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_partial_upstream_acceptance_uses_temporary_branches(tmp_path: Path) -> None: + upstream, base_commit = _make_upstream( + tmp_path, + { + "a.py": "A = 0\n", + "b.py": "B = 0\n", + }, + ) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + _create_vendor(consumer, lock, upstream, base_commit, mode="exact") + destination = consumer / _DESTINATION + _write_files(destination, {"a.py": "A = 1\n", "b.py": "B = 1\n"}) + accepted_destination = _tree_snapshot(destination) + + _git(upstream, "switch", "-c", "vendor-a-and-b") + _vendor(consumer, lock, "export", _VENDOR_NAME, "--repo", upstream) + temporary_ab_commit = _commit(upstream, "export A and B") + _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "vendor-a-and-b", + "--commit", + temporary_ab_commit, + "--repo", + upstream, + ) + assert _vendor_data(lock)["commit"] == temporary_ab_commit + assert _tree_snapshot(destination) == accepted_destination + + _git(upstream, "switch", "main") + _write_files(upstream / _SOURCE, {"a.py": "A = 1\n"}) + canonical_a_commit = _commit(upstream, "accept A upstream") + lock_snapshot = lock.read_bytes() + rejected = _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "main", + "--commit", + canonical_a_commit, + "--repo", + upstream, + check=False, + ) + _assert_failure(rejected, "b.py") + assert lock.read_bytes() == lock_snapshot + assert _tree_snapshot(destination) == accepted_destination + + _git(upstream, "switch", "-c", "vendor-b") + _write_files(upstream / _SOURCE, {"b.py": "B = 1\n"}) + temporary_b_commit = _commit(upstream, "carry B after A") + _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "vendor-b", + "--commit", + temporary_b_commit, + "--repo", + upstream, + ) + assert _vendor_data(lock)["commit"] == temporary_b_commit + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + _git(upstream, "switch", "main") + _write_files(upstream / _SOURCE, {"b.py": "B = 1\n"}) + canonical_ab_commit = _commit(upstream, "accept B upstream") + _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "main", + "--commit", + canonical_ab_commit, + "--repo", + upstream, + ) + vendor = _vendor_data(lock) + assert vendor["branch"] == "main" + assert vendor["commit"] == canonical_ab_commit + assert _tree_snapshot(destination) == accepted_destination + _vendor(consumer, lock, "check", _VENDOR_NAME) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_export_and_pin_retain_then_drop_compatibility_patch(tmp_path: Path) -> None: + upstream, base_commit = _make_upstream( + tmp_path, + { + "compat.py": "MODE = 'upstream'\n", + "feature.py": "FEATURE = 'base'\n", + }, + ) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + destination = consumer / _DESTINATION + (destination / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + _create_vendor(consumer, lock, upstream, base_commit, mode="patched") + original_vendor = _vendor_data(lock) + patch_path = consumer / str(original_vendor["patch"]) + patch_content = patch_path.read_bytes() + patch_mode = stat.S_IMODE(patch_path.stat().st_mode) + patch_digest = original_vendor["patch_digest"] + + no_change = _vendor( + consumer, + lock, + "export", + _VENDOR_NAME, + "--repo", + upstream, + check=False, + ) + _assert_failure(no_change, "no downstream change") + + (destination / "feature.py").write_text("FEATURE = 'exported'\n", encoding="utf-8") + pending_destination = _tree_snapshot(destination) + lock_snapshot = lock.read_bytes() + _git(upstream, "switch", "-c", "feature-fix") + _vendor(consumer, lock, "export", _VENDOR_NAME, "--repo", upstream) + assert (upstream / _SOURCE / "compat.py").read_text(encoding="utf-8") == "MODE = 'upstream'\n" + assert (upstream / _SOURCE / "feature.py").read_text( + encoding="utf-8" + ) == "FEATURE = 'exported'\n" + assert lock.read_bytes() == lock_snapshot + assert _tree_snapshot(destination) == pending_destination + assert patch_path.read_bytes() == patch_content + feature_commit = _commit(upstream, "export feature change") + + _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "feature-fix", + "--commit", + feature_commit, + "--repo", + upstream, + ) + retained_vendor = _vendor_data(lock) + assert retained_vendor["patch"] == original_vendor["patch"] + assert retained_vendor["patch_digest"] == patch_digest + assert patch_path.read_bytes() == patch_content + assert stat.S_IMODE(patch_path.stat().st_mode) == patch_mode + assert _tree_snapshot(destination) == pending_destination + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + _git(upstream, "switch", "-c", "incomplete-pin") + (upstream / _SOURCE / "feature.py").write_text("FEATURE = 'wrong'\n", encoding="utf-8") + incomplete_commit = _commit(upstream, "candidate does not reproduce destination") + retained_lock = lock.read_bytes() + rejected = _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "incomplete-pin", + "--commit", + incomplete_commit, + "--repo", + upstream, + check=False, + ) + _assert_failure(rejected, "feature.py") + assert lock.read_bytes() == retained_lock + assert patch_path.read_bytes() == patch_content + assert stat.S_IMODE(patch_path.stat().st_mode) == patch_mode + assert _tree_snapshot(destination) == pending_destination + + _git(upstream, "switch", "feature-fix") + (upstream / _SOURCE / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + absorbed_commit = _commit(upstream, "absorb TensorRT-LLM compatibility change") + _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "feature-fix", + "--commit", + absorbed_commit, + "--repo", + upstream, + ) + exact_vendor = _vendor_data(lock) + assert exact_vendor["commit"] == absorbed_commit + assert "patch" not in exact_vendor + assert "patch_digest" not in exact_vendor + assert not patch_path.exists() + assert _tree_snapshot(destination) == pending_destination + _vendor(consumer, lock, "check", _VENDOR_NAME) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_pin_restores_absorbed_patch_when_lock_save_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], +) -> None: + upstream, base_commit = _make_upstream(tmp_path, {"compat.py": "MODE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + destination = consumer / _DESTINATION + (destination / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + _create_vendor(consumer, lock, upstream, base_commit, mode="patched") + patch_path = consumer / str(_vendor_data(lock)["patch"]) + patch_content = patch_path.read_bytes() + patch_mode = stat.S_IMODE(patch_path.stat().st_mode) + lock_content = lock.read_bytes() + destination_snapshot = _tree_snapshot(destination) + + _git(upstream, "switch", "-c", "absorb-compat") + (upstream / _SOURCE / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + absorbed_commit = _commit(upstream, "absorb compatibility change") + module = _load_vendor_sources_module() + + def fail_save_lock(_: object) -> None: + raise OSError("injected pin lock-save failure") + + with monkeypatch.context() as failure: + failure.setattr(module, "_save_lock", fail_save_lock) + result = module.main( + [ + "--lock", + str(lock), + "pin", + _VENDOR_NAME, + "--branch", + "absorb-compat", + "--commit", + absorbed_commit, + "--repo", + str(upstream), + ] + ) + + assert result == 1 + assert "injected pin lock-save failure" in capfd.readouterr().err + assert lock.read_bytes() == lock_content + assert patch_path.read_bytes() == patch_content + assert stat.S_IMODE(patch_path.stat().st_mode) == patch_mode + assert _tree_snapshot(destination) == destination_snapshot + + _vendor( + consumer, + lock, + "pin", + _VENDOR_NAME, + "--branch", + "absorb-compat", + "--commit", + absorbed_commit, + "--repo", + upstream, + ) + assert not patch_path.exists() + assert "patch" not in _vendor_data(lock) + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream) + + +def test_pin_directory_sync_failure_retains_absorbed_patch( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], +) -> None: + upstream, base_commit = _make_upstream(tmp_path, {"compat.py": "MODE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + destination = consumer / _DESTINATION + (destination / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + _create_vendor(consumer, lock, upstream, base_commit, mode="patched") + patch_path = consumer / str(_vendor_data(lock)["patch"]) + patch_content = patch_path.read_bytes() + + _git(upstream, "switch", "-c", "absorb-compat") + (upstream / _SOURCE / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + absorbed_commit = _commit(upstream, "absorb compatibility change") + module = _load_vendor_sources_module() + real_fsync = os.fsync + + def fail_directory_sync(file_descriptor: int) -> None: + if stat.S_ISDIR(os.fstat(file_descriptor).st_mode): + raise OSError("injected lock-directory sync failure") + real_fsync(file_descriptor) + + with monkeypatch.context() as failure: + failure.setattr(module.os, "fsync", fail_directory_sync) + result = module.main( + [ + "--lock", + str(lock), + "pin", + _VENDOR_NAME, + "--branch", + "absorb-compat", + "--commit", + absorbed_commit, + "--repo", + str(upstream), + ] + ) + + assert result == 1 + assert "injected lock-directory sync failure" in capfd.readouterr().err + assert patch_path.read_bytes() == patch_content + updated_vendor = _vendor_data(lock) + assert updated_vendor["commit"] == absorbed_commit + assert "patch" not in updated_vendor + _vendor(consumer, lock, "check", _VENDOR_NAME, "--offline") + + +def test_pin_patch_cleanup_failure_succeeds_with_orphan_warning( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], +) -> None: + upstream, base_commit = _make_upstream(tmp_path, {"compat.py": "MODE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + destination = consumer / _DESTINATION + (destination / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + _create_vendor(consumer, lock, upstream, base_commit, mode="patched") + patch_path = consumer / str(_vendor_data(lock)["patch"]) + patch_content = patch_path.read_bytes() + + _git(upstream, "switch", "-c", "absorb-compat") + (upstream / _SOURCE / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + absorbed_commit = _commit(upstream, "absorb compatibility change") + module = _load_vendor_sources_module() + real_save_lock_checked = module._save_lock_checked + real_sync_directory = module._sync_directory + real_unlink = Path.unlink + events: list[str] = [] + + def record_lock_save(lock_file: object, description: str) -> None: + real_save_lock_checked(lock_file, description) + events.append("lock-save") + + def record_directory_sync(path: Path) -> None: + real_sync_directory(path) + events.append("directory-sync") + + def fail_patch_cleanup(path: Path, missing_ok: bool = False) -> None: + if path == patch_path: + events.append("patch-unlink") + raise OSError("injected absorbed-patch cleanup failure") + real_unlink(path, missing_ok=missing_ok) + + with monkeypatch.context() as failure: + failure.setattr(module, "_save_lock_checked", record_lock_save) + failure.setattr(module, "_sync_directory", record_directory_sync) + failure.setattr(Path, "unlink", fail_patch_cleanup) + result = module.main( + [ + "--lock", + str(lock), + "pin", + _VENDOR_NAME, + "--branch", + "absorb-compat", + "--commit", + absorbed_commit, + "--repo", + str(upstream), + ] + ) + + captured = capfd.readouterr() + assert result == 0 + assert "offline enforcement restored" in captured.out + assert "warning:" in captured.err.lower() + assert "lock is committed and valid" in captured.err + assert "unreferenced orphan" in captured.err + assert "delete that file manually" in captured.err.lower() + assert str(patch_path) in captured.err + assert events == ["lock-save", "directory-sync", "patch-unlink"] + assert patch_path.read_bytes() == patch_content + updated_vendor = _vendor_data(lock) + assert updated_vendor["commit"] == absorbed_commit + assert "patch" not in updated_vendor + assert "patch_digest" not in updated_vendor + _vendor(consumer, lock, "check", _VENDOR_NAME, "--offline") + + +def test_pin_interruption_after_lock_save_leaves_valid_lock_and_orphan_patch( + tmp_path: Path, +) -> None: + upstream, base_commit = _make_upstream(tmp_path, {"compat.py": "MODE = 'upstream'\n"}) + consumer, lock = _make_consumer(tmp_path) + _copy_python_sources(upstream, consumer) + destination = consumer / _DESTINATION + (destination / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + _create_vendor(consumer, lock, upstream, base_commit, mode="patched") + patch_path = consumer / str(_vendor_data(lock)["patch"]) + patch_content = patch_path.read_bytes() + patch_mode = stat.S_IMODE(patch_path.stat().st_mode) + + _git(upstream, "switch", "-c", "absorb-compat") + (upstream / _SOURCE / "compat.py").write_text("MODE = 'trtllm'\n", encoding="utf-8") + absorbed_commit = _commit(upstream, "absorb compatibility change") + lock_saved = tmp_path / "pin-lock-saved" + interrupt_script = "\n".join( + [ + "import importlib.util", + "import pathlib", + "import sys", + "import time", + "module_path, lock, upstream, commit, marker = sys.argv[1:]", + "spec = importlib.util.spec_from_file_location('_vendor_sources_interrupted', module_path)", + "assert spec is not None and spec.loader is not None", + "module = importlib.util.module_from_spec(spec)", + "sys.modules[spec.name] = module", + "spec.loader.exec_module(module)", + "save_lock_checked = module._save_lock_checked", + "def save_then_pause(lock_file, description):", + " save_lock_checked(lock_file, description)", + " pathlib.Path(marker).write_text('saved', encoding='utf-8')", + " time.sleep(30)", + "module._save_lock_checked = save_then_pause", + "raise SystemExit(module.main([", + " '--lock', lock, 'pin', 'example', '--branch', 'absorb-compat',", + " '--commit', commit, '--repo', upstream,", + "]))", + ] + ) + process = subprocess.Popen( + [ + sys.executable, + "-c", + interrupt_script, + str(_VENDOR_SOURCES), + str(lock), + str(upstream), + absorbed_commit, + str(lock_saved), + ], + cwd=consumer, + env=_command_env(), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + deadline = time.monotonic() + 10 + while not lock_saved.exists() and process.poll() is None and time.monotonic() < deadline: + time.sleep(0.01) + if not lock_saved.exists(): + stdout, stderr = process.communicate(timeout=5) + pytest.fail( + f"Pin did not reach lock-save barrier.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ) + process.kill() + process.communicate(timeout=5) + finally: + if process.poll() is None: + process.kill() + process.communicate(timeout=5) + + assert process.returncode is not None and process.returncode < 0 + interrupted_vendor = _vendor_data(lock) + assert interrupted_vendor["commit"] == absorbed_commit + assert "patch" not in interrupted_vendor + assert "patch_digest" not in interrupted_vendor + assert patch_path.read_bytes() == patch_content + assert stat.S_IMODE(patch_path.stat().st_mode) == patch_mode + _vendor(consumer, lock, "check", _VENDOR_NAME, "--offline") + _vendor(consumer, lock, "check", _VENDOR_NAME, "--repo", upstream)