From 1fa665b4694a7ab2e9461ccb40635f331d5d64be Mon Sep 17 00:00:00 2001 From: haow Date: Mon, 8 Jun 2026 04:24:23 -0700 Subject: [PATCH 1/4] [None][feat] add CuteDSL FP8/FP16 MLA decode attention backend Add a CuteDSL attention backend that intercepts the MLA decode-only path and dispatches to Blackwell CuTe DSL kernels (FP8 e4m3 / FP16), falling back to TrtllmAttention for all other paths. - attention_backend/cute_dsl.py: CuteDslAttention(TrtllmAttention) - cute_dsl_kernels/blackwell/attention/mla: decode fp8/fp16 kernels + helpers - custom_ops: register cute_dsl_mla_decode_{fp8,fp16}_blackwell ops - utils/__init__: wire up CUTEDSL backend selection - tests: CUTEDSL coverage in test_attention_mla + dedicated decode test Signed-off-by: haow --- .pre-commit-config.yaml | 10 + legacy-files.txt | 5 + pyproject.toml | 5 + ruff-legacy.toml | 5 + .../_torch/attention_backend/__init__.py | 2 + .../_torch/attention_backend/cute_dsl.py | 288 ++ .../_torch/attention_backend/utils.py | 3 + .../_torch/custom_ops/cute_dsl_custom_ops.py | 458 ++ .../blackwell/attention/__init__.py | 2 + .../blackwell/attention/mla/__init__.py | 10 + .../attention/mla/mla_decode_fp16.py | 4267 +++++++++++++++++ .../blackwell/attention/mla/mla_decode_fp8.py | 4212 ++++++++++++++++ .../blackwell/attention/mla/mla_helpers.py | 302 ++ .../_torch/attention/test_attention_mla.py | 41 +- .../attention/test_cute_dsl_mla_decode.py | 165 + 15 files changed, 9766 insertions(+), 9 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/cute_dsl.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py create mode 100644 tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85db633af271..4c19a9e33b05 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -299,6 +299,11 @@ common-files: &common_files | tensorrt_llm/_torch/custom_ops/userbuffers_custom_ops.py | tensorrt_llm/_torch/cute_dsl_kernels/__init__.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/__init__.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py | @@ -1674,6 +1679,11 @@ legacy-files: &legacy_files | tensorrt_llm/_torch/custom_ops/userbuffers_custom_ops.py | tensorrt_llm/_torch/cute_dsl_kernels/__init__.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/__init__.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py | + tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py | tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py | diff --git a/legacy-files.txt b/legacy-files.txt index e646b59a9b31..3ac91b7e4947 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -291,6 +291,11 @@ tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py tensorrt_llm/_torch/custom_ops/userbuffers_custom_ops.py tensorrt_llm/_torch/cute_dsl_kernels/__init__.py tensorrt_llm/_torch/cute_dsl_kernels/blackwell/__init__.py +tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py +tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py +tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py +tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py +tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py diff --git a/pyproject.toml b/pyproject.toml index 6d95d3204a9c..4f2200fe9e0d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -349,6 +349,11 @@ exclude = [ "tensorrt_llm/_torch/custom_ops/userbuffers_custom_ops.py", "tensorrt_llm/_torch/cute_dsl_kernels/__init__.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/__init__.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index c261c908abb1..9cbdb5416744 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -308,6 +308,11 @@ include = [ "tensorrt_llm/_torch/custom_ops/userbuffers_custom_ops.py", "tensorrt_llm/_torch/cute_dsl_kernels/__init__.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/__init__.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py", + "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/custom_pipeline.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py", "tensorrt_llm/_torch/cute_dsl_kernels/blackwell/utils.py", diff --git a/tensorrt_llm/_torch/attention_backend/__init__.py b/tensorrt_llm/_torch/attention_backend/__init__.py index ae0d8b87cd85..0895418af85d 100644 --- a/tensorrt_llm/_torch/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/__init__.py @@ -1,4 +1,5 @@ from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE +from .cute_dsl import CuteDslAttention from .interface import AttentionBackend, AttentionForwardArgs, AttentionMetadata from .sparse import get_sparse_attn_kv_cache_manager from .trtllm import AttentionInputType, TrtllmAttention, TrtllmAttentionMetadata @@ -9,6 +10,7 @@ "AttentionBackend", "AttentionForwardArgs", "AttentionInputType", + "CuteDslAttention", "TrtllmAttention", "TrtllmAttentionMetadata", "VanillaAttention", diff --git a/tensorrt_llm/_torch/attention_backend/cute_dsl.py b/tensorrt_llm/_torch/attention_backend/cute_dsl.py new file mode 100644 index 000000000000..abb9eea3ddf6 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/cute_dsl.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuTeDSL attention backend. + +Subclasses ``TrtllmAttention`` and overrides ``forward`` to dispatch the +MLA decode-only path into one of two Blackwell CuTe DSL kernels via a +single shared dispatcher (``_dispatch_cute_dsl_mla_decode``) parameterised +by the kernel's input dtype: + +- ``kernel_dtype == torch.float8_e4m3fn`` + → ``torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell`` +- ``kernel_dtype == torch.float16`` + → ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` + +``forward`` picks the kernel dtype directly from runtime state +(``has_fp8_kv_cache`` → FP8; otherwise ``q.dtype == torch.float16`` → FP16; +neither match → TRTLLM fallback). Every other code path (context / +chunked prefill / cached-KV MLA context / non-MLA / mixed batches / +unsupported SM) goes through ``super().forward`` unchanged. + +Subclassing ``TrtllmAttention`` matters: ``modules/attention.py`` selects +the MLA chunked-prefill / cached-context fast paths via +``isinstance(self.mha, TrtllmAttention)``, so the CUTEDSL backend must +satisfy that check or those paths would silently fall back to the slow +default context path. This mirrors the MoE CuTeDSL integration +(``CuteDslFusedMoE(CutlassFusedMoE)``). +""" + +import math +import os +from typing import Optional + +import torch + +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger + +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from .interface import AttentionForwardArgs +from .trtllm import TrtllmAttention, TrtllmAttentionMetadata + +_DEBUG_FALLBACK = os.environ.get("TLLM_CUTE_DSL_ATTN_DEBUG_FALLBACK", "0") == "1" + +# cutlass / kernel-class imports were used by the now-removed in-eligibility +# ``can_implement`` checks — the dispatch now goes through +# ``torch.ops.trtllm.cute_dsl_mla_decode_*_blackwell`` and the kernel-level +# Runner runs ``can_implement`` itself. ``IS_CUTLASS_DSL_AVAILABLE`` is +# still consulted in the eligibility preconditions below to short-circuit +# environments without the cutlass package. + + +class CuteDslAttention(TrtllmAttention): + """CuteDSL attention backend. + + Inherits the full ``TrtllmAttention`` machinery (metadata, KV cache, + quant flags, RoPE buffers, MLA helpers) and overrides ``forward`` to + intercept the MLA decode-only batch. Anything that fails eligibility + falls through to ``TrtllmAttention.forward``. + """ + + Metadata = TrtllmAttentionMetadata + + # ------------------------------------------------------------------ + # Shared preconditions (everything except dtype-specific gates). + # ------------------------------------------------------------------ + def _cute_dsl_mla_decode_common_preconditions( + self, + metadata: TrtllmAttentionMetadata, + forward_args: Optional[AttentionForwardArgs], + ) -> bool: + """Checks that are identical for the FP8 and FP16 paths. + + Note: phase routing (is_mla_enable / num_contexts / num_generations) + is already handled in ``forward`` before either dtype's eligibility + method runs, so it isn't repeated here. + + The dtype-specific eligibility methods own everything else + (KV-cache dtype, ``q.dtype``, kernel-level ``can_implement``). + """ + if not IS_CUTLASS_DSL_AVAILABLE: + return False + if get_sm_version() not in (100, 103): + return False + if self.predicted_tokens_per_seq is None or not (1 <= self.predicted_tokens_per_seq <= 4): + return False + if metadata.kv_cache_block_offsets is None: + return False + if forward_args is None or forward_args.latent_cache is None: + return False + return True + + # ================================================================== + # MLA decode dispatch (FP8 / FP16) + # ================================================================== + + def _dispatch_cute_dsl_mla_decode( + self, + q: torch.Tensor, + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + output: torch.Tensor, + kernel_dtype: torch.dtype, + ) -> torch.Tensor: + """MLA decode dispatch shared by FP8 and FP16 paths. + + ``kernel_dtype`` is the in/out tensor dtype the chosen CuTe DSL + kernel expects — ``torch.float8_e4m3fn`` for the FP8 kernel, + ``torch.float16`` for the FP16 kernel. The op called is selected + from it. + + Assumes the MLA module has already (1) built ``q`` as the fused + ``[num_tokens, H * (D_latent + D_rope)]`` tensor with RoPE applied + to the rope half, and (2) appended the new token to the paged + latent cache. + """ + if kernel_dtype == torch.float8_e4m3fn: + op = torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell + elif kernel_dtype == torch.float16: + op = torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell + else: + raise ValueError( + f"CuteDslAttention: unsupported kernel_dtype={kernel_dtype}; " + "expected torch.float8_e4m3fn or torch.float16" + ) + + num_tokens = q.shape[0] + num_seqs = metadata.num_generations + seq_len_q = num_tokens // num_seqs + assert seq_len_q * num_seqs == num_tokens, ( + f"CuteDslAttention MLA decode expects num_tokens " + f"({num_tokens}) divisible by num_generations ({num_seqs})" + ) + + # Both kernels: L == 512, R == 64, in/out dtype == kernel_dtype. + d_latent = self.kv_lora_rank + d_rope = self.qk_rope_head_dim + h = self.num_heads + page_size = metadata.tokens_per_block + + # q → [H, D, S_q, B]. Cast to kernel dtype if upstream q is in + # something else (e.g. bf16 model + FP8 KV — lossy but defined). + q_kernel = q if q.dtype == kernel_dtype else q.to(kernel_dtype) + q_view = q_kernel.view(num_seqs, seq_len_q, h, d_latent + d_rope) + q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0).contiguous() + q_rope = q_view[..., d_latent:].permute(2, 3, 1, 0).contiguous() + + # Paged MLA pool view as the kernel's dtype. + # NOTE: pool tensor handle and block-table layout for MLA depends + # on kv-cache-manager wiring; revisit if ``get_buffers`` exposes a + # different layout. + kv_pool = metadata.kv_cache_manager.get_buffers(self.layer_idx) + kv_pool_typed = kv_pool.view(kernel_dtype) + c_pool_latent = kv_pool_typed[..., :d_latent] + c_pool_rope = kv_pool_typed[..., d_latent : d_latent + d_rope] + + block_offsets = metadata.kv_cache_block_offsets + if block_offsets.dim() == 4: + page_table_layer = block_offsets[self.layer_idx, :, 0, :] + elif block_offsets.dim() == 3: + page_table_layer = block_offsets[:, 0, :] + else: + page_table_layer = block_offsets + # Kernel: [max_pages, B], leading_dim=0 ⇒ pages contiguous per B. + page_table = page_table_layer.transpose(0, 1).contiguous().to(torch.int32) + + cache_seqs = metadata.kv_lens_cuda_runtime.to(torch.int32) + + split_kv = 1 + workspace = torch.empty(0, dtype=torch.float32, device=q.device) + block_split_kvs = torch.empty(0, dtype=torch.int32, device=q.device) + + o_kernel = torch.empty( + (h, d_latent, seq_len_q, num_seqs), + dtype=kernel_dtype, + device=q.device, + ) + lse = torch.empty( + (h, seq_len_q, num_seqs), + dtype=torch.float32, + device=q.device, + ) + + # MLA softmax scale follows the canonical TRT-LLM formula + # (see modules/attention.py:1453): + # softmax_scale = 1 / (sqrt(qk_head_dim) * q_scaling) + # where ``qk_head_dim`` is the *unabsorbed* Q head dim + # ``qk_nope_head_dim + qk_rope_head_dim``. We deliberately do NOT + # use ``(kv_lora_rank + qk_rope_head_dim)`` even though that is + # the absorbed attention's inner dimension — the scale stays bound + # to the original head dim so attention scores match the + # unabsorbed reference. ``self.q_scaling`` carries any YaRN + # ``mscale`` adjustment (set by the MLA module via + # ``q_scaling = 1 / (mscale * mscale)`` — see attention.py:1428). + qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + softmax_scale = float(1.0 / (math.sqrt(qk_head_dim) * self.q_scaling)) + output_scale = 1.0 + + c_latent_kernel = c_pool_latent.unsqueeze(-1).contiguous() + c_rope_kernel = c_pool_rope.unsqueeze(-1).contiguous() + + op( + q_latent, + q_rope, + c_latent_kernel, + c_rope_kernel, + page_table, + cache_seqs, + block_split_kvs, + o_kernel, + lse, + workspace, + self.num_heads, + seq_len_q, + page_size, + True, # is_persistent + True, # is_var_seq + False, # is_var_split_kv + split_kv, + softmax_scale, + output_scale, + ) + + attn_out = o_kernel.permute(3, 2, 0, 1).reshape(num_tokens, h * d_latent) + output.copy_(attn_out.to(output.dtype)) + return output + + # ================================================================== + # forward dispatch + # ================================================================== + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: Optional[AttentionForwardArgs] = None, + **kwargs, + ) -> torch.Tensor: + if forward_args is None and kwargs: + forward_args = AttentionForwardArgs(**kwargs) + kwargs = {} + + # Phase routing — only the MLA decode-only batch is eligible for the + # CuTe DSL fast path. Everything else falls through to TRTLLM: + # - non-MLA attention → TRTLLM + # - prefill / mixed batch (has context) → TRTLLM + # - decode-only MLA → try CuteDSL FP8/FP16, + # TRTLLM on failure + is_decode_only_mla = ( + self.is_mla_enable and metadata.num_contexts == 0 and metadata.num_generations > 0 + ) + + if ( + is_decode_only_mla + and self._cute_dsl_mla_decode_common_preconditions(metadata, forward_args) + and q.shape[0] % metadata.num_generations == 0 + ): + # Direct dtype-based dispatch (no per-dtype eligibility helper). + # The kernel-level Runner runs ``can_implement`` for the chosen + # dtype; anything it rejects falls through to TRTLLM via the + # try/except below. + if getattr(self, "has_fp8_kv_cache", False): + kernel_dtype = torch.float8_e4m3fn + elif q.dtype == torch.float16: + kernel_dtype = torch.float16 + else: + kernel_dtype = None + + if kernel_dtype is not None: + try: + output = q.new_empty( + (q.shape[0], self.num_heads * self.kv_lora_rank), dtype=q.dtype + ) + return self._dispatch_cute_dsl_mla_decode( + q, metadata, forward_args, output, kernel_dtype + ) + except Exception as exc: # noqa: BLE001 + if _DEBUG_FALLBACK: + logger.warning( + "CuteDslAttention: MLA decode fast path " + "(kernel_dtype=%s) failed (%s); falling back " + "to TRTLLM backend.", + kernel_dtype, + exc, + ) + + return super().forward(q, k, v, metadata, forward_args=forward_args, **kwargs) diff --git a/tensorrt_llm/_torch/attention_backend/utils.py b/tensorrt_llm/_torch/attention_backend/utils.py index 2c4bc574bb58..fa699c5d536c 100644 --- a/tensorrt_llm/_torch/attention_backend/utils.py +++ b/tensorrt_llm/_torch/attention_backend/utils.py @@ -48,6 +48,9 @@ def get_attention_backend( elif backend_name == "FLASHINFER_STAR_ATTENTION" and IS_FLASHINFER_AVAILABLE: from .star_flashinfer import StarAttention return StarAttention + elif backend_name == "CUTEDSL": + from .cute_dsl import CuteDslAttention + return CuteDslAttention logger.warning("Falling back to TRTLLM attention backend") return TrtllmAttention diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 01df74678018..fa8ea384c448 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7328,3 +7328,461 @@ def _( max_context_len, dtype=output_dtype, device=q.device) + + # ========================================================================= + # MLA decode (Blackwell) — wraps the CuTe DSL kernels that live at + # tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/. + # Used by the CUTEDSL attention backend (see attention_backend/cute_dsl.py). + # + # One generic Runner ``CuteDSLNVMlaDecodeBlackwellRunner`` services both + # FP8 and FP16 paths — only the cutlass ``in_dtype`` is passed at + # construction; the kernel class is derived from it via + # ``CuteDSLNVMlaDecodeBlackwellRunner._KERNEL_CLASS_BY_DTYPE``. Each + # dtype still has its own ``@torch.library.custom_op`` (distinct op + # name + fake-tensor rule); the ops differ only in which ``in_dtype`` + # they hand the generic Runner. + # + # torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell + # → CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.Float8E4M3FN) + # (→ BlackwellMultiHeadLatentAttentionForwardFP8) + # + # torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell + # → CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.Float16) + # (→ BlackwellMultiHeadLatentAttentionForwardFP16) + # ========================================================================= + + from ..cute_dsl_kernels.blackwell.attention.mla.mla_decode_fp8 import \ + BlackwellMultiHeadLatentAttentionForwardFP8 + from ..cute_dsl_kernels.blackwell.attention.mla.mla_decode_fp16 import \ + BlackwellMultiHeadLatentAttentionForwardFP16 + + _CUTE_DSL_MLA_CLUSTER_SHAPE_MNK = (2, 1, 1) + + class CuteDSLNVMlaDecodeBlackwellRunner(TunableRunner): + """Generic TunableRunner for the Blackwell CuTe DSL MLA decode kernels. + + Works for both FP8 and FP16 — pass the cutlass input dtype at + construction; the kernel class is derived from it: + + CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass.Float8E4M3FN, ...) # → ...ForwardFP8 + CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass.Float16, ...) # → ...ForwardFP16 + + ``get_valid_tactics`` returns the tiler shapes as tactics + (``(mma_qk_tiler_mn, mma_pv_tiler_mn)`` tuples), filtered by the + kernel's static ``can_implement``. The current candidate list + carries only ``((128, 128), (128, 256))`` — the lone combination + both kernels accept today — but more can be added without + touching ``forward``. ``kernel_cache`` is class-level and keyed + by ``(in_dtype, ..., mma_qk_tiler_mn, mma_pv_tiler_mn)``, so FP8 + and FP16 compilations, plus future tilers, coexist without + collisions. + """ + kernel_cache = dict() + + # in_dtype → kernel class. The kernels' own ``can_implement`` is + # what ultimately rejects unsupported dtypes, but this lookup + # picks which kernel we even try to compile. + _KERNEL_CLASS_BY_DTYPE = { + cutlass.Float8E4M3FN: BlackwellMultiHeadLatentAttentionForwardFP8, + cutlass.Float16: BlackwellMultiHeadLatentAttentionForwardFP16, + } + + def __init__( + self, + in_dtype, + num_heads: int, + seq_len_q: int, + page_size: int, + is_persistent: bool = True, + is_var_seq: bool = True, + is_var_split_kv: bool = False, + skip_correction_threshold: float = 0.0, + ): + super().__init__() + kernel_class = self.__class__._KERNEL_CLASS_BY_DTYPE.get(in_dtype) + if kernel_class is None: + raise ValueError( + f"CuteDSLNVMlaDecodeBlackwellRunner: unsupported " + f"in_dtype={in_dtype}. Supported: " + f"{list(self.__class__._KERNEL_CLASS_BY_DTYPE.keys())}") + self.in_dtype = in_dtype + self.kernel_class = kernel_class + self.num_heads = num_heads + self.seq_len_q = seq_len_q + 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.skip_correction_threshold = skip_correction_threshold + + def unique_id(self): + # `kernel_class` is derived from `in_dtype`, so dropping it + # from the key keeps cache slots 1-to-1 with the in_dtype. + # The tilers are NOT here — they're part of the tactic and + # appended into the cache key inside ``forward``. + return ( + self.in_dtype, + self.num_heads, + self.seq_len_q, + self.page_size, + self.is_persistent, + self.is_var_seq, + self.is_var_split_kv, + self.skip_correction_threshold, + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]: + """Filter the candidate tilers via the kernel's + ``can_implement``. Returns a list of + ``(mma_qk_tiler_mn, mma_pv_tiler_mn)`` tuples; AutoTuner picks + one and passes it to ``forward`` as ``tactic``. + """ + if get_sm_version() not in (100, 103): + return [] + q_latent, q_rope, _c_latent, _c_rope, _page_table, cache_seqs, \ + *_rest = inputs + h, latent_dim, seq_len_q, _ = q_latent.shape + rope_dim = q_rope.shape[1] + batch_size = cache_seqs.shape[0] + + # Candidate tilers — widen this list to enable AutoTuner + # exploration over tile shapes. Each entry is + # ``(mma_qk_tiler_mn, mma_pv_tiler_mn)``. + candidate_tiler_tactics = [ + ((128, 128), (128, 256)), + ] + + valid = [] + for mma_qk_tiler_mn, mma_pv_tiler_mn in candidate_tiler_tactics: + if self.kernel_class.can_implement( + batch_size, + seq_len_q, + self.page_size, + h, + latent_dim, + rope_dim, + self.in_dtype, # in_dtype + self.in_dtype, # out_dtype + cutlass.Float32, # acc_dtype + cutlass.Float32, # lse_dtype + mma_qk_tiler_mn, + mma_pv_tiler_mn, + 1, + self.is_persistent, + self.is_var_seq, + self.is_var_split_kv, + self.page_size, + ): + valid.append((mma_qk_tiler_mn, mma_pv_tiler_mn)) + else: + logger.debug( + "CuteDSLNVMlaDecodeBlackwellRunner.can_implement " + "rejected tactic: kernel=%s in_dtype=%s " + "H=%d L=%d R=%d S=%d B=%d page_size=%d " + "mma_qk=%s mma_pv=%s persistent=%s var_seq=%s " + "var_split=%s", self.kernel_class.__name__, + self.in_dtype, h, latent_dim, rope_dim, seq_len_q, + batch_size, self.page_size, mma_qk_tiler_mn, + mma_pv_tiler_mn, self.is_persistent, self.is_var_seq, + self.is_var_split_kv) + return valid + + def get_tuning_config(self) -> TuningConfig: + return TuningConfig() + + def forward( + self, + inputs: List[torch.Tensor], + tactic, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + (q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, + block_split_kvs, o, lse, workspace) = inputs + split_kv = int(kwargs.get("split_kv", 1)) + softmax_scale = float(kwargs.get("softmax_scale", 1.0)) + output_scale = float(kwargs.get("output_scale", 1.0)) + + # Unpack the tactic produced by ``get_valid_tactics``. When + # AutoTuner isn't engaged (e.g. the attention backend calls + # the op without ``choose_one``), tactic may be ``None`` — + # fall back to the default (128,128)/(128,256) shape. + if isinstance(tactic, tuple) and len(tactic) == 2: + mma_qk_tiler_mn, mma_pv_tiler_mn = tactic + else: + mma_qk_tiler_mn, mma_pv_tiler_mn = (128, 128), (128, 256) + mma_qk_tiler_mn = tuple(mma_qk_tiler_mn) + mma_pv_tiler_mn = tuple(mma_pv_tiler_mn) + + torch_stream = torch.cuda.current_stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + cache_key = self.unique_id() + (mma_qk_tiler_mn, mma_pv_tiler_mn) + if cache_key not in CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache: + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + _CUTE_DSL_MLA_CLUSTER_SHAPE_MNK[0] * + _CUTE_DSL_MLA_CLUSTER_SHAPE_MNK[1] * + _CUTE_DSL_MLA_CLUSTER_SHAPE_MNK[2]) + + mla = self.kernel_class( + cutlass.Float32, # acc_dtype + cutlass.Float32, # lse_dtype + mma_qk_tiler_mn, + mma_pv_tiler_mn, + max_active_clusters, + self.page_size, + self.skip_correction_threshold, + self.is_persistent, + self.is_var_seq, + self.is_var_split_kv, + num_heads=self.num_heads, + seq_len_q=self.seq_len_q, + ) + + q_latent_ct = cute.runtime.from_dlpack( + q_latent, + assumed_align=16).mark_layout_dynamic(leading_dim=1) + q_rope_ct = cute.runtime.from_dlpack( + q_rope, assumed_align=16).mark_layout_dynamic(leading_dim=1) + c_latent_ct = cute.runtime.from_dlpack( + c_latent, + assumed_align=16).mark_layout_dynamic(leading_dim=1) + c_rope_ct = cute.runtime.from_dlpack( + c_rope, assumed_align=16).mark_layout_dynamic(leading_dim=1) + page_table_ct = cute.runtime.from_dlpack( + page_table, + assumed_align=16).mark_layout_dynamic(leading_dim=0) + o_ct = cute.runtime.from_dlpack( + o, assumed_align=16).mark_layout_dynamic(leading_dim=1) + lse_ct = cute.runtime.from_dlpack( + lse, assumed_align=16).mark_layout_dynamic(leading_dim=0) + workspace_ct = cute.runtime.from_dlpack( + workspace, assumed_align=16).mark_layout_dynamic() + cache_seqs_ct = cute.runtime.from_dlpack( + cache_seqs, assumed_align=16).mark_layout_dynamic() + block_split_kvs_ct = (cute.runtime.from_dlpack( + block_split_kvs, assumed_align=16).mark_layout_dynamic() + if self.is_var_split_kv else None) + + CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache[cache_key] = \ + cute.compile( + mla, + q_latent_ct, + q_rope_ct, + c_latent_ct, + c_rope_ct, + page_table_ct, + o_ct, + lse_ct, + workspace_ct, + cutlass.Int32(split_kv), + cache_seqs_ct, + block_split_kvs_ct, + cutlass.Float32(softmax_scale), + cutlass.Float32(output_scale), + stream, + options="--opt-level 2", + ) + + compiled_mla = CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache[ + cache_key] + compiled_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs if self.is_var_split_kv else None, + softmax_scale, + output_scale, + stream, + ) + return o, lse + + @torch.library.custom_op( + "trtllm::cute_dsl_mla_decode_fp8_blackwell", + mutates_args=("o", "lse", "workspace"), + device_types="cuda", + ) + def cute_dsl_mla_decode_fp8_blackwell( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + c_latent: torch.Tensor, + c_rope: torch.Tensor, + page_table: torch.Tensor, + cache_seqs: torch.Tensor, + block_split_kvs: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + workspace: torch.Tensor, + num_heads: int, + seq_len_q: int, + page_size: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + split_kv: int, + softmax_scale: float, + output_scale: float, + ) -> None: + """CuTe DSL FP8 MLA decode (Blackwell SM100/SM103). + + ``o``, ``lse``, ``workspace`` are mutated in place. Tensor layouts: + see ``BlackwellMultiHeadLatentAttentionForwardFP8``. + """ + if (sm_version := get_sm_version()) not in (100, 103): + raise ValueError( + f"trtllm::cute_dsl_mla_decode_fp8_blackwell requires SM 100 or " + f"SM 103, got SM {sm_version}") + + runner = CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass.Float8E4M3FN, + num_heads=num_heads, + seq_len_q=seq_len_q, + page_size=page_size, + is_persistent=is_persistent, + is_var_seq=is_var_seq, + is_var_split_kv=is_var_split_kv, + ) + inputs = [ + q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, + block_split_kvs, o, lse, workspace + ] + tuner = AutoTuner.get() + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_mla_decode_fp8_blackwell", + [runner], + runner.get_tuning_config(), + inputs, + ) + runner( + inputs, + tactic=best_tactic, + split_kv=split_kv, + softmax_scale=softmax_scale, + output_scale=output_scale, + ) + + @torch.library.register_fake("trtllm::cute_dsl_mla_decode_fp8_blackwell") + def _( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + c_latent: torch.Tensor, + c_rope: torch.Tensor, + page_table: torch.Tensor, + cache_seqs: torch.Tensor, + block_split_kvs: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + workspace: torch.Tensor, + num_heads: int, + seq_len_q: int, + page_size: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + split_kv: int, + softmax_scale: float, + output_scale: float, + ) -> None: + return None + + @torch.library.custom_op( + "trtllm::cute_dsl_mla_decode_fp16_blackwell", + mutates_args=("o", "lse", "workspace"), + device_types="cuda", + ) + def cute_dsl_mla_decode_fp16_blackwell( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + c_latent: torch.Tensor, + c_rope: torch.Tensor, + page_table: torch.Tensor, + cache_seqs: torch.Tensor, + block_split_kvs: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + workspace: torch.Tensor, + num_heads: int, + seq_len_q: int, + page_size: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + split_kv: int, + softmax_scale: float, + output_scale: float, + ) -> None: + """CuTe DSL FP16 MLA decode (Blackwell SM100/SM103). + + ``o``, ``lse``, ``workspace`` are mutated in place. Tensor layouts: + see ``BlackwellMultiHeadLatentAttentionForwardFP16``. + """ + if (sm_version := get_sm_version()) not in (100, 103): + raise ValueError( + f"trtllm::cute_dsl_mla_decode_fp16_blackwell requires SM 100 " + f"or SM 103, got SM {sm_version}") + + runner = CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass.Float16, + num_heads=num_heads, + seq_len_q=seq_len_q, + page_size=page_size, + is_persistent=is_persistent, + is_var_seq=is_var_seq, + is_var_split_kv=is_var_split_kv, + ) + inputs = [ + q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, + block_split_kvs, o, lse, workspace + ] + tuner = AutoTuner.get() + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_mla_decode_fp16_blackwell", + [runner], + runner.get_tuning_config(), + inputs, + ) + runner( + inputs, + tactic=best_tactic, + split_kv=split_kv, + softmax_scale=softmax_scale, + output_scale=output_scale, + ) + + @torch.library.register_fake("trtllm::cute_dsl_mla_decode_fp16_blackwell") + def _( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + c_latent: torch.Tensor, + c_rope: torch.Tensor, + page_table: torch.Tensor, + cache_seqs: torch.Tensor, + block_split_kvs: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + workspace: torch.Tensor, + num_heads: int, + seq_len_q: int, + page_size: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + split_kv: int, + softmax_scale: float, + output_scale: float, + ) -> None: + return None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py new file mode 100644 index 000000000000..ed6d7e25176a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .mla_decode_fp8 import BlackwellMultiHeadLatentAttentionForwardFP8 +from .mla_decode_fp16 import BlackwellMultiHeadLatentAttentionForwardFP16 + +__all__ = [ + "BlackwellMultiHeadLatentAttentionForwardFP8", + "BlackwellMultiHeadLatentAttentionForwardFP16", +] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py new file mode 100644 index 000000000000..f1cc55fb629b --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py @@ -0,0 +1,4267 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +import math +import os +import sys +from types import SimpleNamespace +from typing import Optional, Tuple, Type + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.cute.testing as testing +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.base_dsl.arch import Arch +from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import BaseDSL +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .mla_helpers import (LOG2_E, MAX_SPLITS, MLAStaticTileScheduler, + MLAStaticTileSchedulerParams, ceil_div, + create_mla_static_tile_scheduler, + create_mla_static_tile_scheduler_params) + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) +""" +A Multi-Head Latent Attention (MLA) example with FP16 data type for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of inference of multi-head latent attention using a TMA + Blackwell +SM100 TensorCore warp-specialized persistent kernel. The implementation integrates the (Qc + Qr)*(Kc + Kr)^T +matrix multiplication, softmax normalization, and softmax((Qc + Qr)*(Kc + Kr)^T)*Vc into a single kernel. +The kernel provides support for page table storage and variable-length KV cache sequences. It implements KV splitting +functionality to minimize latency when processing long KV sequences. + +The kernel implements key optimizations including: +- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue) +- Pipeline stages between different warps for overlapping computation and memory access +- Support for different precision data types +- Two sub-kernels (split KV kernel and reduction kernel) that enable split KV processing + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mla_fp16.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float16 --out_dtype Float16 \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent + +The above example runs Multi-Head Latent Attention (MLA) with the following configuration: +- Batch size: 4 +- Sequence length of Q: 1 +- Sequence length of K: 1024 +- Latent dimension: 512 +- RoPE dimension: 64 +- Number of heads: 128 +- Data types: Float16 (input), Float16 (output), Float32 (accumulation and LSE) + +It utilizes page table storage for the KV cache and enables both variable-length KV cache sequences +and variable split KV processing with persistent scheduling. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mla_fp16.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float16 --out_dtype Float16 \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent --warmup_iterations 3 \ + --iterations 10 --skip_ref_check + +Constraints for this example: +* Data type requirements: + - Input/output: Float16 + - Accumulation and LSE: Float32 +* Fixed architecture parameters: + - Number of attention heads: 128 + - Latent dimension: 512 + - RoPE dimension: 64 +* Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) +* Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +* Query sequence length must be 1-4 +* Only supports 2-CTA instructions +* Variable sequence length requires page table storage enabled +""" + + +class BlackwellMultiHeadLatentAttentionForwardFP16: + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + max_active_clusters: int, + page_size: int, + skip_correction_threshold: float, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + + :param acc_dtype: Data type for accumulation S and O + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Data type for output LSE + :type lse_dtype: Type[cutlass.Numeric] + :param mma_s_tiler: The (H, K) tile shape of the MMA instruction for S + :type mma_s_tiler: Tuple[int, int] + :param mma_p_tiler: The (H, D) tile shape of the MMA instruction for P + :type mma_p_tiler: Tuple[int, int] + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: int + :param page_size: The page size of the page table + :type page_size: int + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split KV + :type is_var_split_kv: bool + """ + + self.latent_dim = 512 + self.rope_dim = 64 + 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.skip_correction_threshold = skip_correction_threshold + self.is_persistent = is_persistent + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), + # while warps 2-3 handle accumulation for second half [n/2, n) + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 + mma_qk_tiler_k = self.rope_dim + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + mma_qk_tiler_k, + ) + self.mma_qk_rope_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + self.rope_dim, + ) + self.mma_pv_tiler = ( + self.mma_pv_tiler_mn[0], + self.mma_pv_tiler_mn[1], + self.mma_qk_tiler[1] * self.mma_qk_tiler[2] // + self.mma_pv_tiler_mn[1], + ) + self.iterations_qk_latent = self.latent_dim // self.mma_qk_tiler[2] + self.iterations_qk_rope = mma_qk_tiler_k // self.mma_qk_tiler[2] + self.iterations_qk = self.iterations_qk_latent + self.iterations_qk_rope + self.iterations_pv_k = self.mma_qk_tiler[1] // self.mma_pv_tiler[2] + self.iterations_pv_n = self.latent_dim // self.mma_pv_tiler[1] + + # Set specialized warp ids + self.compute_warp_ids = (0, 1, 2, 3) + self.correction_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + + self.load_tma_warp_id = 9 + self.load_pt_warp_id = 10 + self.empty_warp_ids = (11, ) + self.threads_per_cta = self.threads_per_warp * len(( + self.mma_warp_id, + self.load_tma_warp_id, + self.load_pt_warp_id, + *self.compute_warp_ids, + *self.correction_warp_ids, + *self.empty_warp_ids, + )) + + # register settings + self.softmax_reg_num = 192 + self.correction_reg_num = 208 + self.other_reg_num = 96 + # Named barriers + self.tmem_ptr_sync_bar = pipeline.NamedBarrier( + barrier_id=1, + num_threads=(self.threads_per_warp + + self.threads_per_warp * self.num_compute_warps * 2), + ) + self.softmax_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=2, + num_threads=(self.threads_per_warp * self.num_compute_warps)) + self.epilogue_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=3, + num_threads=(self.threads_per_warp * self.num_compute_warps)) + + def _setup_attributes(self): + """Set up configurations and parameters for the MLA kernel operation. + + This method initializes and configures various attributes required for the + execution of the multi-head latent attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.load_q_stage = 1 + self.load_kv_stage = 15 + self.mma_s_stage = 2 + self.p_mma_stage = 2 + self.p_cor_stage = 2 + self.mma_o_stage = 1 + self.load_pt_stage = 4 + + self.tmem_o_offset = self.mma_s_stage * self.mma_qk_tiler[ + 1] // self.warps_in_n + self.correction_factor_offset = (self.tmem_o_offset + + self.latent_dim // self.warps_in_n) + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_table: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: cuda.CUstream, + ): + """Execute the Multi-Head Latent Attention operation on the provided tensors. + + The method handles: + 1. Initialization of workspace for temporary split KV buffers + 2. Validation of tensor data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch(split KV kernel and reduction kernel) with appropriate parameters + + :param q_latent: The query tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type q_latent: cute.Tensor + :param q_rope: The query RoPE tensor with shape [num_head, rope_dim, seq_len_q, batch_size] + :type q_rope: cute.Tensor + :param c_latent: The key tensor with shape [seq_len_k, latent_dim, batch_size] + :type c_latent: cute.Tensor + :param c_rope: The key RoPE tensor with shape [seq_len_k, rope_dim, batch_size] + :type c_rope: cute.Tensor + :param page_table: The page table tensor with shape [page_count, batch_size] + :type page_table: cute.Tensor + :param o: The output tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type o: cute.Tensor + :param lse: The LSE tensor with shape [num_head, seq_len_q, batch_size] + :type lse: cute.Tensor + :param workspace: The workspace tensor with 1-d shape prepared for acc_o and acc_lse + :type workspace: cute.Tensor + :param split_kv: The scalar factor for split KV + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax + :type softmax_scale: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + + :raises TypeError: If tensor data types don't match or aren't supported + """ + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q_latent.element_type + self.k_dtype = c_latent.element_type + self.v_dtype = c_latent.element_type + self.o_dtype = o.element_type + + # check type consistency + if cutlass.const_expr(self.q_dtype != self.k_dtype + or self.q_dtype != self.v_dtype): + raise TypeError( + f"Type mismatch: {self.q_dtype} != {self.k_dtype} or {self.q_dtype} != {self.v_dtype}" + ) + # check leading dimensions of input/output + 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") + 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") + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") + + acc_o, acc_lse = self.initialize_workspace( + q_latent.shape[0], + q_latent.shape[1], + q_latent.shape[2], + q_latent.shape[3], + split_kv, + self.acc_dtype, + workspace, + ) + + c_latent_tranpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor(c_latent.iterator, + c_latent_tranpose_layout) + + self.q_major_mode = tcgen05.OperandMajorMode.K + self.k_major_mode = tcgen05.OperandMajorMode.K + self.v_major_mode = tcgen05.OperandMajorMode.MN + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from smem & k-major + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.acc_dtype, + cta_group, + self.mma_qk_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.acc_dtype, + cta_group, + self.mma_pv_tiler[:2], + ) + + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape, ), + ) + + self.epi_tile = self.mma_pv_tiler[:2] + + q_latent_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_tiler, + self.q_dtype, + (self.iterations_qk_latent * self.load_q_stage), + ) + q_latent_smem_layout_staged = cute.logical_divide( + q_latent_smem_layout_staged, + (None, None, None, self.iterations_qk_latent)) + q_rope_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.q_dtype, + self.load_q_stage, + ) + + # rope reuse the same smem layout as latent + kc_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_tiler, + self.k_dtype, + self.load_kv_stage, + ) + kc_page_tile_size = min( + self.page_size, + qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape) + + kc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + (self.mma_qk_tiler[0] // qk_tiled_mma.thr_id.shape, + self.mma_qk_tiler[2]), + self.k_dtype, + self.load_kv_stage, + ) + kc_smem_layout_for_tma = cute.tiled_divide( + kc_smem_layout_for_tma, (kc_page_tile_size, self.mma_qk_tiler[2])) + + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.mma_pv_tiler, + self.q_dtype, + (self.iterations_pv_k * self.p_mma_stage), + ) + p_smem_layout_staged = cute.logical_divide( + p_smem_layout_staged, (None, None, None, self.iterations_pv_k)) + + vc_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.mma_pv_tiler, + self.v_dtype, + self.load_kv_stage, + ) + vc_page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + vc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.MN, + (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, + self.mma_pv_tiler[2]), + self.v_dtype, + self.load_kv_stage, + ) + vc_smem_layout_for_tma = cute.tiled_divide( + vc_smem_layout_for_tma, + ( + pv_tiled_mma.op.shape_mnk[1] // pv_tiled_mma.thr_id.shape, + vc_page_tile_size, + ), + ) + # TMA load for Q latent and rope + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_latent_smem_layout = cute.select(q_latent_smem_layout_staged, + mode=[0, 1, 2]) + tma_atom_q_latent, tma_tensor_q_latent = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_latent, + q_latent_smem_layout, + self.mma_qk_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + q_rope_smem_layout = cute.select(q_rope_smem_layout_staged, + mode=[0, 1, 2]) + tma_atom_q_rope, tma_tensor_q_rope = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_rope, + q_rope_smem_layout, + self.mma_qk_rope_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + # TMA load for c latent and k rope + kc_smem_layout = cute.select(kc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent, tma_tensor_c_latent = self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + tma_atom_c_rope, tma_tensor_c_rope = self.make_paged_tiled_tma_atom( + tma_load_op, + c_rope, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + # TMA load for c latent transpose + vc_smem_layout = cute.select(vc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent_transpose, tma_tensor_c_latent_transpose = ( + self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent_transpose, + vc_smem_layout, + (self.mma_pv_tiler[1], self.mma_pv_tiler[2]), + pv_tiled_mma, + is_k_load=False, + )) + + q_latent_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_latent_smem_layout) * + cute.size(qk_tiled_mma.thr_id.shape) * self.iterations_qk_latent) + q_rope_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_rope_smem_layout) * + cute.size(qk_tiled_mma.thr_id.shape) * self.iterations_qk_rope) + q_copy_size = q_latent_copy_size + q_rope_copy_size + kc_copy_size = cute.size_in_bytes( + self.k_dtype, cute.select(kc_smem_layout_staged, mode=[ + 0, 1, 2 + ])) * cute.size(qk_tiled_mma.thr_id.shape) + vc_copy_size = cute.size_in_bytes( + self.v_dtype, cute.select(vc_smem_layout_staged, mode=[ + 0, 1, 2 + ])) * cute.size(pv_tiled_mma.thr_id.shape) + assert (kc_copy_size == vc_copy_size + ), "kc_copy_size and vc_copy_size must be the same" + + self.tma_copy_q_bytes = q_copy_size + self.tma_copy_kc_bytes = kc_copy_size + + tile_sched_params, grid = self._compute_grid( + o, + split_kv, + self.cluster_shape_mnk, + self.max_active_clusters, + self.is_persistent, + ) + + @cute.struct + class SplitKVKernelSharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.load_q_stage * 2] + load_kv_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.load_kv_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.mma_s_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.p_mma_stage * 2] + p_cor_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.p_cor_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.mma_o_stage * 2] + load_pt_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.load_pt_stage * 2] + # Tmem dealloc cluster barrier + tmem_dealloc_mbar_ptr: cutlass.Int64 + + # Tmem holding buffer + tmem_holding_buf: cutlass.Int32 + # Smem tensors + softmax_smem_exchange: cute.struct.MemRange[self.acc_dtype, + self.num_compute_warps * + self.threads_per_warp] + epilogue_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp] + smem_q_latent: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, + cute.cosize(q_latent_smem_layout_staged)], + 1024, + ] + smem_q_rope: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, + cute.cosize(q_rope_smem_layout_staged)], + 1024, + ] + smem_kc: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, + cute.cosize(kc_smem_layout_staged)], + 1024, + ] + smem_p: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, + cute.cosize(p_smem_layout_staged)], + 1024, + ] + smem_page_table: cute.struct.MemRange[cutlass.Int32, + self.load_pt_stage * + self.mma_qk_tiler[1] // 2] + + softmax_scale_log2 = softmax_scale * LOG2_E + self.split_kv_kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q_latent, + tma_tensor_q_latent, + tma_atom_q_rope, + tma_tensor_q_rope, + tma_atom_c_latent, + tma_tensor_c_latent, + tma_atom_c_rope, + tma_tensor_c_rope, + tma_atom_c_latent_transpose, + tma_tensor_c_latent_transpose, + page_table, + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale_log2, + output_scale, + q_latent_smem_layout_staged, + q_rope_smem_layout_staged, + kc_smem_layout_staged, + p_smem_layout_staged, + vc_smem_layout_staged, + kc_smem_layout_for_tma, + vc_smem_layout_for_tma, + cta_layout_vmnk, + tile_sched_params, + SplitKVKernelSharedStorage, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=SplitKVKernelSharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( + grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], + smem=MAX_SPLITS * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.jit + def make_paged_tiled_tma_atom( + self, + tma_load_op: cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp, + gmem: cute.Tensor, + smem_layout: cute.Layout, + mma_tiler, + tiled_mma: cute.TiledMma, + is_k_load: bool, + ): + ident = cute.make_identity_layout(gmem.shape) + g_tile = cute.composition(ident, mma_tiler) + cta_mn = mma_tiler[0] // tiled_mma.thr_id.shape + cta_v_map = cute.flat_divide(g_tile, (cta_mn, )) + cta_v_map = cute.select(cta_v_map, mode=[0, 2]) + page_tile_size = (min(self.page_size, cta_mn) if is_k_load else min( + self.page_size, mma_tiler[1])) + cta_v_map = cute.zipped_divide( + cta_v_map, + (page_tile_size, mma_tiler[1]) if is_k_load else + (cta_mn, page_tile_size), + ) + cta_v_map = cute.select(cta_v_map, mode=[0]) + from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir + + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + gmem.value, + smem_layout.value, + cta_v_map, + tma_load_op._to_ir(), + num_multicast=1, + ) + return ( + cute.CopyAtom(tma_load_op, + cpasync.CopyBulkTensorTileG2SNonExecTrait(res[0])), + res[1], + ) + + @cute.kernel + def split_kv_kernel( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tma_atom_q_latent: Optional[cute.CopyAtom], + mQL: cute.Tensor, + tma_atom_q_rope: Optional[cute.CopyAtom], + mQR: cute.Tensor, + tma_atom_c_latent: Optional[cute.CopyAtom], + mCL: cute.Tensor, + tma_atom_c_rope: Optional[cute.CopyAtom], + mKR: cute.Tensor, + tma_atom_c_latent_transpose: Optional[cute.CopyAtom], + mCLT: cute.Tensor, + mPT: cute.Tensor, + mO: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mAccO: Optional[cute.Tensor], + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + q_latent_smem_layout_staged: cute.ComposedLayout, + q_rope_smem_layout_staged: cute.ComposedLayout, + kc_smem_layout_staged: cute.ComposedLayout, + p_smem_layout_staged: cute.ComposedLayout, + vc_smem_layout_staged: cute.ComposedLayout, + kc_smem_layout_for_tma: cute.ComposedLayout, + vc_smem_layout_for_tma: cute.ComposedLayout, + cta_layout_vmnk: cute.Layout, + tile_sched_params: MLAStaticTileSchedulerParams, + SharedStorage: cutlass.Constexpr, + ): + """The device split_kv kernel implementation of the Multi-Head Latent Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the MLA computation: + 1. Load warp: Loads Q/C latent/rope data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Compute warps: Compute softmax and do rescaling on accumulators, and store the intermediate/final results + to global memory + + The kernel produces either intermediate or final results of the MLA computation based on the split_kv parameter. + When split_kv is 1, the kernel generates the final results directly. Otherwise, it produces intermediate results + that will later be combined by a reduction kernel. + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases. + + :param tiled_mma_qk: Tiled MMA for Q*K^T + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: Tiled MMA for P*V + :type tiled_mma_pv: cute.TiledMma + :param tma_atom_q_latent: TMA copy atom for query latent tensor + :type tma_atom_q_latent: cute.CopyAtom + :param mQL: query latent tensor + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom + :param mKR: Compressed rope tensor + :type mKR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor + :type mCL: cute.Tensor + :param tma_atom_c_rope: TMA copy atom for c rope tensor + :type tma_atom_c_rope: cute.CopyAtom + :param mCLT: Compressed latent transpose tensor + :type mCLT: cute.Tensor + :param mPT: Page table tensor + :type mPT: cute.Tensor + :param mO: Output tensor + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor + :type mLSE: cute.Tensor + :param mAccO: Intermediate accumulator output tensor + :type mAccO: cute.Tensor + :param mAccLSE: Intermediate accumulator log-sum-exp tensor + :type mAccLSE: cute.Tensor + :param split_kv: The split_kv parameter + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax + :type softmax_scale_log2: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param q_latent_smem_layout_staged: Shared memory layout for query latent tensor + :type q_latent_smem_layout_staged: cute.ComposedLayout + :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor + :type q_rope_smem_layout_staged: cute.ComposedLayout + :param kc_smem_layout_staged: Shared memory layout for key/value latent/rope tensor + :type kc_smem_layout_staged: cute.ComposedLayout + :param p_smem_layout_staged: Shared memory layout for probability matrix + :type p_smem_layout_staged: cute.ComposedLayout + :param vc_smem_layout_staged: Shared memory layout for value tensor + :type vc_smem_layout_staged: cute.ComposedLayout + :param kc_smem_layout_for_tma: Shared memory layout for key/value latent tensor for TMA + :type kc_smem_layout_for_tma: cute.ComposedLayout + :param vc_smem_layout_for_tma: Shared memory layout for value tensor for TMA + :type vc_smem_layout_for_tma: cute.ComposedLayout + :param cta_layout_vmnk: Layout for compute threads + :type cta_layout_vmnk: cute.Layout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: MLAStaticTileSchedulerParams + :param SharedStorage: Shared storage for the kernel + :type SharedStorage: cutlass.Constexpr + """ + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # Prefetch tma descriptor + if warp_idx == self.mma_warp_id: + cpasync.prefetch_descriptor(tma_atom_q_latent) + cpasync.prefetch_descriptor(tma_atom_q_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent) + cpasync.prefetch_descriptor(tma_atom_c_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent_transpose) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_bar, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + load_q_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_q_stage, + self.tma_copy_q_bytes, + ) + load_kv_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_kv_stage, + self.tma_copy_kc_bytes, + ) + mma_s_pipeline = self.make_and_init_mma_s_pipeline( + storage.mma_s_mbar_ptr.data_ptr(), cta_layout_vmnk) + p_mma_pipeline = self.make_and_init_p_mma_pipeline( + storage.p_mma_mbar_ptr.data_ptr(), cta_layout_vmnk) + p_cor_pipeline = self.make_and_init_p_cor_pipeline( + storage.p_cor_mbar_ptr.data_ptr()) + mma_o_pipeline = self.make_and_init_mma_o_pipeline( + storage.mma_o_mbar_ptr.data_ptr(), cta_layout_vmnk) + load_pt_pipeline = self.make_and_init_load_pt_pipeline( + storage.load_pt_mbar_ptr.data_ptr()) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, + is_relaxed=True) + + # Generate smem tensor Q/KC/VC/exchange + # (MMA, MMA_H, MMA_R, PIPE) + sQ = storage.smem_q_latent.get_tensor( + q_latent_smem_layout_staged.outer, + swizzle=q_latent_smem_layout_staged.inner) + sQ_rope = storage.smem_q_rope.get_tensor( + q_rope_smem_layout_staged.outer, + swizzle=q_rope_smem_layout_staged.inner) + # (MMA, MMA_K, MMA_R, PIPE) + sKC = storage.smem_kc.get_tensor(kc_smem_layout_staged.outer, + swizzle=kc_smem_layout_staged.inner) + sKC_for_tma = storage.smem_kc.get_tensor( + kc_smem_layout_for_tma.outer, + swizzle=kc_smem_layout_for_tma.inner, + ) + # (MMA, MMA_D, MMA_K, PIPE) + # reuse smem + sVC_ptr = cute.recast_ptr(sKC.iterator, vc_smem_layout_staged.inner) + sVC = cute.make_tensor(sVC_ptr, vc_smem_layout_staged.outer) + sVC_for_tma = cute.make_tensor(sVC_ptr, vc_smem_layout_for_tma.outer) + # (MMA, MMA_H, MMA_K) + sP = storage.smem_p.get_tensor(p_smem_layout_staged.outer, + swizzle=p_smem_layout_staged.inner) + sPT = storage.smem_page_table.get_tensor( + cute.make_layout((self.mma_qk_tiler[1] // 2, self.load_pt_stage))) + # (compute_threads,) + softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp)) + epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp)) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load warps, including page table and data tensors + # /////////////////////////////////////////////////////////////////////////////// + + if warp_idx >= self.empty_warp_ids[ + 0] and warp_idx <= self.empty_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.other_reg_num) + if warp_idx == self.load_pt_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_pt_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_pt_stage) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + load_pt_common_params = SimpleNamespace( + blk_coord=blk_coord, + load_pt_pipeline=load_pt_pipeline, + mPT=mPT, + sPT=sPT, + tidx=tidx, + page_size=mCL.shape[0], + ) + load_pt_producer_state = self.load_page_table( + load_pt_common_params, + k_index, + k_tile_count, + load_pt_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_pt_pipeline.producer_tail(load_pt_producer_state) + if warp_idx == self.load_tma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_q_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_q_stage) + load_kv_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_kv_stage) + load_pt_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_pt_stage) + load_pt_release_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_pt_stage) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_kv_pipeline=load_kv_pipeline, + mPT=mPT, + sPT=sPT, + load_pt_pipeline=load_pt_pipeline, + ) + tma_qk_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + tma_atom_q_latent=tma_atom_q_latent, + tma_atom_q_rope=tma_atom_q_rope, + tma_atom_c_latent=tma_atom_c_latent, + tma_atom_c_rope=tma_atom_c_rope, + mQL=mQL, + mQR=mQR, + mCL=mCL, + mKR=mKR, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC_for_tma, + ) + tma_pv_params = SimpleNamespace( + tiled_mma_pv=tiled_mma_pv, + tma_atom_c_latent_transpose=tma_atom_c_latent_transpose, + mCL=mCL, + mKR=mKR, + mCLT=mCLT, + sVC=sVC_for_tma, + ) + # Load tma + ( + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) = self.load_tma( + tma_common_params, + tma_qk_params, + tma_pv_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + load_q_pipeline.producer_tail(load_q_producer_state) + load_kv_pipeline.producer_tail(load_kv_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA warp + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + # Alloc tensor memory buffer + tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100")) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + load_q_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_q_stage) + load_kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_kv_stage) + mma_s_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_s_stage) + p_mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_mma_stage) + mma_o_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_o_stage) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord) + if k_tile_count > 0: + mma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_kv_pipeline=load_kv_pipeline, + tmem_ptr=tmem_ptr, + is_leader_cta=is_leader_cta, + L=mCL.shape[1], + ) + mma_qk_params = SimpleNamespace( + mma_s_pipeline=mma_s_pipeline, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC, + ) + mma_pv_params = SimpleNamespace( + p_mma_pipeline=p_mma_pipeline, + mma_o_pipeline=mma_o_pipeline, + sP=sP, + sVC=sVC, + ) + ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma( + mma_common_params, + mma_qk_params, + mma_pv_params, + k_tile_count, + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + mma_s_pipeline.producer_tail(mma_s_producer_state) + mma_o_pipeline.producer_tail(mma_o_producer_state) + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Compute warp + # /////////////////////////////////////////////////////////////////////////////// + if (warp_idx >= self.compute_warp_ids[0] + and warp_idx <= self.compute_warp_ids[-1]): + cute.arch.setmaxregister_increase(self.softmax_reg_num) + mma_s_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_s_stage) + p_mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_mma_stage) + p_cor_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_cor_stage) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=softmax_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, + p_cor_pipeline=p_cor_pipeline, + ) + compute_softmax_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + sP=sP, + mma_s_pipeline=mma_s_pipeline, + p_mma_pipeline=p_mma_pipeline, + softmax_scale_log2=softmax_scale_log2, + ) + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = ( + self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + )) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_cor_pipeline.producer_tail(p_cor_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction warp + # /////////////////////////////////////////////////////////////////////////////// + if (warp_idx >= self.correction_warp_ids[0] + and warp_idx <= self.correction_warp_ids[-1]): + cute.arch.setmaxregister_increase(self.correction_reg_num) + p_cor_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_cor_stage) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=epilogue_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + H=mQL.shape[0], + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, + p_cor_pipeline=p_cor_pipeline, + mma_o_pipeline=mma_o_pipeline, + ) + compute_epilogue_params = SimpleNamespace( + output_scale=output_scale, + softmax_scale_log2=softmax_scale_log2, + mAccLSE=mAccLSE, + mLSE=mLSE, + ) + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + return + + @cute.kernel + def reduction_kernel( + self, + mO: cute.Tensor, + mLSE: cute.Tensor, + mAccO: cute.Tensor, + mAccLSE: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + ): + """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results + from multiple split_kv blocks into final outputs. + + :param mO: Output tensor for storing final results + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor for storing final LSE values + :type mLSE: cute.Tensor + :param mAccO: Accumulated output tensor from split_kv blocks + :type mAccO: cute.Tensor + :param mAccLSE: Accumulated LSE tensor from split_kv blocks + :type mAccLSE: cute.Tensor + :param split_kv: Number of split_kv blocks + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor (for variable split_kv) + :type block_split_kvs: cute.Tensor + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + blk_coord = (bidx, bidy, bidz) + local_split_kv = (block_split_kvs[blk_coord[2]] + if self.is_var_split_kv else split_kv) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], + self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) + local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() + storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) + smem_lse_scale = cute.make_tensor(lse_scale_ptr, + cute.make_layout(MAX_SPLITS)) + + gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) + + local_lse = cute.make_rmem_tensor(cute.make_layout(lse_per_thread), + self.lse_dtype) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = (gLSE[split_kv_idx] if cute.elem_less( + split_kv_idx, local_split_kv) else -self.lse_dtype.inf) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 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 = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = (lse_max + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) + or sum_lse != sum_lse else self.lse_dtype.inf) + if tidx == 0: + mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2(local_lse[i] - + global_lse, + fastmath=True) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.latent_dim, self.threads_per_warp * self.num_compute_warps) + gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] + rAccO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), + self.acc_dtype) + rO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), + self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + return + + @staticmethod + def get_split_kv(B: int, S: int, K: int, mma_qk_tiler_mn: tuple, + max_active_blocks: int) -> int: + """Get the proper split_kv value for the MLA kernel based on parameters. + + :param B: Batch size + :type B: int + :param S: Sequence length + :type S: int + :param K: Sequence length + :type K: int + :param mma_qk_tiler_mn: MLA tiling parameters + :type mma_qk_tiler_mn: tuple + :param max_active_blocks: Maximum number of active blocks + :type max_active_blocks: int + :return: Split_kv value + :rtype: int + """ + max_splits = ceil_div(K, mma_qk_tiler_mn[1]) + blocks_per_batch = max(1, max_active_blocks // B // (S * 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) + max_split_kv = 32 + return min(split_wave_aware, max_split_kv) + + @cute.jit + def get_k_tile_count( + self, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + blk_coord: cute.Coord, + ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. + + :param split_kv: Split_kv value + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param blk_coord: Block coordinate + :type blk_coord: cute.Coord + :return: k_index, k_tile_count, split_kv + :rtype: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + """ + K = cache_seqs[blk_coord[2]] + if cutlass.const_expr(self.is_var_split_kv): + split_kv = block_split_kvs[blk_coord[2]] + + k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) + k_index = blk_coord[3] * k_tile_per_cta + k_tile_count = max( + 0, + min(k_tile_total, k_index + k_tile_per_cta) - k_index) + return k_index, k_tile_count, split_kv + + @cute.jit + def load_page_table( + self, + common_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_pt_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load warp to load page table. Updates the load pt producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_pt_producer_state: The load pt producer state + :type load_pt_producer_state: pipeline.PipelineState + + :return: The load pt producer state + :rtype: pipeline.PipelineState + """ + mPT = common_params.mPT[None, common_params.blk_coord[2]] + page_per_tile = self.mma_qk_tiler[1] // self.page_size + tidx = common_params.tidx % self.threads_per_warp + + load_pt_pipeline = common_params.load_pt_pipeline + while k_tile_count > 0: + load_pt_pipeline.producer_acquire(load_pt_producer_state) + + elem_per_thread = cute.ceil_div(page_per_tile, + self.threads_per_warp) + + # atom_async_copy: async copy atom for page table load + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), + cutlass.Int32, + num_bits_per_copy=cutlass.Int32.width, + ) + mPT_for_copy = cute.flat_divide(mPT, (1, )) + sPT_for_copy = cute.flat_divide(common_params.sPT, (1, )) + # elem_per_thread is a dynamic value depends on the page_size setting. + for i in range(elem_per_thread): + idx = i * self.threads_per_warp + tidx + if cute.elem_less(k_index * page_per_tile + idx, + mPT.shape[0]) and cute.elem_less( + idx, page_per_tile): + cute.copy( + atom_async_copy, + mPT_for_copy[None, k_index * page_per_tile + idx], + sPT_for_copy[None, idx, load_pt_producer_state.index], + ) + else: + sPT_for_copy[None, idx, + load_pt_producer_state.index].fill(0) + mbar_ptr = load_pt_pipeline.producer_get_barrier( + load_pt_producer_state) + load_pt_pipeline.producer_commit(load_pt_producer_state) + load_pt_producer_state.advance() + k_index += 1 + k_tile_count -= 1 + + return load_pt_producer_state + + @cute.jit + def load_tma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_kv_producer_state: pipeline.PipelineState, + load_pt_consumer_state: pipeline.PipelineState, + load_pt_release_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Load wrap to load Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param v_params: The v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_kv_producer_state: The load kv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_consumer_state: The load pt consumer state + :type load_pt_consumer_state: pipeline.PipelineState + :param load_pt_release_state: The load pt release state + :type load_pt_release_state: pipeline.PipelineState + + :return: The load q producer state, load kv producer state, load pt consumer state, and load pt release state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for QK TMA load + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) + mma_qk_tiler_mk_rope = cute.select(self.mma_qk_rope_tiler, mode=[0, 2]) + gQR = cute.flat_divide(qk_params.mQR, mma_qk_tiler_mk_rope) + + thr_mma_qk = qk_params.tiled_mma_qk.get_slice( + common_params.blk_coord[0] % + cute.size(qk_params.tiled_mma_qk.thr_id)) + tSgQL = thr_mma_qk.partition_A(gQL) + tSgQR = thr_mma_qk.partition_A(gQR) + + cta_m = min( + qk_params.tiled_mma_qk.op.shape_mnk[0] // + qk_params.tiled_mma_qk.thr_id.shape, + self.page_size, + ) + page_tile_size = min(self.page_size, cta_m) + gCL = cute.tiled_divide(qk_params.mCL, + (page_tile_size, self.mma_qk_tiler[2])) + tSgCL = (gCL[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] if cta_m < self.page_size else gCL[None, 0, None, None]) + gKR = cute.tiled_divide(qk_params.mKR, + (page_tile_size, self.mma_qk_tiler[2])) + tSgKR = (gKR[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] if cta_m < self.page_size else gKR[None, 0, None, None]) + + # tma partition for q, k latent/rope + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tQsQ, tQLgQL_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_latent, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ, 0, 3), + cute.group_modes(tSgQL, 0, 3), + ) + + tQsQ_rope, tQRgQR_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_rope, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ_rope, 0, 3), + cute.group_modes(tSgQR, 0, 3), + ) + + tKCsKC, tCLgCL = cpasync.tma_partition( + qk_params.tma_atom_c_latent, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgCL, + ) + + _, tKRgKR = cpasync.tma_partition( + qk_params.tma_atom_c_rope, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgKR, + ) + + tQLgQL = tQLgQL_mkl[None, None, None, common_params.blk_coord[1], + common_params.blk_coord[2]] + tQRgQR = tQRgQR_mkl[None, None, None, common_params.blk_coord[1], + common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for V TMA load + page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + gCLT = cute.flat_divide(v_params.mCLT, + (self.mma_pv_tiler[1], page_tile_size)) + cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape + gCLT = cute.logical_divide(gCLT, + (cta_n, ))[(None, + common_params.blk_coord[0]), + None, None, None, None] + tOgCLT = cute.tiled_divide(gCLT, (cta_n, page_tile_size)) + tOgCLT = tOgCLT[None, 0, 0, None, None, None] + + # tma partition for vc + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tVCsVC, tCLTgCLT = cpasync.tma_partition( + v_params.tma_atom_c_latent_transpose, + 0, + cute.make_layout(1), + v_params.sVC, + tOgCLT, + ) + + # set extra params + common_params.mPT = mPT + qk_params.tQLgQL = tQLgQL + qk_params.tQRgQR = tQRgQR + qk_params.tCLgCL = tCLgCL + qk_params.tKRgKR = tKRgKR + qk_params.tQsQ = tQsQ + qk_params.tQsQ_rope = tQsQ_rope + qk_params.tKCsKC = tKCsKC + v_params.tCLTgCLT = tCLTgCLT + v_params.tVCsVC = tVCsVC + + load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = ( + self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_q=True, + )) + k_index += 1 + k_tile_count -= 1 + while k_tile_count > 0: + load_q_producer_state, load_kv_producer_state, load_pt_consumer_state = ( + self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_q=False, + )) + load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index - 1, + load_kv_producer_state, + load_pt_release_state, + ) + k_index += 1 + k_tile_count -= 1 + + # load last v tile + load_kv_producer_state, load_pt_release_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index - 1, + load_kv_producer_state, + load_pt_release_state, + ) + return ( + load_q_producer_state, + load_kv_producer_state, + load_pt_consumer_state, + load_pt_release_state, + ) + + @cute.jit + def load_tma_qk_one_k_tile( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_kv_producer_state: pipeline.PipelineState, + load_pt_consumer_state: pipeline.PipelineState, + load_q: bool, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, + pipeline.PipelineState]: + """Load one k-tile of Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_kv_producer_state: The load kv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_consumer_state: The load pt consumer state + :type load_pt_consumer_state: pipeline.PipelineState + :param load_q: Whether to load q + :type load_q: bool + + :return: The load q producer state, load kv producer state, and load pt consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = ceil_div(self.mma_qk_tiler[1] // self.page_size, + qk_params.tiled_mma_qk.thr_id.shape) + common_params.load_pt_pipeline.consumer_wait(load_pt_consumer_state) + page_table_stage = load_pt_consumer_state.index + load_pt_consumer_state.advance() + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), + cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = ( + common_params.sPT[0, + page_table_stage] if self.mma_qk_tiler[1] // + self.page_size == 1 else + common_params.sPT[i + + common_params.blk_coord[0] * page_per_tile, + page_table_stage]) + # load q once at first iteration + if cutlass.const_expr(load_q): + common_params.load_q_pipeline.producer_acquire( + load_q_producer_state) + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_q_pipeline.producer_get_barrier( + load_q_producer_state) + for i in cutlass.range(self.iterations_qk_latent): + # load q latent + cute.copy( + qk_params.tma_atom_q_latent, + qk_params.tQLgQL[None, 0, i], + qk_params.tQsQ[None, (i, 0)], + tma_bar_ptr=tma_bar_ptr, + ) + for i in cutlass.range(self.iterations_qk_rope): + # load q rope + cute.copy( + qk_params.tma_atom_q_rope, + qk_params.tQRgQR[None, 0, i], + qk_params.tQsQ_rope[None, i], + tma_bar_ptr=tma_bar_ptr, + ) + load_q_producer_state.advance() + load_kv_pipeline = common_params.load_kv_pipeline + tma_bar_ptr = load_kv_pipeline.producer_get_barrier( + load_kv_producer_state) + for i in cutlass.range(self.iterations_qk_latent): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier( + load_kv_producer_state) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_tile): + # load k latent + cute.copy( + qk_params.tma_atom_c_latent, + qk_params.tCLgCL[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_kv_producer_state.advance() + + for i in cutlass.range(self.iterations_qk_rope): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier( + load_kv_producer_state) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_tile): + # load k rope + cute.copy( + qk_params.tma_atom_c_rope, + qk_params.tKRgKR[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_kv_producer_state.advance() + + return load_q_producer_state, load_kv_producer_state, load_pt_consumer_state + + @cute.jit + def load_tma_v_one_k_tile( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + load_kv_producer_state: pipeline.PipelineState, + load_pt_release_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of compressed latent transpose tensor(v). Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The load tma v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param load_kv_producer_state: The load qkv producer state + :type load_kv_producer_state: pipeline.PipelineState + :param load_pt_release_state: The load pt release state + :type load_pt_release_state: pipeline.PipelineState + + :return: The load kv producer state and load pt release state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = self.mma_pv_tiler[ + 2] * self.iterations_pv_k // self.page_size + page_per_subtile = ceil_div(page_per_tile, self.iterations_pv_k) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), + cutlass.Int32) + page_table_stage = load_pt_release_state.index + for i in cutlass.range(page_per_tile): + k_idx[i] = (common_params.sPT[0, page_table_stage] if page_per_tile + == 1 else common_params.sPT[i, page_table_stage]) + common_params.load_pt_pipeline.consumer_release(load_pt_release_state) + load_pt_release_state.advance() + load_kv_pipeline = common_params.load_kv_pipeline + tma_bar_ptr = load_kv_pipeline.producer_get_barrier( + load_kv_producer_state) + for i in cutlass.range(self.iterations_pv_k): + for j in cutlass.range(self.iterations_pv_n): + # get the mbar ptr from pipeline. + tma_bar_ptr = load_kv_pipeline.producer_get_barrier( + load_kv_producer_state) + load_kv_pipeline.producer_acquire(load_kv_producer_state) + for k in cutlass.range(page_per_subtile): + k_idx_i = k_idx[ + k + i // ceil_div(self.iterations_pv_k, page_per_tile) * + page_per_subtile] + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[ + None, + j, + i % ceil_div(self.iterations_pv_k, page_per_tile), + k_idx_i, + ], + v_params.tVCsVC[None, 0, k, + load_kv_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + + load_kv_producer_state.advance() + return load_kv_producer_state, load_pt_release_state + + @cute.jit + def mma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + pv_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_kv_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """MMA warp to compute the result of Q*K^T and P*V. Updates the tiled mma and pipeline states. + + :param common_params: The common parameters for mma qk and pv + :type common_params: SimpleNamespace + :param qk_params: The mma qk parameters + :type qk_params: SimpleNamespace + :param pv_params: The mma pv parameters + :type pv_params: SimpleNamespace + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + :param p_mma_consumer_state: The p mma consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The mma o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the tiled mma pv, the load q consumer state, the load kv consumer state, the mma s producer state, the p mma consumer state, and the mma o producer state + :rtype: tuple[cute.TiledMma, cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + tSrQ = tiled_mma_qk.make_fragment_A(qk_params.sQ) + tSrQ_rope = tiled_mma_qk.make_fragment_A(qk_params.sQ_rope) + tSrKC = tiled_mma_qk.make_fragment_B(qk_params.sKC) + tOrP = tiled_mma_pv.make_fragment_A(pv_params.sP) + tOrVC = tiled_mma_pv.make_fragment_B(pv_params.sVC) + + tStS_shape = tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1])) + tStS_staged_fake = tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage)) + # use real tmem ptr for tStS + tStS_staged = cute.make_tensor(common_params.tmem_ptr, + tStS_staged_fake.layout) + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1])) + # mma O has 1 stage. + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO_staged = cute.make_tensor( + tStS_staged.iterator + self.tmem_o_offset, tOtO_layout) + + # set more parameters + qk_params.tSrQ = tSrQ + qk_params.tSrQ_rope = tSrQ_rope + qk_params.tSrKC = tSrKC + qk_params.tStS_staged = tStS_staged + pv_params.tOrP = tOrP + pv_params.tOrVC = tOrVC + pv_params.tOtO_staged = tOtO_staged + + # mma O accumulates on K, so the accumulate flag is set to False once before all K blocks. + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + if common_params.is_leader_cta: + load_q_release_state = load_q_consumer_state.clone() + + ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + wait_q=True, + ) + k_tile_count -= 1 + while k_tile_count > 0: + ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + wait_q=False, + ) + ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + k_tile_count -= 1 + + # release q consumer states + load_q_pipeline.consumer_release(load_q_release_state) + load_q_release_state.advance() + ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + return ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def mma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + tiled_mma_qk: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_kv_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + wait_q: bool, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states. + + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the load q consumer state, the load kv consumer state, and the mma s producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + tStS = qk_params.tStS_staged[None, None, None, + mma_s_producer_state.index] + + qk_params.mma_s_pipeline.producer_acquire(mma_s_producer_state) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + load_kv_pipeline = common_params.load_kv_pipeline + if cutlass.const_expr(wait_q): + load_q_pipeline.consumer_wait(load_q_consumer_state) + load_q_consumer_state.advance() + for q_stage in range(self.iterations_qk_latent): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + kc_stage = load_kv_consumer_state.index + for k_block in cutlass.range(cute.size(qk_params.tSrQ.shape[2])): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ[None, None, k_block, q_stage], + qk_params.tSrKC[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + for q_stage in range(self.iterations_qk_rope): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + kc_stage = load_kv_consumer_state.index + for k_block in cutlass.range(self.rope_dim // + tiled_mma_qk.shape_mnk[2]): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ_rope[None, None, k_block, q_stage], + qk_params.tSrKC[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + + qk_params.mma_s_pipeline.producer_commit(mma_s_producer_state) + mma_s_producer_state.advance() + return ( + tiled_mma_qk, + load_q_consumer_state, + load_kv_consumer_state, + mma_s_producer_state, + ) + + @cute.jit + def mma_pv( + self, + common_params: SimpleNamespace, + pv_params: SimpleNamespace, + tiled_mma_pv: cute.TiledMma, + load_kv_consumer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for P*V. Updates the tiled mma pv and pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param pv_params: The pv parameters + :type pv_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_kv_consumer_state: The load kv consumer state + :type load_kv_consumer_state: pipeline.PipelineState + :param p_mma_consumer_state: The P MMA consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The MMA o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma pv, the load qkv consumer state, the P MMA consumer state, and the MMA o producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + pv_params.mma_o_pipeline.producer_acquire(mma_o_producer_state) + pv_params.p_mma_pipeline.consumer_wait(p_mma_consumer_state) + load_kv_pipeline = common_params.load_kv_pipeline + for p_stage in range(self.iterations_pv_k): + accumulate_flag = tiled_mma_pv.get(tcgen05.Field.ACCUMULATE) + for acc_stage in range(self.iterations_pv_n): + load_kv_pipeline.consumer_wait(load_kv_consumer_state) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, accumulate_flag) + vc_stage = load_kv_consumer_state.index + tOtO = pv_params.tOtO_staged[None, None, None, acc_stage] + for k_block in cutlass.range(pv_params.tOrP.shape[2]): + cute.gemm( + tiled_mma_pv, + tOtO, + pv_params.tOrP[ + None, + None, + k_block, + (p_stage, p_mma_consumer_state.index), + ], + pv_params.tOrVC[None, None, k_block, vc_stage], + tOtO, + ) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, True) + load_kv_pipeline.consumer_release(load_kv_consumer_state) + load_kv_consumer_state.advance() + pv_params.p_mma_pipeline.consumer_release(p_mma_consumer_state) + p_mma_consumer_state.advance() + pv_params.mma_o_pipeline.producer_commit(mma_o_producer_state) + mma_o_producer_state.advance() + + return ( + tiled_mma_pv, + load_kv_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def compute( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, + pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + + :return: The MMA s consumer state, the P MMA producer state, and the P correction producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) + + row_max = -self.acc_dtype.inf + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # no mask applied + while k_tile_count > 1: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + False, + False, + ) + k_index = k_index + 1 + k_tile_count = k_tile_count - 1 + + # mask applied + if cutlass.const_expr(common_params.mAccO is not None): + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + k_index == k_tile_total - 1, + True, + ) + else: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + True, + True, + ) + + return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state + + @cute.jit + def correction( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + p_cor_consumer_state: pipeline.PipelineState, + mma_o_consumer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + :param mma_o_consumer_state: The MMA o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, and the MMA o consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction = ( + self.get_correction_factor(common_params, p_cor_consumer_state)) + if k_tile_count_init != k_tile_count: + mma_o_consumer_state = self.rescale( + common_params, + mma_o_consumer_state, + correction_factor, + no_correction, + ) + k_tile_count = k_tile_count - 1 + if k_tile_count == 0: + mma_o_consumer_state = self.epilogue( + common_params, + epilogue_params, + mma_o_consumer_state, + row_sum, + row_max, + ) + + return p_cor_consumer_state, mma_o_consumer_state + + @cute.jit + def exchange_p_cor_metadata( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + correction_factor: cutlass.Float32, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + row_max_new: cutlass.Float32, + tAcc: cute.Tensor, + tidx: cutlass.Int32, + p_cor_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Compute the correction factor for the last k tile.""" + no_correction = 0 + if ( + row_max_new - row_max + ) * softmax_params.softmax_scale_log2 <= self.skip_correction_threshold: + no_correction = 1 + row_max_new = row_max + + # pad for 4x32b + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.mma_s_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, + corr_layout, + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype) + corr_tmem_store_tiled_copy = tcgen05.make_tmem_copy( + corr_tmem_store_atom, tCor) + corr_tmem_store_thr_copy = corr_tmem_store_tiled_copy.get_slice(tidx) + cCor_for_copy = corr_tmem_store_thr_copy.partition_S(cCor) + tCor_for_copy = corr_tmem_store_thr_copy.partition_D(tCor) + rCor = cute.make_fragment_like(cCor_for_copy[None, None, None, 0], + self.acc_dtype) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout) + rCor[0] = row_sum + rCor[1] = row_max_new + rCor[2] = correction_factor + rCor_int[3] = no_correction + + cute.copy( + corr_tmem_store_tiled_copy, + rCor, + tCor_for_copy[None, None, None, p_cor_producer_state.index], + ) + # fence between tmem store and correction warp + cute.arch.fence_view_async_tmem_store() + common_params.p_cor_pipeline.producer_commit(p_cor_producer_state) + p_cor_producer_state.advance() + return p_cor_producer_state, row_max_new + + @cute.jit + def softmax( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + row_max: cutlass.Float32, + row_sum: cutlass.Float32, + correction_factor: cutlass.Float32, + is_last_tile: bool, + is_local_last_tile: cutlass.Boolean, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ]: + """Softmax for one k-tile. Updates the related pipeline states and returns the computed results. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + :param row_max: The row max + :type row_max: cutlass.Float32 + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param is_last_tile: Whether the last tile + :type is_last_tile: bool + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean + + :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] + """ + + softmax_params.p_mma_pipeline.producer_acquire(p_mma_producer_state) + softmax_params.mma_s_pipeline.consumer_wait(mma_s_consumer_state) + + # load S from tmem + tStS_shape = softmax_params.tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1])) + tStS_staged_fake = softmax_params.tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage)) + tStS_staged = cute.make_tensor(common_params.tmem_ptr, + tStS_staged_fake.layout) + tStS = tStS_staged[None, None, None, mma_s_consumer_state.index] + + tAcc = tStS[(None, None), 0, 0] + cta_qk_tiler = ( + self.mma_qk_tiler[0] // self.cluster_shape_mnk[0], + self.mma_qk_tiler[1], + self.mma_qk_tiler[2], + ) + cS = cute.make_identity_tensor(cute.select(cta_qk_tiler, mode=[0, 1])) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.acc_dtype) + tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + + tidx = common_params.tidx % (self.num_compute_warps * + self.threads_per_warp) + + tmem_thr_copy = tmem_tiled_copy.get_slice(tidx) + tTR_tAcc = tmem_thr_copy.partition_S(tAcc) + tTR_tS = tmem_thr_copy.partition_D(cS) + + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if is_last_tile: + tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) else -self.acc_dtype.inf) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, + row_max_new, 0) + + elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + tmem_load_red_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(64), + redOp=tcgen05.TmemLoadRedOp.MAX), + self.acc_dtype, + ) + tmem_red_tiled_copy = tcgen05.make_tmem_copy( + tmem_load_red_atom, tAcc) + tmem_red_thr_copy = tmem_red_tiled_copy.get_slice(tidx) + tTR_tAcc_red = tmem_red_thr_copy.partition_S(tAcc) + tTR_tS_red = tmem_red_thr_copy.partition_D(cS) + tTR_rAcc_red = cute.make_fragment_like(tTR_tS_red, self.acc_dtype) + tTR_rMax = cute.make_rmem_tensor( + cute.make_layout((1, tTR_tS_red.shape[1], tTR_tS_red.shape[2])), + self.acc_dtype, + ) + cute.copy( + tmem_red_tiled_copy, + tTR_tAcc_red, + (tTR_rAcc_red, tTR_rMax), + ) + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if is_last_tile: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) else -self.acc_dtype.inf) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, + row_max_new, 0) + else: + row_max_new = cute.arch.fmax(row_max_new, tTR_rMax[0]) + + # if warps in N is 2, reduce row_max across warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_max_new + self.softmax_exchange_sync_bar.wait() + row_max_new = cute.arch.fmax( + row_max_new, + common_params.smem_exchange[(tidx + 64) % + (self.num_compute_warps * + self.threads_per_warp)], + ) + + # find correction factor + correction_factor = cute.math.exp2( + (row_max - row_max_new) * softmax_params.softmax_scale_log2, + fastmath=True) + # split kv case + if cutlass.const_expr(not is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # softmax + fma_b = softmax_params.softmax_scale_log2 + fma_c = (0.0 - row_max_new) * softmax_params.softmax_scale_log2 + + for i in cutlass.range(cute.size(tTR_rAcc), + vectorize=True, + unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * fma_b + fma_c + tTR_rAcc[i] = cute.math.exp2(tTR_rAcc[i], fastmath=True) + + tTR_rS = cute.make_fragment_like(tTR_tS, self.q_dtype) + + # quantize + tTR_rS.store(tTR_rAcc.load().to(self.q_dtype)) + + # create sP + sP = softmax_params.sP[None, None, None, + (None, p_mma_producer_state.index)] + sP_mk_view = cute.make_tensor( + sP.iterator, + cute.make_layout( + ( + (sP.shape[0][0], sP.shape[1]), + (sP.shape[0][1], sP.shape[2], sP.shape[3]), + ), + stride=( + (sP.stride[0][0], sP.stride[1]), + (sP.stride[0][1], sP.stride[2], sP.stride[3]), + ), + ), + ) + # change to PISL + sP_wo_swizzle_iter = cute.recast_ptr(sP.iterator, swizzle_=None) + swizzle_bits = (int( + math.log2(self.mma_pv_tiler[2] * self.q_dtype.width // 8 // 32)) + + 1) + swizzle_base = 3 if self.q_dtype.width == 16 else 4 + sP_swizzle = cute.make_swizzle(swizzle_bits, swizzle_base, 3) + sP_mk_view = cute.make_tensor( + sP_wo_swizzle_iter, + cute.make_composed_layout(sP_swizzle, 0, sP_mk_view.layout), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, + tmem_tiled_copy) + smem_thr_copy = smem_tiled_copy.get_slice(tidx) + rP_copy_view = smem_thr_copy.retile(tTR_rS) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + + # fence between smem store and mma o + cute.arch.fence_view_async_shared() + softmax_params.p_mma_pipeline.producer_commit(p_mma_producer_state) + p_mma_producer_state.advance() + + # row_sum, using `add_packed_f32x2` to reduce the number of instructions + row_sum = row_sum * correction_factor + row_sum_vec = (0.0, 0.0) + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + row_sum_vec = cute.arch.add_packed_f32x2( + row_sum_vec, (tTR_rAcc[i], tTR_rAcc[i + 1])) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + + # split kv case + if cutlass.const_expr(is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # store correction factor/row_sum/row_max to tmem for correction warp + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # fence between tmem load and mma s + cute.arch.fence_view_async_tmem_load() + + softmax_params.mma_s_pipeline.consumer_release(mma_s_consumer_state) + mma_s_consumer_state.advance() + + return ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max_new, + row_sum, + correction_factor, + ) + + @cute.jit + def _tmem_load_partition( + self, common_params: SimpleNamespace, tiled_mma_pv: cute.TiledMma, + iter_n: int + ) -> tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, + cute.TiledMma]: + """Tensor memory load partition for rescale and epilogue. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param iter_n: The iteration number + :type iter_n: int + + :return: The tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv + :rtype: tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma] + """ + + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1])) + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO = cute.make_tensor(common_params.tmem_ptr + self.tmem_o_offset, + tOtO_layout) + tOtO = tOtO[None, None, None, iter_n] + + tAcc = tOtO[(None, None), 0, 0] + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.acc_dtype) + tmem_load_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + tmem_load_thr_copy = tmem_load_tiled_copy.get_slice( + common_params.tidx % + (self.num_compute_warps * self.threads_per_warp)) + + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + # Flatten divide and partition global tensors for O + cta_pv_tiler_mn = cute.select(cta_pv_tiler, mode=[0, 1]) + + gO = None + if cutlass.const_expr(common_params.mAccO is not None): + gO = cute.local_tile( + common_params.mAccO[None, common_params.blk_coord[3], None, + None, None], + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor( + common_params.mAccO[None, common_params.blk_coord[3], None, + None, None].shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + else: + gO = cute.local_tile( + common_params.mO, + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor(common_params.mO.shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + tTR_tAcc = tmem_load_thr_copy.partition_S(tAcc) + tTR_gO = tmem_load_thr_copy.partition_D(gO) + tTR_cO = tmem_load_thr_copy.partition_D(cO) + tTR_rAcc = cute.make_fragment_like(tTR_gO, self.acc_dtype) + return tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc + + def get_correction_factor( + self, + common_params: SimpleNamespace, + p_cor_consumer_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Int32, + ]: + """Get the correction factor from the P correction consumer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, the row_sum, the row_max, and the correction factor + :rtype: tuple[pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Int32] + """ + common_params.p_cor_pipeline.consumer_wait(p_cor_consumer_state) + tidx = common_params.tidx % (self.num_compute_warps * + self.threads_per_warp) + # load correction factor + _, tAcc, _, _, _, _ = self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, 0) + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.p_cor_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, corr_layout) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype) + corr_tmem_load_tiled_copy = tcgen05.make_tmem_copy( + corr_tmem_load_atom, tCor) + corr_tmem_load_thr_copy = corr_tmem_load_tiled_copy.get_slice(tidx) + tCor_for_copy = corr_tmem_load_thr_copy.partition_S(tCor) + cCor_for_copy = corr_tmem_load_thr_copy.partition_D(cCor) + rCor = cute.make_fragment_like(cCor_for_copy[None, None, None, 0], + self.acc_dtype) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout) + cute.copy( + corr_tmem_load_tiled_copy, + tCor_for_copy[None, None, None, p_cor_consumer_state.index], + rCor, + ) + row_sum = rCor[0] + row_max = rCor[1] + correction_factor = rCor[2] + no_correction = rCor_int[3] + + common_params.p_cor_pipeline.consumer_release(p_cor_consumer_state) + p_cor_consumer_state.advance() + return p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction + + @cute.jit + def rescale( + self, + common_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + correction_factor: cutlass.Float32, + no_correction: cutlass.Int32, + ) -> pipeline.PipelineState: + """Rescale for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param no_correction: Whether to apply correction factor + :type no_correction: cutlass.Int32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + skip_correction = cute.arch.vote_all_sync(no_correction == 1) + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + if not skip_correction: + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition(common_params, + common_params.tiled_mma_pv, + iter_n)) + + # tmem store tiled copy + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), + self.acc_dtype) + tmem_store_tiled_copy = tcgen05.make_tmem_copy( + tmem_store_atom, tAcc) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + # rescale, using `mul_packed_f32x2` to reduce the number of instructions + for i in cutlass.range(cute.size(tTR_rAcc), + vectorize=True, + unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * correction_factor + + # store o to tensor memory for next k tile + cute.copy(tmem_store_tiled_copy, tTR_rAcc, tTR_tAcc) + + cute.arch.fence_view_async_tmem_store() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + @cute.jit + def epilogue( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + ) -> pipeline.PipelineState: + """Epilogue for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param row_max: The row max + :type row_max: cutlass.Float32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + + tidx = common_params.tidx % (self.num_compute_warps * + self.threads_per_warp) + + # exchange row_sum between warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_sum + self.epilogue_exchange_sync_bar.wait() + # (64, 2) + row_sum = (row_sum + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp)]) + # mma_o pipeline consumer wait + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition(common_params, + common_params.tiled_mma_pv, iter_n)) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + + # apply output scale and normalize by row_sum + for i in cutlass.range(cute.size(tTR_rAcc), + vectorize=True, + unroll_full=True): + tTR_rAcc[i] = (tTR_rAcc[i] * epilogue_params.output_scale * + cute.arch.rcp_approx(row_sum)) + + # store o to global memory + tR2G_rO_src = None + tR2G_rO_dst = tTR_gO + if cutlass.const_expr(common_params.mAccO is None): + tR2G_rO_src = cute.make_fragment_like(tTR_gO, self.o_dtype) + # using final output dtype for o + tR2G_rO_src.store(tTR_rAcc.load().to(self.o_dtype)) + else: + # using accumulate dtype for o + tR2G_rO_src = tTR_rAcc + + if cute.elem_less(tTR_cO[0][0], common_params.H): + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority. + NO_ALLOCATE, + ) + + # store the lse to global memory + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + gLSE = None + cLSE = None + if cutlass.const_expr(epilogue_params.mAccLSE is None): + gLSE = cute.local_tile( + epilogue_params.mLSE, + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor(epilogue_params.mLSE.shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + + else: + gLSE = cute.local_tile( + epilogue_params.mAccLSE[None, common_params.blk_coord[3], + None, None], + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor( + epilogue_params.mAccLSE[None, + common_params.blk_coord[3], + None, None].shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + lse = (cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max) + if cutlass.const_expr(self.warps_in_n == 2): + if cute.elem_less(cLSE[tidx][0], common_params.H): + gLSE[tidx] = lse + + cute.arch.fence_view_async_tmem_load() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + def make_and_init_load_pt_pipeline(self, load_pt_mbar_ptr): + """Create and initialize the load page table pipeline. + + :param load_pt_mbar_ptr: The load page table mbar pointer + :type load_pt_mbar_ptr: cute.Tensor + + :return: The load page table pipeline + :rtype: pipeline.PipelineAsync + """ + load_pt_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len([self.load_pt_warp_id]), + ) + load_pt_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len([self.load_tma_warp_id]), + ) + return pipeline.PipelineCpAsync.create( + barrier_storage=load_pt_mbar_ptr, + num_stages=self.load_pt_stage, + producer_group=load_pt_producer_group, + consumer_group=load_pt_consumer_group, + defer_sync=True, + ) + + def make_and_init_load_qkv_pipeline(self, load_qkv_mbar_ptr, + cta_layout_vmnk, load_stages, + tx_count) -> pipeline.PipelineTmaUmma: + """Create and initialize the tma load qkv pipeline. + + :param load_qkv_mbar_ptr: The load qkv mbar pointer + :type load_qkv_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + :param load_stages: The load stages + :type load_stages: list[int] + :param tx_count: The tx count + :type tx_count: int + + :return: The tma load qkv pipeline + :rtype: pipeline.PipelineTmaUmma + """ + load_qkv_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_tma_warp_id])) + load_qkv_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_qkv_mbar_ptr, + num_stages=load_stages, + producer_group=load_qkv_producer_group, + consumer_group=load_qkv_consumer_group, + tx_count=tx_count, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_mma_s_pipeline( + self, mma_s_mbar_ptr, + cta_layout_vmnk) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma s pipeline. + + :param mma_s_mbar_ptr: The mma s mbar pointer + :type mma_s_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma s pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_s_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + consumer_thread_size = (self.threads_per_warp * + len(self.compute_warp_ids) * + self.cluster_shape_mnk[0]) + mma_s_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_s_mbar_ptr, + num_stages=self.mma_s_stage, + producer_group=mma_s_producer_group, + consumer_group=mma_s_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_mma_pipeline( + self, p_mma_mbar_ptr, + cta_layout_vmnk) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p mma pipeline. + + :param p_mma_mbar_ptr: The p mma mbar pointer + :type p_mma_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The p mma pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = (self.threads_per_warp * + len(self.compute_warp_ids) * + self.cluster_shape_mnk[0]) + p_mma_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_mma_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + return pipeline.PipelineAsyncUmma.create( + barrier_storage=p_mma_mbar_ptr, + num_stages=self.p_mma_stage, + producer_group=p_mma_producer_group, + consumer_group=p_mma_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_cor_pipeline( + self, p_cor_mbar_ptr) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p correction pipeline. + + :param p_cor_mbar_ptr: The p correction mbar pointer + :type p_cor_mbar_ptr: cute.Tensor + + :return: The p correction pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = self.threads_per_warp * len( + self.compute_warp_ids) + p_cor_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_cor_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + return pipeline.PipelineAsync.create( + barrier_storage=p_cor_mbar_ptr, + num_stages=self.p_cor_stage, + producer_group=p_cor_producer_group, + consumer_group=p_cor_consumer_group, + defer_sync=True, + ) + + def make_and_init_mma_o_pipeline( + self, mma_o_mbar_ptr, + cta_layout_vmnk) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma o pipeline. + + :param mma_o_mbar_ptr: The mma o mbar pointer + :type mma_o_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma o pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_o_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + consumer_thread_size = (self.threads_per_warp * + len(self.compute_warp_ids) * + self.cluster_shape_mnk[0]) + mma_o_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_o_mbar_ptr, + num_stages=self.mma_o_stage, + producer_group=mma_o_producer_group, + consumer_group=mma_o_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + @staticmethod + def _compute_grid( + o: cute.Tensor, + split_kv: cutlass.Int32, + cluster_shape_mnk: Tuple[int, int, int], + max_active_clusters: int, + is_persistent: bool, + ) -> Tuple[MLAStaticTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid shape for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + + :return: Tile scheduler parameters and grid shape. + :rtype: tuple[MLAStaticTileSchedulerParams, tuple[int, int, int]] + """ + o_shape = o.shape + tile_sched_params = create_mla_static_tile_scheduler_params( + is_persistent, + cute.size(o_shape[3]), + cute.size(o_shape[2]), + cluster_shape_mnk, + split_kv, + ) + grid = MLAStaticTileScheduler.get_grid_shape(tile_sched_params, + max_active_clusters) + + return tile_sched_params, grid + + @staticmethod + def get_workspace_size( + H: int, + S: int, + D: int, + B: int, + split_kv: int, + acc_dtype: Type[cutlass.Numeric], + ) -> int: + """Get the extra workspace(device memory) size for the MLA kernel when split_kv is not 1. + + :param H: The height of the output tensor C + :type H: int + :param S: The sequence length of the output tensor C + :type S: int + :param D: The depth of the output tensor C + :type D: int + :param B: The batch size of the output tensor C + :type B: int + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + + :return: The workspace size for the MLA kernel + :rtype: int + """ + if split_kv == 1: + return 0 + return B * H * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + acc_dtype: Type[cutlass.Numeric], + workspace: cute.Tensor, + ) -> tuple[cute.Tensor, cute.Tensor]: + """Initialize the workspace for the MLA kernel. Construct the intermediate tensors + acc_o and acc_lse. + + :param H: The height of the output tensor C + :type H: cutlass.Int32 + :param D: The depth of the output tensor C + :type D: cutlass.Int32 + :param S: The sequence length of the output tensor C + :type S: cutlass.Int32 + :param B: The batch size of the output tensor C + :type B: cutlass.Int32 + :param split_kv: The split key-value of the output tensor C + :type split_kv: cutlass.Int32 + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + :param workspace: The workspace tensor + :type workspace: cute.Tensor + + :return: The output tensor C and the workspace tensor + :rtype: tuple[cute.Tensor, cute.Tensor] + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + align = 256 // self.q_dtype.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=acc_dtype) + 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 + + cute.cosize(acc_o_layout) * acc_dtype.width // 8, + dtype=acc_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @staticmethod + def can_implement( + B: int, + S: int, + K: int, + H: int, + L: int, + R: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + ) -> bool: + """Check if the MLA kernel can be implemented. + + :param B: The batch size of the output tensor C + :type B: int + :param S: The sequence length of the output tensor C + :type S: int + :param K: The width of the output tensor KV + :type K: int + :param H: The number of heads of the output tensor C + :type H: int + :param L: The number of latent dimensions of the tensor KV + :type L: int + :param R: The number of rope dimensions of the tensor C_rope + :type R: int + :param in_dtype: The data type of the input tensor + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: The data type of the output tensor + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: The data type of the log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: The tile shape of the query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: The tile shape of the probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: The page size of the page table + :type page_size: int + + :return: Whether the MLA kernel can be implemented + :rtype: bool + """ + if L != 512 or R != 64: + return False + if in_dtype not in [cutlass.Float16]: + return False + if out_dtype not in [cutlass.Float16]: + return False + if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: + return False + # page size equals 1 is prohibited by tma specification, not 128B aligned. + if mma_qk_tiler_mn[1] % page_size != 0 or page_size == 1: + return False + if mma_qk_tiler_mn[0] != mma_pv_tiler_mn[0] or mma_qk_tiler_mn[0] != 128: + return False + if is_var_split_kv and not is_var_seq: + return False + if H > 128 or (H < 128 and split_kv != 1): + return False + if S < 1 or S > 4: + return False + if K <= 0: + return False + return True + + +def run( + batch_size: int, + seq_len_q: int, + seq_len_k: int, + num_heads: int, + latent_dim: int, + rope_dim: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + softmax_scale: float, + output_scale: float, + skip_correction_threshold: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool, + **kwargs, +): + """Execute Multi-Head Latent Attention (MLA) on Blackwell architecture and validate results. + + This function creates random input tensors for query latent/rope, compressed latent/rope, and value, + then performs the complete MLA computation pipeline. It supports configurable data types, tiling parameters, + page table, variable sequence length, and variable split_kv. Results can be validated against a PyTorch reference + implementation or run multiple times for performance measurement. + + :param batch_size: Batch size + :type batch_size: int + :param seq_len_q: Sequence length of Q + :type seq_len_q: int + :param seq_len_k: Sequence length of K + :type seq_len_k: int + :param num_heads: Number of heads + :type num_heads: int + :param latent_dim: dimension of query/compressed latent + :type latent_dim: int + :param rope_dim: dimension of query/compressed rope + :type rope_dim: int + :param in_dtype: Input data type for query/compressed latent/rope tensors + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: Output data type for attention output + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for query-key matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Accumulator data type for log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: Matrix multiply accumulate tile shape (M, N) for query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: Matrix multiply accumulate tile shape (M, N) for probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: Split key-value + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: Page size of the page table + :type page_size: int + :param softmax_scale: Attention score scaling factor + :type softmax_scale: float + :param output_scale: Output scaling factor + :type output_scale: float + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param tolerance: Maximum acceptable error for validation + :type tolerance: float + :param warmup_iterations: Number of warmup iterations + :type warmup_iterations: int + :param iterations: Number of iterations to run for performance testing + :type iterations: int + :param skip_ref_check: Skip validation against reference implementation + :type skip_ref_check: bool + :param use_cold_l2: Whether to use cold L2 cache + :type use_cold_l2: bool + + :raises ValueError: If input shapes are incompatible or head dimension is unsupported + :raises RuntimeError: If GPU is unavailable for computation + """ + + print("Running Blackwell MLA test with:") + print(f" batch_size: {batch_size}") + print(f" seq_len_q: {seq_len_q}") + print(f" seq_len_k: {seq_len_k}") + print(f" num_heads: {num_heads}") + print(f" latent_dim: {latent_dim}") + print(f" rope_dim: {rope_dim}") + print(f" in_dtype: {in_dtype}") + print(f" out_dtype: {out_dtype}") + print(f" acc_dtype: {acc_dtype}") + print(f" mma_qk_tiler_mn: {mma_qk_tiler_mn}") + print(f" mma_pv_tiler_mn: {mma_pv_tiler_mn}") + print(f" split_kv: {split_kv}") + print(f" is_persistent: {is_persistent}") + print(f" is_var_seq: {is_var_seq}") + print(f" is_var_split_kv: {is_var_split_kv}") + print(f" page_size: {page_size}") + print(f" softmax_scale: {softmax_scale}") + print(f" output_scale: {output_scale}") + print(f" skip_correction_threshold: {skip_correction_threshold}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + + import cutlass.torch as cutlass_torch + import torch + + # Prepare pytorch tensors: Q, K, V (random from 0 to 2) and O (all zero) + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not BlackwellMultiHeadLatentAttentionForwardFP16.can_implement( + batch_size, + seq_len_q, + seq_len_k, + num_heads, + latent_dim, + rope_dim, + in_dtype, + out_dtype, + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + split_kv, + is_persistent, + is_var_seq, + is_var_split_kv, + page_size, + ): + raise TypeError( + f"Unsupported testcase {batch_size}, {seq_len_q}, {seq_len_k}, {num_heads}, {latent_dim}, {rope_dim}, {in_dtype}, {out_dtype}, {acc_dtype}, {lse_dtype}, {mma_qk_tiler_mn}, {mma_pv_tiler_mn}, {split_kv}, {is_persistent}, {is_var_seq}, {is_var_split_kv}, {page_size}" + ) + + torch.manual_seed(1111) + + def create_data_tensor( + B, + HK, + D, + dtype, + is_dynamic_layout=True, + page_table=None, + cache_seqs=None, + is_lse=False, + seq_len_q=None, + ): + shape = (B, HK, D) + if page_table is not None: + if cache_seqs is not None: + max_seq_len = torch.max(cache_seqs) + shape = (B * ceil_div(max_seq_len, page_size), page_size, D) + else: + shape = (B * ceil_div(HK, page_size), page_size, D) + + if seq_len_q is not None: + shape = (B, seq_len_q, HK, D) + + permute_order = (1, 2, 0) + stride_order = (2, 0, 1) + leading_dim = 1 + if is_lse: + shape = (B, seq_len_q, HK) + permute_order = (2, 1, 0) + stride_order = (2, 1, 0) + leading_dim = 0 + elif seq_len_q is not None: + permute_order = (2, 3, 1, 0) + stride_order = (3, 2, 0, 1) + leading_dim = 1 + + init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) + + torch_dtype = (cutlass_torch.dtype(dtype) + if dtype != cutlass.Float8E4M3FN else torch.int8) + + # Create dtype torch tensor (cpu) + torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch_dtype, + permute_order=permute_order, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=init_config, + ) + + # Create dtype torch tensor (gpu) + torch_tensor_gpu = torch_tensor_cpu.cuda() + + # Create f32 torch tensor (cpu) + f32_torch_tensor = torch_tensor_cpu.to(dtype=torch.float32) + + # Create dtype cute tensor (gpu) + cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16) + cute_tensor.element_type = dtype + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic( + leading_dim=leading_dim) + if not is_lse: + cute_tensor = cute_tensor.mark_compact_shape_dynamic( + mode=leading_dim, + stride_order=stride_order, + divisibility=(128 // dtype.width), + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=is_dynamic_layout, + ) + + return f32_torch_tensor, cute_tensor, torch_tensor_gpu + + def create_cache_seqs(batch_size, seq_len_k, is_var_seq): + cache_seqs_ref = torch.ones(batch_size, dtype=torch.int32) * seq_len_k + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack(cache_seqs_gpu, + assumed_align=16).mark_layout_dynamic() + if is_var_seq: + max_seq_len = seq_len_k + min_seq_len = int(seq_len_k * 0.8) + cache_seqs_ref = cutlass_torch.create_and_permute_torch_tensor( + (batch_size, ), + torch.int32, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig(min_val=min_seq_len, + max_val=max_seq_len + + 1), + ) + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack( + cache_seqs_gpu, + assumed_align=16, + ).mark_layout_dynamic() + return cache_seqs_ref, cache_seqs, cache_seqs_gpu + + def create_page_table(batch_size, seq_len_k, is_var_seq, page_size): + max_seq_len = seq_len_k if not is_var_seq else torch.max(cache_seqs_ref) + page_count = ceil_div(max_seq_len, page_size) + page_table_ref = torch.empty([batch_size, page_count], + dtype=torch.int32) + # use transposed index for page table to make sure the value is in bound of `batch_size * seq_len_block`. In practice, the value could be any positive values. This setting is only for testing purpose. + for b in range(batch_size): + for j in range(page_count): + page_table_ref[b, j] = b + j * batch_size + page_table_gpu = page_table_ref.permute(1, 0).cuda() + page_table = from_dlpack( + page_table_gpu, assumed_align=16).mark_layout_dynamic(leading_dim=0) + return page_table_ref, page_table, page_table_gpu + + def create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ): + block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu = None, None, None + # check if split_kv is valid otherwise do auto setting of split_kv + if is_var_split_kv: + block_split_kvs_ref = torch.zeros([batch_size], dtype=torch.int32) + for b in range(batch_size): + block_split_kvs_ref[b] = ( + BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[b].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + )) + split_kv = torch.max(block_split_kvs_ref).item() + block_split_kvs_gpu = block_split_kvs_ref.cuda() + block_split_kvs = from_dlpack( + block_split_kvs_gpu, assumed_align=16).mark_layout_dynamic() + elif split_kv <= 0: + split_kv = BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[0].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + return split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu + + def create_workspace(num_heads, seq_len_q, latent_dim, batch_size, split_kv, + acc_dtype): + workspace_size = ( + BlackwellMultiHeadLatentAttentionForwardFP16.get_workspace_size( + num_heads, + seq_len_q, + latent_dim, + batch_size, + split_kv, + acc_dtype, + )) + + workspace, workspace_torch = None, None + if workspace_size > 0: + workspace_torch = torch.empty([workspace_size], + dtype=torch.int8).cuda() + workspace = from_dlpack(workspace_torch, assumed_align=32) + return workspace, workspace_torch + + cache_seqs_ref, cache_seqs, cache_seqs_torch = create_cache_seqs( + batch_size, seq_len_k, is_var_seq) + page_table_ref, page_table, page_table_torch = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size) + cluster_shape_mnk = (2, 1, 1) + hardware_info = utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mnk[0] * cluster_shape_mnk[1]) + split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_torch = ( + create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + )) + + q_latent_ref, q_latent, q_latent_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + q_rope_ref, q_rope, q_rope_torch = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + c_latent_ref, c_latent, c_latent_torch = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + c_rope_ref, c_rope, c_rope_torch = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + o_ref, o, o_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + lse_ref, lse, lse_torch = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace(num_heads, seq_len_q, + latent_dim, batch_size, + split_kv, acc_dtype) + + mla = BlackwellMultiHeadLatentAttentionForwardFP16( + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + max_active_clusters, + page_size, + skip_correction_threshold, + is_persistent, + is_var_seq, + is_var_split_kv, + ) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + stream = cuda.CUstream(torch_stream.cuda_stream) + + # compile mla kernel + compiled_mla = cute.compile( + mla, + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + options="--opt-level 2", + ) + + def torch_reference_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + softmax_scale=1.0, + output_scale=1.0, + ): + # expand and concat q_latent and q_rope to have the dimension of sequence length for q + q_ref = torch.cat([q_latent, q_rope], dim=1).permute(3, 2, 0, 1) + # expand and concat c_latent and c_rope to have the dimension of num_heads for k and v + page_count = page_table_ref.shape[1] + k_ref_paged = (torch.cat([c_latent, c_rope], + dim=1).permute(2, 0, 1).reshape( + batch_size * page_count, page_size, + latent_dim + rope_dim)) + v_ref_paged = c_latent.permute(2, 0, 1).reshape(batch_size * page_count, + page_size, latent_dim) + + if is_var_seq: + max_seq_len = torch.max(cache_seqs_ref) + else: + max_seq_len = seq_len_k + + k_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim + rope_dim]) + v_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim]) + k_ref = torch.index_select( + k_ref_paged, 0, torch.flatten(page_table_ref)).reshape( + batch_size, 1, -1, latent_dim + rope_dim)[:, :, :max_seq_len, :] + v_ref = torch.index_select(v_ref_paged, 0, + torch.flatten(page_table_ref)).reshape( + batch_size, 1, -1, + latent_dim)[:, :, :max_seq_len, :] + for b in range(batch_size): + k_ref[b, :, cache_seqs_ref[b]:, :] = 0 + v_ref[b, :, cache_seqs_ref[b]:, :] = 0 + import torch.nn.functional as F + + o_ref = F.scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + attn_mask=None, + dropout_p=0.0, + scale=softmax_scale, + is_causal=False, + ) + s_ref = torch.einsum("bhld,bhsd->bhls", q_ref, k_ref) + s_ref_max, s_ref_max_pos = torch.max(s_ref, dim=-1, keepdim=True) + softmax_scale_log2 = LOG2_E * softmax_scale + s_ref_sum = torch.sum(torch.exp2( + (s_ref - s_ref_max) * softmax_scale_log2), + dim=-1, + keepdim=True) + + lse_ref = s_ref_max * softmax_scale_log2 + torch.log2(s_ref_sum) + lse_ref = lse_ref.squeeze(3).permute(2, 1, 0) + o_ref = o_ref * output_scale + o_ref = o_ref.permute(2, 3, 1, 0) + + return o_ref, lse_ref + + if skip_correction_threshold > 0.0: + print( + "Skipping correction verification since skip_correction_threshold is greater than 0.0..." + ) + skip_ref_check = True + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + torch.cuda.synchronize() + + print("Verifying results...") + if in_dtype == cutlass.Float8E4M3FN: + tolerance = 0.13 + o_ref, lse_ref = torch_reference_mla( + q_latent_ref, + q_rope_ref, + c_latent_ref, + c_rope_ref, + page_table, + cache_seqs, + softmax_scale, + output_scale, + ) + + if out_dtype in [cutlass.Float8E5M2, cutlass.Float8E4M3FN]: + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + cutlass.Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o, o_fp32) + o = o_fp32_torch.cpu() + ref_fp8, _ = cutlass_torch.cute_tensor_like( + torch.empty(*o_ref.permute(3, 2, 0, 1).shape, + dtype=torch.uint8).permute(2, 3, 1, 0), + out_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + o_ref_gpu = o_ref.cuda() + o_ref_f32 = from_dlpack(o_ref_gpu).mark_layout_dynamic( + leading_dim=1) + + # convert ref : f32 -> fp8 -> f32 + cute.testing.convert(o_ref_f32, ref_fp8) + cute.testing.convert(ref_fp8, o_ref_f32) + + o_ref = o_ref_gpu.cpu() + else: + o = o_torch.cpu().to(torch.float32) + lse = lse_torch.cpu() + lse_ref = lse_ref.to(cutlass.torch.dtype(lse_dtype)) + # Assert close results + torch.testing.assert_close(o, o_ref, atol=tolerance, rtol=1e-05) + torch.testing.assert_close(lse, lse_ref, atol=tolerance, rtol=1e-05) + print("Results verified successfully!") + + def generate_tensors(): + _, cache_seqs, _ = create_cache_seqs(batch_size, seq_len_k, is_var_seq) + _, page_table, _ = create_page_table(batch_size, seq_len_k, is_var_seq, + page_size) + _split_kv, _, block_split_kvs, _ = create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + + _, q_latent, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, q_rope, _ = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + _, c_latent, _ = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, c_rope, _ = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, o, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, lse, _ = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace(num_heads, seq_len_q, + latent_dim, batch_size, + _split_kv, acc_dtype) + return testing.JitArguments( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + _split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + q_latent_torch.numel() * q_latent_torch.element_size() + + q_rope_torch.numel() * q_rope_torch.element_size() + + c_latent_torch.numel() * c_latent_torch.element_size() + + c_rope_torch.numel() * c_rope_torch.element_size() + + o_torch.numel() * o_torch.element_size() + + lse_torch.numel() * lse_torch.element_size() + + cache_seqs_torch.numel() * cache_seqs_torch.element_size()) + one_workspace_bytes += (page_table_torch.numel() * + page_table_torch.element_size()) + if is_var_split_kv: + one_workspace_bytes += (block_split_kvs_torch.numel() * + block_split_kvs_torch.element_size()) + if workspace_torch is not None: + one_workspace_bytes += (workspace_torch.numel() * + workspace_torch.element_size()) + workspace_count = testing.get_workspace_count(one_workspace_bytes, + warmup_iterations, + iterations) + + avg_time_us = testing.benchmark( + compiled_mla, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return avg_time_us # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers.") + + def parse_mma_tiler(s: str) -> Tuple[int, int, Tuple[int, int]]: + ret = parse_comma_separated_ints(s) + if len(ret) != 2: + raise argparse.ArgumentTypeError( + "Invalid format. Expected 2 comma-separated integers.") + return (ret[0], ret[1]) + + parser = argparse.ArgumentParser(description="Example of MLA on Blackwell.") + + parser.add_argument( + "--in_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Input data type", + ) + + parser.add_argument( + "--out_dtype", + type=cutlass.dtype, + default=cutlass.Float16, + help="Output data type", + ) + + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Accumulator data type", + ) + + parser.add_argument( + "--lse_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="LSE data type", + ) + parser.add_argument( + "--mma_qk_tiler_mn", + type=parse_mma_tiler, + default=(128, 128), + help="MMA tile shape (H, K)", + ) + parser.add_argument( + "--mma_pv_tiler_mn", + type=parse_mma_tiler, + default=(128, 256), + help="MMA tile shape (H, D)", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size", + ) + + parser.add_argument( + "--seq_len_q", + type=int, + default=1, + help="Sequence length of Q", + ) + + parser.add_argument( + "--seq_len_k", + type=int, + default=128, + help="Sequence length of K/V", + ) + + parser.add_argument( + "--num_heads", + type=int, + default=128, + help="Number of heads of Q", + ) + + parser.add_argument( + "--latent_dim", + type=int, + default=512, + help="Latent dimension of Q/C", + ) + + parser.add_argument( + "--rope_dim", + type=int, + default=64, + help="Rope dimension of Q/C", + ) + + parser.add_argument( + "--is_var_seq", + action="store_true", + help="Use variable length of sequence length or not", + ) + + parser.add_argument( + "--is_var_split_kv", + action="store_true", + help="Use variable length of split kv or not", + ) + + parser.add_argument( + "--page_size", + type=int, + default=128, + help="Page size of page table", + ) + + parser.add_argument( + "--split_kv", + type=int, + default=-1, + help="Split KV setting", + ) + + parser.add_argument( + "--softmax_scale", + type=float, + default=0.0416, + help="Scaling factor to scale softmax", + ) + + parser.add_argument( + "--output_scale", + type=float, + default=1.0, + help="Scaling factor to scale output", + ) + + parser.add_argument( + "--skip_correction_threshold", + type=float, + default=0.0, + help="Skip correction threshold", + ) + + parser.add_argument("--tolerance", + type=float, + default=1e-02, + help="Tolerance for validation") + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + help="Use cold L2 cache", + ) + + args = parser.parse_args() + + run( + args.batch_size, + args.seq_len_q, + args.seq_len_k, + args.num_heads, + args.latent_dim, + args.rope_dim, + args.in_dtype, + args.out_dtype, + args.acc_dtype, + args.lse_dtype, + args.mma_qk_tiler_mn, + args.mma_pv_tiler_mn, + args.split_kv, + args.is_persistent, + args.is_var_seq, + args.is_var_split_kv, + args.page_size, + args.softmax_scale, + args.output_scale, + args.skip_correction_threshold, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py new file mode 100644 index 000000000000..560284cf17cf --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py @@ -0,0 +1,4212 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import argparse +import math +import os +import sys +from types import SimpleNamespace +from typing import Optional, Tuple, Type + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.cpasync as cpasync +import cutlass.cute.testing as testing +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.arch import Arch +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.nvgpu.tcgen05 import OperandMajorMode +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import BaseDSL +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from .mla_helpers import (LOG2_E, MAX_SPLITS, MLAStaticTileScheduler, + MLAStaticTileSchedulerParams, ceil_div, + create_mla_static_tile_scheduler, + create_mla_static_tile_scheduler_params) + +if __name__ == "__main__": + current_dir = os.path.dirname(os.path.abspath(__file__)) + sys.path.insert(0, os.path.join(current_dir, "../..")) +""" +A Multi-Head Latent Attention (MLA) example using fp8 as input/output for the NVIDIA Blackwell SM100 architecture using CUTE DSL + +This example demonstrates an implementation of inference of multi-head latent attention using a TMA + Blackwell +SM100 TensorCore warp-specialized persistent kernel. The implementation integrates the (Qc + Qr)*(Kc + Kr)^T +matrix multiplication, softmax normalization, and softmax((Qc + Qr)*(Kc + Kr)^T)*Vc into a single kernel. +The kernel provides support for page table storage and variable-length KV cache sequences. It implements KV splitting +functionality to minimize latency when processing long KV sequences. + +The kernel implements key optimizations including: +- Warp specialization for different computation phases (load, MMA, softmax, correction, epilogue) +- Pipeline stages between different warps for overlapping computation and memory access +- Support for different precision data types +- Two sub-kernels (split KV kernel and reduction kernel) that enable split KV processing + +To run this example: + +.. code-block:: bash + + python examples/blackwell/mla_fp8.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float8E4M3FN --out_dtype Float8E4M3FN \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent + +The above example runs Multi-Head Latent Attention (MLA) with the following configuration: +- Batch size: 4 +- Sequence length of Q: 1 +- Sequence length of K: 1024 +- Latent dimension: 512 +- RoPE dimension: 64 +- Number of heads: 128 +- Data types: Float8E4M3FN (input), Float8E4M3FN (output), Float32 (accumulation and LSE) + +It utilizes page table storage for the KV cache and enables both variable-length KV cache sequences +and variable split KV processing with persistent scheduling. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/mla_fp8.py \ + --batch_size 4 --latent_dim 512 --rope_dim 64 \ + --num_heads 128 --seq_len_q 1 --seq_len_k 1024 \ + --in_dtype Float8E4M3FN --out_dtype Float8E4M3FN \ + --acc_dtype Float32 --lse_dtype Float32 \ + --is_var_seq --is_var_split_kv \ + --is_persistent --warmup_iterations 3 \ + --iterations 10 --skip_ref_check + +Constraints for this example: +* Data type requirements: + - Input/output: Float8E4M3FN + - Accumulation and LSE: Float32 +* Fixed architecture parameters: + - Number of attention heads: 128 + - Latent dimension: 512 + - RoPE dimension: 64 +* Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) +* Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +* Query sequence length must be 1-4 +* Only supports 2-CTA instructions +* Variable sequence length requires page table storage enabled +""" + + +class BlackwellMultiHeadLatentAttentionForwardFP8: + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + max_active_clusters: int, + page_size: int, + skip_correction_threshold: float, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + + :param acc_dtype: Data type for accumulation S and O + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Data type for output LSE + :type lse_dtype: Type[cutlass.Numeric] + :param mma_s_tiler: The (H, K) tile shape of the MMA instruction for S + :type mma_s_tiler: Tuple[int, int] + :param mma_p_tiler: The (H, D) tile shape of the MMA instruction for P + :type mma_p_tiler: Tuple[int, int] + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: int + :param page_size: The page size + :type page_size: int + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split KV + :type is_var_split_kv: bool + """ + + self.latent_dim = 512 + self.rope_dim = 64 + 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.skip_correction_threshold = skip_correction_threshold + self.is_persistent = is_persistent + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), + # while warps 2-3 handle accumulation for second half [n/2, n) + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 + mma_qk_tiler_k = self.rope_dim * 2 + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + mma_qk_tiler_k, + ) + self.mma_qk_rope_tiler = ( + self.mma_qk_tiler_mn[0], + self.mma_qk_tiler_mn[1], + self.rope_dim, + ) + self.mma_pv_tiler = ( + self.mma_pv_tiler_mn[0], + self.mma_pv_tiler_mn[1], + self.mma_qk_tiler[1] * self.mma_qk_tiler[2] // + self.mma_pv_tiler_mn[1], + ) + self.iterations_qk_latent = self.latent_dim // self.mma_qk_tiler[2] + self.iterations_qk_rope = 1 + self.iterations_qk = self.iterations_qk_latent + self.iterations_qk_rope + self.iterations_pv_k = self.mma_qk_tiler[1] // self.mma_pv_tiler[2] + self.iterations_pv_n = self.latent_dim // self.mma_pv_tiler[1] + + # Set specialized warp ids + self.compute_warp_ids = (0, 1, 2, 3) + self.correction_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.load_tma_k_warp_id = 9 + self.load_tma_v_warp_id = 10 + self.empty_warp_ids = (11, ) + self.threads_per_cta = self.threads_per_warp * len(( + self.mma_warp_id, + self.load_tma_k_warp_id, + self.load_tma_v_warp_id, + *self.compute_warp_ids, + *self.correction_warp_ids, + *self.empty_warp_ids, + )) + + # register settings + self.softmax_reg_num = 192 + self.correction_reg_num = 256 + self.other_reg_num = 48 + # Named barriers + self.tmem_ptr_sync_bar = pipeline.NamedBarrier( + barrier_id=1, + num_threads=(self.threads_per_warp + + self.threads_per_warp * self.num_compute_warps * 2), + ) + self.softmax_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=2, + num_threads=(self.threads_per_warp * self.num_compute_warps)) + self.epilogue_exchange_sync_bar = pipeline.NamedBarrier( + barrier_id=3, + num_threads=(self.threads_per_warp * self.num_compute_warps)) + + def _setup_attributes(self): + """Set up configurations and parameters for the MLA kernel operation. + + This method initializes and configures various attributes required for the + execution of the multi-head latent attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + self.load_q_stage = 1 + self.load_k_stage = 3 + self.load_v_stage = 2 + self.mma_s_stage = 2 + self.p_mma_stage = 2 + self.p_cor_stage = 2 + self.mma_o_stage = 2 + + self.tmem_o_offset = self.mma_s_stage * self.mma_qk_tiler[ + 1] // self.warps_in_n + self.correction_factor_offset = (self.tmem_o_offset + + self.latent_dim // self.warps_in_n) + + @cute.jit + def __call__( + self, + q_latent: cute.Tensor, + q_rope: cute.Tensor, + c_latent: cute.Tensor, + c_rope: cute.Tensor, + page_table: cute.Tensor, + o: cute.Tensor, + lse: cute.Tensor, + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + stream: cuda.CUstream, + ): + """Execute the Multi-Head Latent Attention operation on the provided tensors. + + The method handles: + 1. Initialization of workspace for temporary split KV buffers + 2. Validation of tensor data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch(split KV kernel and reduction kernel) with appropriate parameters + + :param q_latent: The query tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type q_latent: cute.Tensor + :param q_rope: The query RoPE tensor with shape [num_head, rope_dim, seq_len_q, batch_size] + :type q_rope: cute.Tensor + :param c_latent: The key tensor with shape [seq_len_k, latent_dim, batch_size] + :type c_latent: cute.Tensor + :param c_rope: The key RoPE tensor with shape [seq_len_k, rope_dim, batch_size] + :type c_rope: cute.Tensor + :param page_table: The page table tensor with shape [page_count, batch_size] + :type page_table: cute.Tensor + :param o: The output tensor with shape [num_head, latent_dim, seq_len_q, batch_size] + :type o: cute.Tensor + :param lse: The LSE tensor with shape [num_head, seq_len_q, batch_size] + :type lse: cute.Tensor + :param workspace: The workspace tensor with 1-d shape prepared for acc_o and acc_lse + :type workspace: cute.Tensor + :param split_kv: The scalar factor for split KV + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax + :type softmax_scale: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param stream: The CUDA stream to execute the kernel on + :type stream: cuda.CUstream + + :raises TypeError: If tensor data types don't match or aren't supported + """ + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q_latent.element_type + self.k_dtype = c_latent.element_type + self.v_dtype = c_latent.element_type + self.o_dtype = o.element_type + + # check type consistency + if cutlass.const_expr(self.q_dtype != self.k_dtype + or self.q_dtype != self.v_dtype): + raise TypeError( + f"Type mismatch: {self.q_dtype} != {self.k_dtype} or {self.q_dtype} != {self.v_dtype}" + ) + # check leading dimensions of input/output + 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") + 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") + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") + + acc_o, acc_lse = self.initialize_workspace( + q_latent.shape[0], + q_latent.shape[1], + q_latent.shape[2], + q_latent.shape[3], + split_kv, + self.acc_dtype, + workspace, + ) + + c_latent_tranpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor(c_latent.iterator, + c_latent_tranpose_layout) + + self.q_major_mode = OperandMajorMode.K + self.k_major_mode = OperandMajorMode.K + self.v_major_mode = OperandMajorMode.MN + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from smem & k-major + p_major_mode = OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.acc_dtype, + cta_group, + self.mma_qk_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.acc_dtype, + cta_group, + self.mma_pv_tiler[:2], + ) + + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape, ), + ) + + self.epi_tile = self.mma_pv_tiler[:2] + + q_latent_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_tiler, + self.q_dtype, + (self.iterations_qk_latent * self.load_q_stage), + ) + q_latent_smem_layout_staged = cute.logical_divide( + q_latent_smem_layout_staged, + (None, None, None, self.iterations_qk_latent)) + q_rope_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.q_dtype, + self.load_q_stage, + ) + + kc_latent_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_tiler, + self.k_dtype, + (self.iterations_qk_latent * self.load_k_stage), + ) + kc_page_tile_size = min( + self.page_size, + qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape) + kc_latent_smem_layout_staged = cute.logical_divide( + kc_latent_smem_layout_staged, + (None, None, None, self.iterations_qk_latent)) + + kc_latent_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + (self.mma_qk_tiler[0] // qk_tiled_mma.thr_id.shape, + self.mma_qk_tiler[2]), + self.k_dtype, + (self.iterations_qk_latent * self.load_k_stage), + ) + kc_latent_smem_layout_for_tma = cute.tiled_divide( + kc_latent_smem_layout_for_tma, + (kc_page_tile_size, self.mma_qk_tiler[2])) + kc_latent_smem_layout_for_tma = cute.logical_divide( + kc_latent_smem_layout_for_tma, + (None, None, None, self.iterations_qk_latent)) + + kc_rope_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.mma_qk_rope_tiler, + self.k_dtype, + self.load_k_stage, + ) + kc_rope_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.K, + ( + self.mma_qk_rope_tiler[0] // qk_tiled_mma.thr_id.shape, + self.mma_qk_rope_tiler[2], + ), + self.k_dtype, + (self.iterations_qk_rope * self.load_k_stage), + ) + kc_rope_smem_layout_for_tma = cute.tiled_divide( + kc_rope_smem_layout_for_tma, + (kc_page_tile_size, self.mma_qk_rope_tiler[2])) + + p_smem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.mma_pv_tiler, + self.q_dtype, + (self.iterations_pv_k * self.p_mma_stage), + ) + p_smem_layout_staged = cute.logical_divide( + p_smem_layout_staged, (None, None, None, self.iterations_pv_k)) + + vc_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.mma_pv_tiler, + self.v_dtype, + (self.iterations_pv_k * self.iterations_pv_n * self.load_v_stage), + ) + vc_smem_layout_staged = cute.logical_divide( + cute.logical_divide( + vc_smem_layout_staged, + (None, None, None, self.iterations_pv_k * self.iterations_pv_n), + ), + (None, None, None, (self.iterations_pv_n, None)), + ) + vc_page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + vc_smem_layout_for_tma = sm100_utils.make_smem_layout( + OperandMajorMode.MN, + (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, + self.mma_pv_tiler[2]), + self.v_dtype, + (self.iterations_pv_k * self.iterations_pv_n * self.load_v_stage), + ) + vc_smem_layout_for_tma = cute.tiled_divide( + vc_smem_layout_for_tma, + ( + pv_tiled_mma.op.shape_mnk[1] // pv_tiled_mma.thr_id.shape, + vc_page_tile_size, + ), + ) + vc_smem_layout_for_tma = cute.logical_divide( + cute.logical_divide( + vc_smem_layout_for_tma, + (None, None, None, self.iterations_pv_k * self.iterations_pv_n), + ), + (None, None, None, (self.iterations_pv_n, None)), + ) + # TMA load for Q latent and rope + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_smem_layout = cute.select(q_latent_smem_layout_staged, mode=[0, 1, 2]) + + tma_atom_q_latent, tma_tensor_q_latent = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_latent, + q_smem_layout, + self.mma_qk_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + q_rope_smem_layout = cute.select(q_rope_smem_layout_staged, + mode=[0, 1, 2]) + tma_atom_q_rope, tma_tensor_q_rope = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q_rope, + q_rope_smem_layout, + self.mma_qk_rope_tiler, + qk_tiled_mma, + cta_layout_vmnk.shape, + ) + # TMA load for c latent and k rope + kc_smem_layout = cute.select(kc_latent_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent, tma_tensor_c_latent = self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent, + kc_smem_layout, + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + kc_rope_smem_layout = cute.select(kc_rope_smem_layout_for_tma, mode=[0]) + tma_atom_c_rope, tma_tensor_c_rope = self.make_paged_tiled_tma_atom( + tma_load_op, + c_rope, + kc_rope_smem_layout, + (self.mma_qk_rope_tiler[1], self.mma_qk_rope_tiler[2]), + qk_tiled_mma, + is_k_load=True, + ) + + # TMA load for c latent transpose + vc_smem_layout = cute.select(vc_smem_layout_for_tma, mode=[0]) + tma_atom_c_latent_transpose, tma_tensor_c_latent_transpose = ( + self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent_transpose, + vc_smem_layout, + (self.mma_pv_tiler[1], self.mma_pv_tiler[2]), + pv_tiled_mma, + is_k_load=False, + )) + + q_latent_copy_size = (cute.size_in_bytes(self.q_dtype, q_smem_layout) * + cute.size(qk_tiled_mma.thr_id.shape) * + self.iterations_qk_latent) + q_rope_copy_size = ( + cute.size_in_bytes(self.q_dtype, q_rope_smem_layout) * + cute.size(qk_tiled_mma.thr_id.shape) * self.iterations_qk_rope) + kc_latent_copy_size = (cute.size_in_bytes( + self.k_dtype, + cute.select(kc_latent_smem_layout_staged, mode=[0, 1, 2]), + ) * cute.size(qk_tiled_mma.thr_id.shape) * self.iterations_qk_latent) + kc_rope_copy_size = (cute.size_in_bytes( + self.k_dtype, + cute.select(kc_rope_smem_layout_staged, mode=[0, 1, 2]), + ) * cute.size(qk_tiled_mma.thr_id.shape) * self.iterations_qk_rope) + vc_copy_size = (cute.size_in_bytes( + self.v_dtype, cute.select(vc_smem_layout_staged, mode=[0, 1, 2])) * + cute.size(pv_tiled_mma.thr_id.shape) * + self.iterations_pv_n * self.iterations_pv_k) + + self.tma_copy_q_bytes = q_latent_copy_size + q_rope_copy_size + self.tma_copy_kc_bytes = kc_latent_copy_size + kc_rope_copy_size + self.tma_copy_vc_bytes = vc_copy_size + + tile_sched_params, grid = self._compute_grid( + o, + split_kv, + self.cluster_shape_mnk, + self.max_active_clusters, + self.is_persistent, + ) + + @cute.struct + class SplitKVKernelSharedStorage: + # Pipeline barriers + load_q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.load_q_stage * 2] + load_k_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.load_k_stage * 2] + load_v_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.load_v_stage * 2] + mma_s_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.mma_s_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.p_mma_stage * 2] + p_cor_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.p_cor_stage * 2] + mma_o_mbar_ptr: cute.struct.MemRange[cutlass.Int64, + self.mma_o_stage * 2] + + # Smem tensors + smem_p: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, + cute.cosize(p_smem_layout_staged)], + 1024, + ] + smem_kc_latent: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, + cute.cosize(kc_latent_smem_layout_staged)], + 1024, + ] + + smem_kc_rope: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, + cute.cosize(kc_rope_smem_layout_staged)], + 1024, + ] + smem_q_latent: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, + cute.cosize(q_latent_smem_layout_staged)], + 1024, + ] + smem_q_rope: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, + cute.cosize(q_rope_smem_layout_staged)], + 1024, + ] + smem_vc: cute.struct.Align[ + cute.struct.MemRange[self.v_dtype, + cute.cosize(vc_smem_layout_staged)], + 1024, + ] + softmax_smem_exchange: cute.struct.MemRange[self.acc_dtype, + self.num_compute_warps * + self.threads_per_warp] + epilogue_smem_exchange: cute.struct.MemRange[ + self.acc_dtype, self.num_compute_warps * self.threads_per_warp] + + # Tmem dealloc cluster barrier + tmem_dealloc_mbar_ptr: cutlass.Int64 + + # Tmem holding buffer + tmem_holding_buf: cutlass.Int32 + + softmax_scale_log2 = softmax_scale * LOG2_E + + self.split_kv_kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q_latent, + tma_tensor_q_latent, + tma_atom_q_rope, + tma_tensor_q_rope, + tma_atom_c_latent, + tma_tensor_c_latent, + tma_atom_c_rope, + tma_tensor_c_rope, + tma_atom_c_latent_transpose, + tma_tensor_c_latent_transpose, + page_table, + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale_log2, + output_scale, + q_latent_smem_layout_staged, + q_rope_smem_layout_staged, + kc_latent_smem_layout_staged, + kc_rope_smem_layout_staged, + p_smem_layout_staged, + vc_smem_layout_staged, + kc_latent_smem_layout_for_tma, + kc_rope_smem_layout_for_tma, + vc_smem_layout_for_tma, + cta_layout_vmnk, + tile_sched_params, + SplitKVKernelSharedStorage, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=SplitKVKernelSharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( + grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], + smem=MAX_SPLITS * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.jit + def make_paged_tiled_tma_atom( + self, + tma_load_op: cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp, + gmem: cute.Tensor, + smem_layout: cute.Layout, + mma_tiler, + tiled_mma: cute.TiledMma, + is_k_load: bool, + ): + ident = cute.make_identity_layout(gmem.shape) + g_tile = cute.composition(ident, mma_tiler) + cta_mn = mma_tiler[0] // tiled_mma.thr_id.shape + cta_v_map = cute.flat_divide(g_tile, (cta_mn, )) + cta_v_map = cute.select(cta_v_map, mode=[0, 2]) + page_tile_size = (min(self.page_size, cta_mn) if is_k_load else min( + self.page_size, mma_tiler[1])) + cta_v_map = cute.zipped_divide( + cta_v_map, + (page_tile_size, mma_tiler[1]) if is_k_load else + (cta_mn, page_tile_size), + ) + cta_v_map = cute.select(cta_v_map, mode=[0]) + from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir + + res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( + gmem.value, + smem_layout.value, + cta_v_map, + tma_load_op._to_ir(), + num_multicast=1, + ) + return ( + cute.CopyAtom(tma_load_op, + cpasync.CopyBulkTensorTileG2SNonExecTrait(res[0])), + res[1], + ) + + @cute.kernel + def split_kv_kernel( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tma_atom_q_latent: Optional[cute.CopyAtom], + mQL: cute.Tensor, + tma_atom_q_rope: Optional[cute.CopyAtom], + mQR: cute.Tensor, + tma_atom_c_latent: Optional[cute.CopyAtom], + mCL: cute.Tensor, + tma_atom_c_rope: Optional[cute.CopyAtom], + mKR: cute.Tensor, + tma_atom_c_latent_transpose: Optional[cute.CopyAtom], + mCLT: cute.Tensor, + mPT: cute.Tensor, + mO: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mAccO: Optional[cute.Tensor], + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, + q_latent_smem_layout_staged: cute.ComposedLayout, + q_rope_smem_layout_staged: cute.ComposedLayout, + kc_latent_smem_layout_staged: cute.ComposedLayout, + kc_rope_smem_layout_staged: cute.ComposedLayout, + p_smem_layout_staged: cute.ComposedLayout, + vc_smem_layout_staged: cute.ComposedLayout, + kc_latent_smem_layout_for_tma: Optional[cute.ComposedLayout], + kc_rope_smem_layout_for_tma: Optional[cute.ComposedLayout], + vc_smem_layout_for_tma: Optional[cute.ComposedLayout], + cta_layout_vmnk: cute.Layout, + tile_sched_params: MLAStaticTileSchedulerParams, + SharedStorage: cutlass.Constexpr, + ): + """The device split_kv kernel implementation of the Multi-Head Latent Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the MLA computation: + 1. Load warp: Loads Q/C latent/rope data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Compute warps: Compute softmax and do rescaling on accumulators, and store the intermediate/final results + to global memory + + The kernel produces either intermediate or final results of the MLA computation based on the split_kv parameter. + When split_kv is 1, the kernel generates the final results directly. Otherwise, it produces intermediate results + that will later be combined by a reduction kernel. + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases. + + :param tiled_mma_qk: Tiled MMA for Q*K^T + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: Tiled MMA for P*V + :type tiled_mma_pv: cute.TiledMma + :param tma_atom_q_latent: TMA copy atom for query latent tensor + :type tma_atom_q_latent: cute.CopyAtom + :param mQL: query latent tensor + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom + :param mKR: Compressed rope tensor + :type mKR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor + :type mCL: cute.Tensor + :param tma_atom_c_rope: TMA copy atom for c rope tensor + :type tma_atom_c_rope: cute.CopyAtom + :param mCLT: Compressed latent transpose tensor + :type mCLT: cute.Tensor + :param mPT: Page table tensor + :type mPT: cute.Tensor + :param mO: Output tensor + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor + :type mLSE: cute.Tensor + :param mAccO: Intermediate accumulator output tensor + :type mAccO: cute.Tensor + :param mAccLSE: Intermediate accumulator log-sum-exp tensor + :type mAccLSE: cute.Tensor + :param split_kv: The split_kv parameter + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax + :type softmax_scale_log2: cutlass.Float32 + :param output_scale: The scale factor for the output + :type output_scale: cutlass.Float32 + :param q_latent_smem_layout_staged: Shared memory layout for query tensor + :type q_latent_smem_layout_staged: cute.ComposedLayout + :param q_rope_smem_layout_staged: Shared memory layout for query rope tensor + :type q_rope_smem_layout_staged: cute.ComposedLayout + :param kc_latent_smem_layout_staged: Shared memory layout for key tensor + :type kc_latent_smem_layout_staged: cute.ComposedLayout + :param kc_rope_smem_layout_staged: Shared memory layout for key rope tensor + :type kc_rope_smem_layout_staged: cute.ComposedLayout + :param p_smem_layout_staged: Shared memory layout for probability matrix + :type p_smem_layout_staged: cute.ComposedLayout + :param vc_smem_layout_staged: Shared memory layout for value tensor + :type vc_smem_layout_staged: cute.ComposedLayout + :param cta_layout_vmnk: Layout for compute threads + :type cta_layout_vmnk: cute.Layout + :param tile_sched_params: Scheduling parameters for work distribution + :type tile_sched_params: MLAStaticTileSchedulerParams + :param SharedStorage: Shared storage for the kernel + :type SharedStorage: cutlass.Constexpr + """ + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # Prefetch tma descriptor + if warp_idx == self.mma_warp_id: + cpasync.prefetch_descriptor(tma_atom_q_latent) + cpasync.prefetch_descriptor(tma_atom_q_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent) + cpasync.prefetch_descriptor(tma_atom_c_rope) + cpasync.prefetch_descriptor(tma_atom_c_latent_transpose) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_ptr_sync_bar, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + load_q_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_q_stage, + self.tma_copy_q_bytes, + ) + load_k_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_k_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_k_stage, + self.tma_copy_kc_bytes, + ) + load_v_pipeline = self.make_and_init_load_qkv_pipeline( + storage.load_v_mbar_ptr.data_ptr(), + cta_layout_vmnk, + self.load_v_stage, + self.tma_copy_vc_bytes, + ) + mma_s_pipeline = self.make_and_init_mma_s_pipeline( + storage.mma_s_mbar_ptr.data_ptr(), cta_layout_vmnk) + p_mma_pipeline = self.make_and_init_p_mma_pipeline( + storage.p_mma_mbar_ptr.data_ptr(), cta_layout_vmnk) + p_cor_pipeline = self.make_and_init_p_cor_pipeline( + storage.p_cor_mbar_ptr.data_ptr()) + mma_o_pipeline = self.make_and_init_mma_o_pipeline( + storage.mma_o_mbar_ptr.data_ptr(), cta_layout_vmnk) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, + is_relaxed=True) + + # Generate smem tensor Q/KC/VC/exchange + # (MMA, MMA_H, MMA_R, PIPE) + sQ = storage.smem_q_latent.get_tensor( + q_latent_smem_layout_staged.outer, + swizzle=q_latent_smem_layout_staged.inner) + sQ_rope = storage.smem_q_rope.get_tensor( + q_rope_smem_layout_staged.outer, + swizzle=q_rope_smem_layout_staged.inner) + # (MMA, MMA_K, MMA_R, PIPE) + sKC = storage.smem_kc_latent.get_tensor( + kc_latent_smem_layout_staged.outer, + swizzle=kc_latent_smem_layout_staged.inner, + ) + sKC_rope = storage.smem_kc_rope.get_tensor( + kc_rope_smem_layout_staged.outer, + swizzle=kc_rope_smem_layout_staged.inner) + sKC_for_tma = storage.smem_kc_latent.get_tensor( + kc_latent_smem_layout_for_tma.outer, + swizzle=kc_latent_smem_layout_for_tma.inner, + ) + sKC_rope_for_tma = storage.smem_kc_rope.get_tensor( + kc_rope_smem_layout_for_tma.outer, + swizzle=kc_rope_smem_layout_for_tma.inner) + # (MMA, MMA_D, MMA_K, PIPE) + sVC = storage.smem_vc.get_tensor(vc_smem_layout_staged.outer, + swizzle=vc_smem_layout_staged.inner) + sVC_for_tma = storage.smem_vc.get_tensor( + vc_smem_layout_for_tma.outer, swizzle=vc_smem_layout_for_tma.inner) + # (MMA, MMA_H, MMA_K) + sP = storage.smem_p.get_tensor(p_smem_layout_staged.outer, + swizzle=p_smem_layout_staged.inner) + # (compute_threads,) + softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp)) + epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( + cute.make_layout(self.num_compute_warps * self.threads_per_warp)) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk) + + # /////////////////////////////////////////////////////////////////////////////// + # Load warps, including page table and data tensors + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.empty_warp_ids[ + 0] and warp_idx <= self.empty_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.other_reg_num) + + if warp_idx == self.load_tma_k_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_q_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_q_stage) + load_k_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_k_stage) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_k_pipeline=load_k_pipeline, + load_v_pipeline=load_v_pipeline, + mPT=mPT, + ) + tma_qk_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + tma_atom_q_latent=tma_atom_q_latent, + tma_atom_q_rope=tma_atom_q_rope, + tma_atom_c_latent=tma_atom_c_latent, + tma_atom_c_rope=tma_atom_c_rope, + mQL=mQL, + mQR=mQR, + mCL=mCL, + mKR=mKR, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC_for_tma, + sKC_rope=sKC_rope_for_tma, + ) + # Load tma + load_q_producer_state, load_k_producer_state = self.load_tma_qk( + tma_common_params, + tma_qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_k_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + load_q_pipeline.producer_tail(load_q_producer_state) + load_k_pipeline.producer_tail(load_k_producer_state) + + if warp_idx == self.load_tma_v_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + load_v_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.load_v_stage) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, + cache_seqs, + block_split_kvs, + blk_coord, + ) + if k_tile_count > 0: + # Construct fixed common/tma_qk/tma_pv params for load_tma + tma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_v_pipeline=load_v_pipeline, + mPT=mPT, + ) + tma_pv_params = SimpleNamespace( + tiled_mma_pv=tiled_mma_pv, + tma_atom_c_latent_transpose=tma_atom_c_latent_transpose, + mCLT=mCLT, + sVC=sVC_for_tma, + ) + # Load tma + load_v_producer_state = self.load_tma_v( + tma_common_params, + tma_pv_params, + k_index, + k_tile_count, + load_v_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + load_v_pipeline.producer_tail(load_v_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA warp + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.other_reg_num) + # Alloc tensor memory buffer + tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100")) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + load_q_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_q_stage) + load_k_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_k_stage) + load_v_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.load_v_stage) + mma_s_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_s_stage) + p_mma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_mma_stage) + mma_o_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.mma_o_stage) + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord) + if k_tile_count > 0: + mma_common_params = SimpleNamespace( + blk_coord=blk_coord, + local_split_kv=local_split_kv, + load_q_pipeline=load_q_pipeline, + load_k_pipeline=load_k_pipeline, + load_v_pipeline=load_v_pipeline, + tmem_ptr=tmem_ptr, + is_leader_cta=is_leader_cta, + L=mCL.shape[1], + ) + mma_qk_params = SimpleNamespace( + mma_s_pipeline=mma_s_pipeline, + sQ=sQ, + sQ_rope=sQ_rope, + sKC=sKC, + sKC_rope=sKC_rope, + ) + mma_pv_params = SimpleNamespace( + p_mma_pipeline=p_mma_pipeline, + mma_o_pipeline=mma_o_pipeline, + sP=sP, + sVC=sVC, + ) + ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma( + mma_common_params, + mma_qk_params, + mma_pv_params, + k_tile_count, + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + mma_s_pipeline.producer_tail(mma_s_producer_state) + mma_o_pipeline.producer_tail(mma_o_producer_state) + + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Compute warp + # /////////////////////////////////////////////////////////////////////////////// + if (warp_idx >= self.compute_warp_ids[0] + and warp_idx <= self.compute_warp_ids[-1]): + cute.arch.setmaxregister_increase(self.softmax_reg_num) + mma_s_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_s_stage) + p_mma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_mma_stage) + p_cor_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.p_cor_stage) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=softmax_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, + p_cor_pipeline=p_cor_pipeline, + ) + compute_softmax_params = SimpleNamespace( + tiled_mma_qk=tiled_mma_qk, + sP=sP, + mma_s_pipeline=mma_s_pipeline, + p_mma_pipeline=p_mma_pipeline, + softmax_scale_log2=softmax_scale_log2, + ) + mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state = ( + self.compute( + compute_common_params, + compute_softmax_params, + k_index=k_index, + k_tile_count=k_tile_count, + mma_s_consumer_state=mma_s_consumer_state, + p_mma_producer_state=p_mma_producer_state, + p_cor_producer_state=p_cor_producer_state, + )) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + p_cor_pipeline.producer_tail(p_cor_producer_state) + + # /////////////////////////////////////////////////////////////////////////////// + # Correction warp + # /////////////////////////////////////////////////////////////////////////////// + if (warp_idx >= self.correction_warp_ids[0] + and warp_idx <= self.correction_warp_ids[-1]): + cute.arch.setmaxregister_increase(self.correction_reg_num) + p_cor_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.p_cor_stage) + mma_o_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.mma_o_stage) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tile_sched = create_mla_static_tile_scheduler( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord) + if k_tile_count > 0: + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, + smem_exchange=epilogue_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], + L=mCL.shape[1], + H=mQL.shape[0], + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, + p_cor_pipeline=p_cor_pipeline, + mma_o_pipeline=mma_o_pipeline, + ) + compute_epilogue_params = SimpleNamespace( + output_scale=output_scale, + softmax_scale_log2=softmax_scale_log2, + mAccLSE=mAccLSE, + mLSE=mLSE, + ) + p_cor_consumer_state, mma_o_consumer_state = self.correction( + compute_common_params, + compute_epilogue_params, + k_tile_count=k_tile_count, + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + return + + @cute.kernel + def reduction_kernel( + self, + mO: cute.Tensor, + mLSE: cute.Tensor, + mAccO: cute.Tensor, + mAccLSE: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + ): + """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results + from multiple split_kv blocks into final outputs. + + :param mO: Output tensor for storing final results + :type mO: cute.Tensor + :param mLSE: Log-sum-exp tensor for storing final LSE values + :type mLSE: cute.Tensor + :param mAccO: Accumulated output tensor from split_kv blocks + :type mAccO: cute.Tensor + :param mAccLSE: Accumulated LSE tensor from split_kv blocks + :type mAccLSE: cute.Tensor + :param split_kv: Number of split_kv blocks + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor (for variable split_kv) + :type block_split_kvs: cute.Tensor + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + blk_coord = (bidx, bidy, bidz) + local_split_kv = (block_split_kvs[blk_coord[2]] + if self.is_var_split_kv else split_kv) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], + self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) + local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() + storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) + smem_lse_scale = cute.make_tensor(lse_scale_ptr, + cute.make_layout(MAX_SPLITS)) + + gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) + lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) + + local_lse = cute.make_rmem_tensor(cute.make_layout(lse_per_thread), + self.lse_dtype) + lse_max = -self.lse_dtype.inf + # find the max lse + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + local_lse[i] = (gLSE[split_kv_idx] if cute.elem_less( + split_kv_idx, local_split_kv) else -self.lse_dtype.inf) + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) + lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 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 = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse + global_lse = (lse_max + cute.math.log2(sum_lse, fastmath=True) + if not sum_lse == self.lse_dtype(0.0) + or sum_lse != sum_lse else self.lse_dtype.inf) + if tidx == 0: + mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[split_kv_idx] = cute.math.exp2(local_lse[i] - + global_lse, + fastmath=True) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( + self.latent_dim, self.threads_per_warp * self.num_compute_warps) + gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] + rAccO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), + self.acc_dtype) + rO = cute.make_rmem_tensor(cute.make_layout(elements_per_thread), + self.o_dtype) + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): + element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + return + + @staticmethod + def get_split_kv(B: int, S: int, K: int, mma_qk_tiler_mn: tuple, + max_active_blocks: int) -> int: + """Get the proper split_kv value for the MLA kernel based on parameters. + + :param B: Batch size + :type B: int + :param S: Sequence length + :type S: int + :param K: Sequence length + :type K: int + :param mma_qk_tiler_mn: MLA tiling parameters + :type mma_qk_tiler_mn: tuple + :param max_active_blocks: Maximum number of active blocks + :type max_active_blocks: int + :return: Split_kv value + :rtype: int + """ + max_splits = ceil_div(K, mma_qk_tiler_mn[1]) + blocks_per_batch = max(1, max_active_blocks // B // (S * 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) + max_split_kv = 32 + return min(split_wave_aware, max_split_kv) + + @cute.jit + def get_k_tile_count( + self, + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, + block_split_kvs: cute.Tensor, + blk_coord: cute.Coord, + ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Get the current k_index, k_tile_count, and local split_kv value for the MLA kernel. + + :param split_kv: Split_kv value + :type split_kv: cutlass.Int32 + :param cache_seqs: Cache sequence lengths tensor + :type cache_seqs: cute.Tensor + :param block_split_kvs: Per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param blk_coord: Block coordinate + :type blk_coord: cute.Coord + :return: k_index, k_tile_count, split_kv + :rtype: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32] + """ + K = cache_seqs[blk_coord[2]] + if cutlass.const_expr(self.is_var_split_kv): + split_kv = block_split_kvs[blk_coord[2]] + + k_tile_total = cute.ceil_div(K, self.mma_qk_tiler[1]) + k_tile_per_cta = cute.ceil_div(k_tile_total, split_kv) + k_index = blk_coord[3] * k_tile_per_cta + k_tile_count = max( + 0, + min(k_tile_total, k_index + k_tile_per_cta) - k_index) + return k_index, k_tile_count, split_kv + + @cute.jit + def load_tma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState | None = None, + load_k_producer_state: pipeline.PipelineState | None = None, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load wrap to load Q/K tensors. Updates the load qk producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_k_producer_state: The load k producer state + :type load_k_producer_state: pipeline.PipelineState + + :return: The load q producer state and load k producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for QK TMA load + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) + mma_qk_tiler_mk_rope = cute.select(self.mma_qk_rope_tiler, mode=[0, 2]) + gQR = cute.flat_divide(qk_params.mQR, mma_qk_tiler_mk_rope) + + thr_mma_qk = qk_params.tiled_mma_qk.get_slice( + common_params.blk_coord[0] % + cute.size(qk_params.tiled_mma_qk.thr_id)) + tSgQL = thr_mma_qk.partition_A(gQL) + tSgQR = thr_mma_qk.partition_A(gQR) + + cta_m = min( + qk_params.tiled_mma_qk.op.shape_mnk[0] // + qk_params.tiled_mma_qk.thr_id.shape, + self.page_size, + ) + page_tile_size = min(self.page_size, cta_m) + gCL = cute.tiled_divide(qk_params.mCL, + (page_tile_size, self.mma_qk_tiler[2])) + tSgCL = (gCL[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] if cta_m < self.page_size else gCL[None, 0, None, None]) + gKR = cute.tiled_divide(qk_params.mKR, + (page_tile_size, self.mma_qk_rope_tiler[2])) + tSgKR = (gKR[ + None, + common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, + None, + None, + ] if cta_m < self.page_size else gKR[None, 0, None, None]) + # tma partition for q, k latent/rope + + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tQsQ, tQLgQL_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_latent, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ, 0, 3), + cute.group_modes(tSgQL, 0, 3), + ) + + tQsQ_rope, tQRgQR_mkl = cpasync.tma_partition( + qk_params.tma_atom_q_rope, + 0, + cute.make_layout(1), + cute.group_modes(qk_params.sQ_rope, 0, 3), + cute.group_modes(tSgQR, 0, 3), + ) + tKCsKC, tCLgCL = cpasync.tma_partition( + qk_params.tma_atom_c_latent, + 0, + cute.make_layout(1), + qk_params.sKC, + tSgCL, + ) + + tKCsKC_rope, tKRgKR = cpasync.tma_partition( + qk_params.tma_atom_c_rope, + 0, + cute.make_layout(1), + qk_params.sKC_rope, + tSgKR, + ) + + tQLgQL = tQLgQL_mkl[None, None, None, common_params.blk_coord[1], + common_params.blk_coord[2]] + tQRgQR = tQRgQR_mkl[None, None, None, common_params.blk_coord[1], + common_params.blk_coord[2]] + + # set extra params + common_params.mPT = mPT + qk_params.tQLgQL = tQLgQL + qk_params.tQRgQR = tQRgQR + qk_params.tCLgCL = tCLgCL + qk_params.tKRgKR = tKRgKR + qk_params.tQsQ = tQsQ + qk_params.tQsQ_rope = tQsQ_rope + qk_params.tKCsKC = tKCsKC + qk_params.tKCsKC_rope = tKCsKC_rope + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + load_q_producer_state, load_k_producer_state = self.load_tma_qk_one_k_tile( + common_params, + qk_params, + k_index, + k_tile_count, + load_q_producer_state, + load_k_producer_state, + load_q=k_tile_count_init == k_tile_count, + ) + k_index += 1 + k_tile_count -= 1 + + return load_q_producer_state, load_k_producer_state + + @cute.jit + def load_tma_v( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_v_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load wrap to load V tensors. Updates the load v producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_v_producer_state: The load v producer state + :type load_v_producer_state: pipeline.PipelineState + + :return: The load v producer state + :rtype: pipeline.PipelineState + """ + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + + # Flatten divide and partition global tensors for V TMA load + page_tile_size = min(self.page_size, self.mma_pv_tiler[2]) + gCLT = cute.flat_divide(v_params.mCLT, + (self.mma_pv_tiler[1], page_tile_size)) + cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape + gCLT = cute.logical_divide(gCLT, + (cta_n, ))[(None, + common_params.blk_coord[0]), + None, None, None, None] + tOgCLT = cute.tiled_divide(gCLT, (cta_n, page_tile_size)) + tOgCLT = tOgCLT[None, 0, 0, None, None, None] + # tma partition for vc + # smem: ((atom_v, rest_v), STAGE) + # gmem: ((atom_v, rest_v), RestM, RestK, RestL) + tVCsVC, tCLTgCLT = cpasync.tma_partition( + v_params.tma_atom_c_latent_transpose, + 0, + cute.make_layout(1), + v_params.sVC, + tOgCLT, + ) + + # set extra params + common_params.mPT = mPT + v_params.tCLTgCLT = tCLTgCLT + v_params.tVCsVC = tVCsVC + + while k_tile_count > 0: + load_v_producer_state = self.load_tma_v_one_k_tile( + common_params, + v_params, + k_index, + load_v_producer_state, + ) + k_index += 1 + k_tile_count -= 1 + return load_v_producer_state + + @cute.jit + def load_tma_qk_one_k_tile( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + load_q_producer_state: pipeline.PipelineState, + load_k_producer_state: pipeline.PipelineState, + load_q: bool, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Load one k-tile of Q/C latent/rope tensors. Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param load_q_producer_state: The load q producer state + :type load_q_producer_state: pipeline.PipelineState + :param load_k_producer_state: The load kv producer state + :type load_k_producer_state: pipeline.PipelineState + :param load_q: Whether to load q + :type load_q: bool + + :return: The load q producer state and load kv producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + page_per_tile = ceil_div(self.mma_qk_tiler[1] // self.page_size, + qk_params.tiled_mma_qk.thr_id.shape) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), + cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = (common_params.mPT[k_index] if self.mma_qk_tiler[1] // + self.page_size == 1 else common_params.mPT[ + (k_index * qk_params.tiled_mma_qk.thr_id.shape + + common_params.blk_coord[0]) * page_per_tile + i]) + # load q once at first iteration + load_q_pipeline = common_params.load_q_pipeline + if load_q: + # get the mbar ptr from pipeline. + tma_bar_ptr = load_q_pipeline.producer_get_barrier( + load_q_producer_state) + # expect the extra bytes for q. + load_q_pipeline.producer_acquire(load_q_producer_state) + for i in cutlass.range_constexpr(self.iterations_qk_latent): + # load q latent + cute.copy( + qk_params.tma_atom_q_latent, + qk_params.tQLgQL[None, 0, i], + qk_params.tQsQ[None, (i, 0)], + tma_bar_ptr=tma_bar_ptr, + ) + for i in cutlass.range_constexpr(self.iterations_qk_rope): + # load q rope + cute.copy( + qk_params.tma_atom_q_rope, + qk_params.tQRgQR[None, 0, i], + qk_params.tQsQ_rope[None, i], + tma_bar_ptr=tma_bar_ptr, + ) + load_q_producer_state.advance() + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_k_pipeline.producer_get_barrier( + load_k_producer_state) + common_params.load_k_pipeline.producer_acquire(load_k_producer_state) + for i in range(self.iterations_qk_latent): + for k in range(page_per_tile): + # load k latent + cute.copy( + qk_params.tma_atom_c_latent, + qk_params.tCLgCL[None, i, k_idx[k]], + qk_params.tKCsKC[None, k, 0, + (i, load_k_producer_state.index)], + tma_bar_ptr=tma_bar_ptr, + ) + + for i in cutlass.range_constexpr(self.iterations_qk_rope): + for k in cutlass.range_constexpr(page_per_tile): + # load k rope + cute.copy( + qk_params.tma_atom_c_rope, + qk_params.tKRgKR[None, i, k_idx[k]], + qk_params.tKCsKC_rope[None, k, 0, + load_k_producer_state.index], + tma_bar_ptr=tma_bar_ptr, + ) + load_k_producer_state.advance() + + return load_q_producer_state, load_k_producer_state + + @cute.jit + def load_tma_v_one_k_tile( + self, + common_params: SimpleNamespace, + v_params: SimpleNamespace, + k_index: cutlass.Int32, + load_v_producer_state: pipeline.PipelineState, + ) -> pipeline.PipelineState: + """Load one k-tile of compressed latent transpose tensor(v). Updates the load qkv producer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param v_params: The load tma v parameters + :type v_params: SimpleNamespace + :param k_index: The k index + :type k_index: cutlass.Int32 + :param load_v_producer_state: The load v producer state + :type load_v_producer_state: pipeline.PipelineState + + :return: The load qkv producer state + :rtype: pipeline.PipelineState + """ + page_per_tile = self.mma_pv_tiler[ + 2] * self.iterations_pv_k // self.page_size + page_per_subtile = ceil_div(page_per_tile, self.iterations_pv_k) + k_idx = cute.make_rmem_tensor(cute.make_layout(page_per_tile), + cutlass.Int32) + for i in cutlass.range_constexpr(page_per_tile): + k_idx[i] = (common_params.mPT[k_index] if page_per_tile == 1 else + common_params.mPT[k_index * page_per_tile + i]) + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_v_pipeline.producer_get_barrier( + load_v_producer_state) + common_params.load_v_pipeline.producer_acquire(load_v_producer_state) + for j in cutlass.range_constexpr(self.iterations_pv_n): + for i in cutlass.range_constexpr(self.iterations_pv_k): + if cutlass.const_expr(page_per_tile > 1): + for k in cutlass.range_constexpr(page_per_subtile): + k_idx_i = k_idx[k + i * page_per_subtile] + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[None, j, 0, k_idx_i], + v_params.tVCsVC[None, 0, k, + ((j, i), + load_v_producer_state.index)], + tma_bar_ptr=tma_bar_ptr, + ) + else: + cute.copy( + v_params.tma_atom_c_latent_transpose, + v_params.tCLTgCLT[None, j, i, k_idx[0]], + v_params.tVCsVC[None, 0, 0, + ((j, i), load_v_producer_state.index)], + tma_bar_ptr=tma_bar_ptr, + ) + load_v_producer_state.advance() + return load_v_producer_state + + @cute.jit + def mma( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + pv_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_k_consumer_state: pipeline.PipelineState, + load_v_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """MMA warp to compute the result of Q*K^T and P*V. Updates the tiled mma and pipeline states. + + :param common_params: The common parameters for mma qk and pv + :type common_params: SimpleNamespace + :param qk_params: The mma qk parameters + :type qk_params: SimpleNamespace + :param pv_params: The mma pv parameters + :type pv_params: SimpleNamespace + :param k_tile_count: The k tile count + :type k_tile_count: cutlass.Int32 + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_k_consumer_state: The load k consumer state + :type load_k_consumer_state: pipeline.PipelineState + :param load_v_consumer_state: The load v consumer state + :type load_v_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + :param p_mma_consumer_state: The p mma consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The mma o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the tiled mma pv, the load q consumer state, the load k consumer state, the load v consumer state, the mma s producer state, the p mma consumer state, and the mma o producer state + :rtype: tuple[cute.TiledMma, cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + tSrQ = tiled_mma_qk.make_fragment_A(qk_params.sQ) + tSrQ_rope = tiled_mma_qk.make_fragment_A(qk_params.sQ_rope) + tSrKC = tiled_mma_qk.make_fragment_B(qk_params.sKC) + tSrKC_rope = tiled_mma_qk.make_fragment_B(qk_params.sKC_rope) + tOrP = tiled_mma_pv.make_fragment_A(pv_params.sP) + tOrVC = tiled_mma_pv.make_fragment_B(pv_params.sVC) + + tStS_shape = tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1])) + tStS_staged_fake = tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage)) + # use real tmem ptr for tStS + tStS_staged = cute.make_tensor(common_params.tmem_ptr, + tStS_staged_fake.layout) + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1])) + # mma O has 1 stage. + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO_staged = cute.make_tensor( + tStS_staged.iterator + self.tmem_o_offset, tOtO_layout) + + # set more parameters + qk_params.tSrQ = tSrQ + qk_params.tSrQ_rope = tSrQ_rope + qk_params.tSrKC = tSrKC + qk_params.tSrKC_rope = tSrKC_rope + qk_params.tStS_staged = tStS_staged + pv_params.tOrP = tOrP + pv_params.tOrVC = tOrVC + pv_params.tOtO_staged = tOtO_staged + + # mma O accumulates on K, so the accumulate flag is set to False once before all K blocks. + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + if common_params.is_leader_cta: + load_q_release_state = load_q_consumer_state.clone() + ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + wait_q=True, + ) + k_tile_count -= 1 + + while k_tile_count > 0: + ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) = self.mma_qk( + common_params, + qk_params, + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + wait_q=False, + ) + ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + k_tile_count -= 1 + # release q consumer states + load_q_pipeline.consumer_release(load_q_release_state) + load_q_release_state.advance() + ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) = self.mma_pv( + common_params, + pv_params, + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + return ( + tiled_mma_qk, + tiled_mma_pv, + load_q_consumer_state, + load_k_consumer_state, + load_v_consumer_state, + mma_s_producer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def mma_qk( + self, + common_params: SimpleNamespace, + qk_params: SimpleNamespace, + tiled_mma_qk: cute.TiledMma, + load_q_consumer_state: pipeline.PipelineState, + load_k_consumer_state: pipeline.PipelineState, + mma_s_producer_state: pipeline.PipelineState, + wait_q: bool, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states. + + :param qk_params: The qk parameters + :type qk_params: SimpleNamespace + :param tiled_mma_qk: The tiled mma qk + :type tiled_mma_qk: cute.TiledMma + :param load_q_consumer_state: The load q consumer state + :type load_q_consumer_state: pipeline.PipelineState + :param load_k_consumer_state: The load k consumer state + :type load_k_consumer_state: pipeline.PipelineState + :param mma_s_producer_state: The mma s producer state + :type mma_s_producer_state: pipeline.PipelineState + + :return: The tiled mma qk, the load q consumer state, the load k consumer state, and the mma s producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + tStS = qk_params.tStS_staged[None, None, None, + mma_s_producer_state.index] + + qk_params.mma_s_pipeline.producer_acquire(mma_s_producer_state) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, False) + load_q_pipeline = common_params.load_q_pipeline + load_k_pipeline = common_params.load_k_pipeline + if cutlass.const_expr(wait_q): + load_q_pipeline.consumer_wait(load_q_consumer_state) + load_k_pipeline.consumer_wait(load_k_consumer_state) + for q_stage in range(self.iterations_qk_latent): + kc_stage = load_k_consumer_state.index + for k_block in cutlass.range_constexpr( + cute.size(qk_params.tSrQ.shape[2])): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ[None, None, k_block, (q_stage, 0)], + qk_params.tSrKC[None, None, k_block, (q_stage, kc_stage)], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + + for q_stage in range(self.iterations_qk_rope): + kc_stage = load_k_consumer_state.index + for k_block in cutlass.range_constexpr(self.rope_dim // + tiled_mma_qk.shape_mnk[2]): + cute.gemm( + tiled_mma_qk, + tStS, + qk_params.tSrQ_rope[None, None, k_block, q_stage], + qk_params.tSrKC_rope[None, None, k_block, kc_stage], + tStS, + ) + tiled_mma_qk.set(tcgen05.Field.ACCUMULATE, True) + load_k_pipeline.consumer_release(load_k_consumer_state) + load_k_consumer_state.advance() + if cutlass.const_expr(wait_q): + load_q_consumer_state.advance() + + qk_params.mma_s_pipeline.producer_commit(mma_s_producer_state) + mma_s_producer_state.advance() + return ( + tiled_mma_qk, + load_q_consumer_state, + load_k_consumer_state, + mma_s_producer_state, + ) + + @cute.jit + def mma_pv( + self, + common_params: SimpleNamespace, + pv_params: SimpleNamespace, + tiled_mma_pv: cute.TiledMma, + load_v_consumer_state: pipeline.PipelineState, + p_mma_consumer_state: pipeline.PipelineState, + mma_o_producer_state: pipeline.PipelineState, + ) -> tuple[ + cute.TiledMma, + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + ]: + """Compute one k-tile of mma for P*V. Updates the tiled mma pv and pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param pv_params: The pv parameters + :type pv_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param load_v_consumer_state: The load v consumer state + :type load_v_consumer_state: pipeline.PipelineState + :param p_mma_consumer_state: The P MMA consumer state + :type p_mma_consumer_state: pipeline.PipelineState + :param mma_o_producer_state: The MMA o producer state + :type mma_o_producer_state: pipeline.PipelineState + + :return: The tiled mma pv, the load v consumer state, the P MMA consumer state, and the MMA o producer state + :rtype: tuple[cute.TiledMma, pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + pv_params.p_mma_pipeline.consumer_wait(p_mma_consumer_state) + load_v_pipeline = common_params.load_v_pipeline + accumulate_flag = tiled_mma_pv.get(tcgen05.Field.ACCUMULATE) + mma_o_pipeline = pv_params.mma_o_pipeline + + load_v_pipeline.consumer_wait(load_v_consumer_state) + vc_stage = load_v_consumer_state.index + for acc_stage in range(self.iterations_pv_n): + mma_o_pipeline.producer_acquire(mma_o_producer_state) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, accumulate_flag) + for p_stage in range(self.iterations_pv_k): + tOtO = pv_params.tOtO_staged[None, None, None, acc_stage] + for k_block in cutlass.range_constexpr(pv_params.tOrP.shape[2]): + cute.gemm( + tiled_mma_pv, + tOtO, + pv_params.tOrP[ + None, + None, + k_block, + (p_stage, p_mma_consumer_state.index), + ], + pv_params.tOrVC[None, None, k_block, + ((acc_stage, p_stage), vc_stage)], + tOtO, + ) + tiled_mma_pv.set(tcgen05.Field.ACCUMULATE, True) + + mma_o_pipeline.producer_commit(mma_o_producer_state) + mma_o_producer_state.advance() + load_v_pipeline.consumer_release(load_v_consumer_state) + load_v_consumer_state.advance() + pv_params.p_mma_pipeline.consumer_release(p_mma_consumer_state) + p_mma_consumer_state.advance() + + return ( + tiled_mma_pv, + load_v_consumer_state, + p_mma_consumer_state, + mma_o_producer_state, + ) + + @cute.jit + def compute( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + k_tile_count: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState, + pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + + :return: The MMA s consumer state, the P MMA producer state, and the P correction producer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) + + row_max = -self.acc_dtype.inf + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # no mask applied + while k_tile_count > 1: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + False, + False, + ) + k_index = k_index + 1 + k_tile_count = k_tile_count - 1 + + # mask applied + if cutlass.const_expr(common_params.mAccO is not None): + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + k_index == k_tile_total - 1, + True, + ) + else: + ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + ) = self.softmax( + common_params, + softmax_params, + k_index, + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max, + row_sum, + correction_factor, + True, + True, + ) + + return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state + + @cute.jit + def correction( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + k_tile_count: cutlass.Int32, + p_cor_consumer_state: pipeline.PipelineState, + mma_o_consumer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, pipeline.PipelineState]: + """Compute warp to compute the result of softmax, rescale, and epilogue. Updates the related pipeline states. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param k_tile_count: The number of k-tiles + :type k_tile_count: cutlass.Int32 + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + :param mma_o_consumer_state: The MMA o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, and the MMA o consumer state + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState] + """ + + k_tile_count_init = k_tile_count + while k_tile_count > 0: + p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction = ( + self.get_correction_factor(common_params, p_cor_consumer_state)) + if k_tile_count_init != k_tile_count: + mma_o_consumer_state = self.rescale( + common_params, + mma_o_consumer_state, + correction_factor, + no_correction, + ) + k_tile_count = k_tile_count - 1 + if k_tile_count == 0: + mma_o_consumer_state = self.epilogue( + common_params, + epilogue_params, + mma_o_consumer_state, + row_sum, + row_max, + ) + return p_cor_consumer_state, mma_o_consumer_state + + @cute.jit + def exchange_p_cor_metadata( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + correction_factor: cutlass.Float32, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + row_max_new: cutlass.Float32, + tAcc: cute.Tensor, + tidx: cutlass.Int32, + p_cor_producer_state: pipeline.PipelineState, + ) -> tuple[pipeline.PipelineState, cutlass.Float32]: + """Compute the correction factor for the last k tile.""" + no_correction = 0 + if ( + row_max_new - row_max + ) * softmax_params.softmax_scale_log2 <= self.skip_correction_threshold: + no_correction = 1 + row_max_new = row_max + + # pad for 4x32b + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.mma_s_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, + corr_layout, + ) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype) + corr_tmem_store_tiled_copy = tcgen05.make_tmem_copy( + corr_tmem_store_atom, tCor) + corr_tmem_store_thr_copy = corr_tmem_store_tiled_copy.get_slice(tidx) + cCor_for_copy = corr_tmem_store_thr_copy.partition_S(cCor) + tCor_for_copy = corr_tmem_store_thr_copy.partition_D(tCor) + rCor = cute.make_fragment_like(cCor_for_copy[None, None, None, 0], + self.acc_dtype) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout) + rCor[0] = row_sum + rCor[1] = row_max_new + rCor[2] = correction_factor + rCor_int[3] = no_correction + + cute.copy( + corr_tmem_store_tiled_copy, + rCor, + tCor_for_copy[None, None, None, p_cor_producer_state.index], + ) + # fence between tmem store and correction warp + cute.arch.fence_view_async_tmem_store() + common_params.p_cor_pipeline.producer_commit(p_cor_producer_state) + p_cor_producer_state.advance() + return p_cor_producer_state, row_max_new + + @cute.jit + def softmax( + self, + common_params: SimpleNamespace, + softmax_params: SimpleNamespace, + k_index: cutlass.Int32, + mma_s_consumer_state: pipeline.PipelineState, + p_mma_producer_state: pipeline.PipelineState, + p_cor_producer_state: pipeline.PipelineState, + row_max: cutlass.Float32, + row_sum: cutlass.Float32, + correction_factor: cutlass.Float32, + is_last_tile: bool, + is_local_last_tile: cutlass.Boolean, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + ]: + """Softmax for one k-tile. Updates the related pipeline states and returns the computed results. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param softmax_params: The softmax parameters + :type softmax_params: SimpleNamespace + :param k_index: The index of the k-tile + :type k_index: cutlass.Int32 + :param mma_s_consumer_state: The MMA s consumer state + :type mma_s_consumer_state: pipeline.PipelineState + :param p_mma_producer_state: The P MMA producer state + :type p_mma_producer_state: pipeline.PipelineState + :param p_cor_producer_state: The P correction producer state + :type p_cor_producer_state: pipeline.PipelineState + :param row_max: The row max + :type row_max: cutlass.Float32 + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param is_last_tile: Whether the last tile + :type is_last_tile: bool + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean + + :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] + """ + + softmax_params.p_mma_pipeline.producer_acquire(p_mma_producer_state) + softmax_params.mma_s_pipeline.consumer_wait(mma_s_consumer_state) + + # load S from tmem + tStS_shape = softmax_params.tiled_mma_qk.partition_shape_C( + cute.select(self.mma_qk_tiler, mode=[0, 1])) + tStS_staged_fake = softmax_params.tiled_mma_qk.make_fragment_C( + cute.append(tStS_shape, self.mma_s_stage)) + tStS_staged = cute.make_tensor(common_params.tmem_ptr, + tStS_staged_fake.layout) + tStS = tStS_staged[None, None, None, mma_s_consumer_state.index] + + tAcc = tStS[(None, None), 0, 0] + cta_qk_tiler = ( + self.mma_qk_tiler[0] // self.cluster_shape_mnk[0], + self.mma_qk_tiler[1], + self.mma_qk_tiler[2], + ) + cS = cute.make_identity_tensor(cute.select(cta_qk_tiler, mode=[0, 1])) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.acc_dtype) + tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + + tidx = common_params.tidx % (self.num_compute_warps * + self.threads_per_warp) + + tmem_thr_copy = tmem_tiled_copy.get_slice(tidx) + tTR_tAcc = tmem_thr_copy.partition_S(tAcc) + tTR_tS = tmem_thr_copy.partition_D(cS) + + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if is_last_tile: + tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) else -self.acc_dtype.inf) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, + row_max_new, 0) + elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + tmem_load_red_atom = cute.make_copy_atom( + tcgen05.copy.LdRed32x32bOp(tcgen05.copy.Repetition(64), + redOp=tcgen05.TmemLoadRedOp.MAX), + self.acc_dtype, + ) + tmem_red_tiled_copy = tcgen05.make_tmem_copy( + tmem_load_red_atom, tAcc) + tmem_red_thr_copy = tmem_red_tiled_copy.get_slice(tidx) + tTR_tAcc_red = tmem_red_thr_copy.partition_S(tAcc) + tTR_tS_red = tmem_red_thr_copy.partition_D(cS) + tTR_rAcc_red = cute.make_fragment_like(tTR_tS_red, self.acc_dtype) + tTR_rMax = cute.make_rmem_tensor( + cute.make_layout((1, tTR_tS_red.shape[1], tTR_tS_red.shape[2])), + self.acc_dtype, + ) + cute.copy( + tmem_red_tiled_copy, + tTR_tAcc_red, + (tTR_rAcc_red, tTR_rMax), + ) + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if is_last_tile: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less( + tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, + common_params.K, + ) else -self.acc_dtype.inf) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, + row_max_new, 0) + else: + row_max_new = cute.arch.fmax(row_max_new, tTR_rMax[0]) + + # if warps in N is 2, reduce row_max across warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_max_new + self.softmax_exchange_sync_bar.wait() + row_max_new = cute.arch.fmax( + row_max_new, + common_params.smem_exchange[(tidx + 64) % + (self.num_compute_warps * + self.threads_per_warp)], + ) + + # find correction factor + correction_factor = cute.math.exp2( + (row_max - row_max_new) * softmax_params.softmax_scale_log2, + fastmath=True) + # split kv case + if cutlass.const_expr(not is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # softmax + fma_b = softmax_params.softmax_scale_log2 + fma_c = (0.0 - row_max_new) * softmax_params.softmax_scale_log2 + + for i in cutlass.range(cute.size(tTR_rAcc), + vectorize=True, + unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * fma_b + fma_c + tTR_rAcc[i] = cute.math.exp2(tTR_rAcc[i], fastmath=True) + + tTR_rS = cute.make_fragment_like(tTR_tS, self.q_dtype) + + # quantize + tTR_rS.store(tTR_rAcc.load().to(self.q_dtype)) + + # create sP + sP = softmax_params.sP[None, None, None, + (None, p_mma_producer_state.index)] + sP_mk_view = cute.make_tensor( + sP.iterator, + cute.make_layout( + ( + (sP.shape[0][0], sP.shape[1]), + (sP.shape[0][1], sP.shape[2], sP.shape[3]), + ), + stride=( + (sP.stride[0][0], sP.stride[1]), + (sP.stride[0][1], sP.stride[2], sP.stride[3]), + ), + ), + ) + # change to PISL + sP_wo_swizzle_iter = cute.recast_ptr(sP.iterator, swizzle_=None) + swizzle_bits = (int( + math.log2(self.mma_pv_tiler[2] * self.q_dtype.width // 8 // 32)) + + 1) + swizzle_base = 3 if self.q_dtype.width == 16 else 4 + sP_swizzle = cute.make_swizzle(swizzle_bits, swizzle_base, 3) + sP_mk_view = cute.make_tensor( + sP_wo_swizzle_iter, + cute.make_composed_layout(sP_swizzle, 0, sP_mk_view.layout), + ) + universal_copy_bits = 128 + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=universal_copy_bits, + ) + smem_tiled_copy = cute.make_tiled_copy_D(smem_copy_atom, + tmem_tiled_copy) + smem_thr_copy = smem_tiled_copy.get_slice(tidx) + rP_copy_view = smem_thr_copy.retile(tTR_rS) + sP_copy_view = smem_thr_copy.partition_D(sP_mk_view) + cute.copy(smem_tiled_copy, rP_copy_view, sP_copy_view) + + # fence between smem store and mma o + cute.arch.fence_view_async_shared() + softmax_params.p_mma_pipeline.producer_commit(p_mma_producer_state) + p_mma_producer_state.advance() + + # row_sum, using `add_packed_f32x2` to reduce the number of instructions + row_sum = row_sum * correction_factor + row_sum_vec = (0.0, 0.0) + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + row_sum_vec = cute.arch.add_packed_f32x2( + row_sum_vec, (tTR_rAcc[i], tTR_rAcc[i + 1])) + row_sum = row_sum_vec[0] + row_sum_vec[1] + row_sum + + # split kv case + if cutlass.const_expr(is_local_last_tile): + p_cor_producer_state, row_max_new = self.exchange_p_cor_metadata( + common_params, + softmax_params, + correction_factor, + row_sum, + row_max, + row_max_new, + tAcc, + tidx, + p_cor_producer_state, + ) + + # store correction factor/row_sum/row_max to tmem for correction warp + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + + # fence between tmem load and mma s + cute.arch.fence_view_async_tmem_load() + + softmax_params.mma_s_pipeline.consumer_release(mma_s_consumer_state) + mma_s_consumer_state.advance() + + return ( + mma_s_consumer_state, + p_mma_producer_state, + p_cor_producer_state, + row_max_new, + row_sum, + correction_factor, + ) + + @cute.jit + def _tmem_load_partition( + self, common_params: SimpleNamespace, tiled_mma_pv: cute.TiledMma, + iter_n: int + ) -> tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, + cute.TiledMma]: + """Tensor memory load partition for rescale and epilogue. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param tiled_mma_pv: The tiled mma pv + :type tiled_mma_pv: cute.TiledMma + :param iter_n: The iteration number + :type iter_n: int + + :return: The tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv, the tiled mma pv + :rtype: tuple[cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma, cute.TiledMma] + """ + + tOtO_shape = tiled_mma_pv.partition_shape_C( + cute.select(self.mma_pv_tiler, mode=[0, 1])) + tOtO = tiled_mma_pv.make_fragment_C(tOtO_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + common_params.L // self.mma_pv_tiler[1], + stride=self.mma_pv_tiler[1] // self.warps_in_n, + ), + ) + tOtO = cute.make_tensor(common_params.tmem_ptr + self.tmem_o_offset, + tOtO_layout) + tOtO = tOtO[None, None, None, iter_n] + + tAcc = tOtO[(None, None), 0, 0] + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.acc_dtype) + tmem_load_tiled_copy = tcgen05.make_tmem_copy(tmem_load_atom, tAcc) + tmem_load_thr_copy = tmem_load_tiled_copy.get_slice( + common_params.tidx % + (self.num_compute_warps * self.threads_per_warp)) + + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + # Flatten divide and partition global tensors for O + cta_pv_tiler_mn = cute.select(cta_pv_tiler, mode=[0, 1]) + + gO = None + if cutlass.const_expr(common_params.mAccO is not None): + gO = cute.local_tile( + common_params.mAccO[None, common_params.blk_coord[3], None, + None, None], + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor( + common_params.mAccO[None, common_params.blk_coord[3], None, + None, None].shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + else: + gO = cute.local_tile( + common_params.mO, + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + cO = cute.local_tile( + cute.make_identity_tensor(common_params.mO.shape), + cta_pv_tiler_mn, + ( + common_params.blk_coord[0], + iter_n, + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + ) + tTR_tAcc = tmem_load_thr_copy.partition_S(tAcc) + tTR_gO = tmem_load_thr_copy.partition_D(gO) + tTR_cO = tmem_load_thr_copy.partition_D(cO) + tTR_rAcc = cute.make_fragment_like(tTR_gO, self.acc_dtype) + return tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc + + def get_correction_factor( + self, + common_params: SimpleNamespace, + p_cor_consumer_state: pipeline.PipelineState, + ) -> tuple[ + pipeline.PipelineState, + cutlass.Float32, + cutlass.Float32, + cutlass.Float32, + cutlass.Int32, + ]: + """Get the correction factor from the P correction consumer state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param p_cor_consumer_state: The P correction consumer state + :type p_cor_consumer_state: pipeline.PipelineState + + :return: The P correction consumer state, the row_sum, the row_max, and the correction factor + :rtype: tuple[pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32, cutlass.Int32] + """ + common_params.p_cor_pipeline.consumer_wait(p_cor_consumer_state) + tidx = common_params.tidx % (self.num_compute_warps * + self.threads_per_warp) + # load correction factor + _, tAcc, _, _, _, _ = self._tmem_load_partition( + common_params, common_params.tiled_mma_pv, 0) + corr_layout = cute.make_layout( + (tAcc.shape[0], (4, tAcc.shape[1][1]), self.p_cor_stage), + stride=(tAcc.stride[0], (1, tAcc.stride[1][1]), 4), + ) + tCor = cute.make_tensor( + common_params.tmem_ptr + self.correction_factor_offset, corr_layout) + cCor = cute.make_identity_tensor(tCor.shape) + corr_tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(4)), self.acc_dtype) + corr_tmem_load_tiled_copy = tcgen05.make_tmem_copy( + corr_tmem_load_atom, tCor) + corr_tmem_load_thr_copy = corr_tmem_load_tiled_copy.get_slice(tidx) + tCor_for_copy = corr_tmem_load_thr_copy.partition_S(tCor) + cCor_for_copy = corr_tmem_load_thr_copy.partition_D(cCor) + rCor = cute.make_fragment_like(cCor_for_copy[None, None, None, 0], + self.acc_dtype) + rCor_int = cute.make_tensor( + cute.recast_ptr(rCor.iterator, dtype=cutlass.Int32), rCor.layout) + cute.copy( + corr_tmem_load_tiled_copy, + tCor_for_copy[None, None, None, p_cor_consumer_state.index], + rCor, + ) + row_sum = rCor[0] + row_max = rCor[1] + correction_factor = rCor[2] + no_correction = rCor_int[3] + + common_params.p_cor_pipeline.consumer_release(p_cor_consumer_state) + p_cor_consumer_state.advance() + return p_cor_consumer_state, row_sum, row_max, correction_factor, no_correction + + @cute.jit + def rescale( + self, + common_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + correction_factor: cutlass.Float32, + no_correction: cutlass.Int32, + ) -> pipeline.PipelineState: + """Rescale for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param no_correction: Whether to apply correction factor + :type no_correction: cutlass.Int32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + skip_correction = cute.arch.vote_all_sync(no_correction == 1) + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + if not skip_correction: + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition(common_params, + common_params.tiled_mma_pv, + iter_n)) + + # tmem store tiled copy + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), + self.acc_dtype) + tmem_store_tiled_copy = tcgen05.make_tmem_copy( + tmem_store_atom, tAcc) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + # rescale, using `mul_packed_f32x2` to reduce the number of instructions + for i in cutlass.range(cute.size(tTR_rAcc), + vectorize=True, + unroll_full=True): + tTR_rAcc[i] = tTR_rAcc[i] * correction_factor + + # store o to tensor memory for next k tile + cute.copy(tmem_store_tiled_copy, tTR_rAcc, tTR_tAcc) + + cute.arch.fence_view_async_tmem_store() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + @cute.jit + def epilogue( + self, + common_params: SimpleNamespace, + epilogue_params: SimpleNamespace, + mma_o_consumer_state: pipeline.PipelineState, + row_sum: cutlass.Float32, + row_max: cutlass.Float32, + ) -> pipeline.PipelineState: + """Epilogue for one k-tile. Updates the related pipeline state. + + :param common_params: The common parameters + :type common_params: SimpleNamespace + :param epilogue_params: The epilogue parameters + :type epilogue_params: SimpleNamespace + :param mma_o_consumer_state: The mma o consumer state + :type mma_o_consumer_state: pipeline.PipelineState + :param row_sum: The row sum + :type row_sum: cutlass.Float32 + :param row_max: The row max + :type row_max: cutlass.Float32 + + :return: The MMA o consumer state + :rtype: pipeline.PipelineState + """ + + tidx = common_params.tidx % (self.num_compute_warps * + self.threads_per_warp) + + # exchange row_sum between warps (0, 1) and (2, 3) + if cutlass.const_expr(self.warps_in_n == 2): + common_params.smem_exchange[tidx] = row_sum + self.epilogue_exchange_sync_bar.wait() + # (64, 2) + row_sum = (row_sum + common_params.smem_exchange[ + (tidx + 64) % (self.num_compute_warps * self.threads_per_warp)]) + # mma_o pipeline consumer wait + for iter_n in cutlass.range_constexpr(self.iterations_pv_n): + common_params.mma_o_pipeline.consumer_wait(mma_o_consumer_state) + # tmem load tiled copy and partition results. + tmem_load_tiled_copy, tAcc, tTR_tAcc, tTR_gO, tTR_cO, tTR_rAcc = ( + self._tmem_load_partition(common_params, + common_params.tiled_mma_pv, iter_n)) + + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + + # apply output scale and normalize by row_sum + for i in cutlass.range(cute.size(tTR_rAcc), + vectorize=True, + unroll_full=True): + tTR_rAcc[i] = (tTR_rAcc[i] * epilogue_params.output_scale * + cute.arch.rcp_approx(row_sum)) + + # store o to global memory + tR2G_rO_src = None + tR2G_rO_dst = tTR_gO + if cutlass.const_expr(common_params.mAccO is None): + tR2G_rO_src = cute.make_fragment_like(tTR_gO, self.o_dtype) + # using final output dtype for o + tR2G_rO_src.store(tTR_rAcc.load().to(self.o_dtype)) + else: + # using accumulate dtype for o + tR2G_rO_src = tTR_rAcc + + if cute.elem_less(tTR_cO[0][0], common_params.H): + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority. + NO_ALLOCATE, + ) + + # store the lse to global memory + cta_pv_tiler = ( + self.mma_pv_tiler[0] // self.cluster_shape_mnk[0], + self.mma_pv_tiler[1], + self.mma_pv_tiler[2], + ) + gLSE = None + cLSE = None + if cutlass.const_expr(epilogue_params.mAccLSE is None): + gLSE = cute.local_tile( + epilogue_params.mLSE, + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor(epilogue_params.mLSE.shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + + else: + gLSE = cute.local_tile( + epilogue_params.mAccLSE[None, common_params.blk_coord[3], + None, None], + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + cLSE = cute.local_tile( + cute.make_identity_tensor( + epilogue_params.mAccLSE[None, + common_params.blk_coord[3], + None, None].shape), + (cta_pv_tiler[0], 1, 1), + ( + common_params.blk_coord[0], + common_params.blk_coord[1], + common_params.blk_coord[2], + ), + (1, 1, 1), + ) + lse = (cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max) + if cutlass.const_expr(self.warps_in_n == 2): + if cute.elem_less(cLSE[tidx][0], common_params.H): + gLSE[tidx] = lse + + cute.arch.fence_view_async_tmem_load() + common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) + mma_o_consumer_state.advance() + + return mma_o_consumer_state + + def make_and_init_load_qkv_pipeline(self, load_qkv_mbar_ptr, + cta_layout_vmnk, load_stages, + tx_count) -> pipeline.PipelineTmaUmma: + """Create and initialize the tma load qkv pipeline. + + :param load_qkv_mbar_ptr: The load qkv mbar pointer + :type load_qkv_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + :param load_stages: The load stages + :type load_stages: list[int] + :param tx_count: The tx count + :type tx_count: int + + :return: The tma load qkv pipeline + :rtype: pipeline.PipelineTmaUmma + """ + load_qkv_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_tma_k_warp_id])) + load_qkv_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_qkv_mbar_ptr, + num_stages=load_stages, + producer_group=load_qkv_producer_group, + consumer_group=load_qkv_consumer_group, + tx_count=tx_count, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_mma_s_pipeline( + self, mma_s_mbar_ptr, + cta_layout_vmnk) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma s pipeline. + + :param mma_s_mbar_ptr: The mma s mbar pointer + :type mma_s_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma s pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_s_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + consumer_thread_size = (self.threads_per_warp * + len(self.compute_warp_ids) * + self.cluster_shape_mnk[0]) + mma_s_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_s_mbar_ptr, + num_stages=self.mma_s_stage, + producer_group=mma_s_producer_group, + consumer_group=mma_s_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_mma_pipeline( + self, p_mma_mbar_ptr, + cta_layout_vmnk) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p mma pipeline. + + :param p_mma_mbar_ptr: The p mma mbar pointer + :type p_mma_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The p mma pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = (self.threads_per_warp * + len(self.compute_warp_ids) * + self.cluster_shape_mnk[0]) + p_mma_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_mma_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + return pipeline.PipelineAsyncUmma.create( + barrier_storage=p_mma_mbar_ptr, + num_stages=self.p_mma_stage, + producer_group=p_mma_producer_group, + consumer_group=p_mma_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + def make_and_init_p_cor_pipeline( + self, p_cor_mbar_ptr) -> pipeline.PipelineAsyncUmma: + """Create and initialize the p correction pipeline. + + :param p_cor_mbar_ptr: The p correction mbar pointer + :type p_cor_mbar_ptr: cute.Tensor + + :return: The p correction pipeline + :rtype: pipeline.PipelineAsyncUmma + """ + + producer_thread_size = self.threads_per_warp * len( + self.compute_warp_ids) + p_cor_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + p_cor_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + producer_thread_size, + ) + return pipeline.PipelineAsync.create( + barrier_storage=p_cor_mbar_ptr, + num_stages=self.p_cor_stage, + producer_group=p_cor_producer_group, + consumer_group=p_cor_consumer_group, + defer_sync=True, + ) + + def make_and_init_mma_o_pipeline( + self, mma_o_mbar_ptr, + cta_layout_vmnk) -> pipeline.PipelineUmmaAsync: + """Create and initialize the mma o pipeline. + + :param mma_o_mbar_ptr: The mma o mbar pointer + :type mma_o_mbar_ptr: cute.Tensor + :param cta_layout_vmnk: The cta layout vmnk + :type cta_layout_vmnk: tuple[int, int, int] + + :return: The mma o pipeline + :rtype: pipeline.PipelineUmmaAsync + """ + + mma_o_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id])) + consumer_thread_size = (self.threads_per_warp * + len(self.compute_warp_ids) * + self.cluster_shape_mnk[0]) + mma_o_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + consumer_thread_size, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_o_mbar_ptr, + num_stages=self.mma_o_stage, + producer_group=mma_o_producer_group, + consumer_group=mma_o_consumer_group, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + + @staticmethod + def _compute_grid( + o: cute.Tensor, + split_kv: cutlass.Int32, + cluster_shape_mnk: Tuple[int, int, int], + max_active_clusters: int, + is_persistent: bool, + ) -> Tuple[MLAStaticTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid shape for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + + :return: Tile scheduler parameters and grid shape. + :rtype: tuple[MLAStaticTileSchedulerParams, tuple[int, int, int]] + """ + o_shape = o.shape + tile_sched_params = create_mla_static_tile_scheduler_params( + is_persistent, + cute.size(o_shape[3]), + cute.size(o_shape[2]), + cluster_shape_mnk, + split_kv, + ) + grid = MLAStaticTileScheduler.get_grid_shape(tile_sched_params, + max_active_clusters) + + return tile_sched_params, grid + + @staticmethod + def get_workspace_size( + H: int, + S: int, + D: int, + B: int, + split_kv: int, + acc_dtype: Type[cutlass.Numeric], + ) -> int: + """Get the extra workspace(device memory) size for the MLA kernel when split_kv is not 1. + + :param H: The height of the output tensor C + :type H: int + :param S: The sequence length of the output tensor C + :type S: int + :param D: The depth of the output tensor C + :type D: int + :param B: The batch size of the output tensor C + :type B: int + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + + :return: The workspace size for the MLA kernel + :rtype: int + """ + if split_kv == 1: + return 0 + return B * H * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( + self, + H: cutlass.Int32, + D: cutlass.Int32, + S: cutlass.Int32, + B: cutlass.Int32, + split_kv: cutlass.Int32, + acc_dtype: Type[cutlass.Numeric], + workspace: cute.Tensor, + ) -> tuple[cute.Tensor, cute.Tensor]: + """Initialize the workspace for the MLA kernel. Construct the intermediate tensors + acc_o and acc_lse. + + :param H: The height of the output tensor C + :type H: cutlass.Int32 + :param D: The depth of the output tensor C + :type D: cutlass.Int32 + :param S: The sequence length of the output tensor C + :type S: cutlass.Int32 + :param B: The batch size of the output tensor C + :type B: cutlass.Int32 + :param split_kv: The split key-value of the output tensor C + :type split_kv: cutlass.Int32 + :param acc_dtype: The data type of the output tensor C + :type acc_dtype: Type[cutlass.Numeric] + :param workspace: The workspace tensor + :type workspace: cute.Tensor + + :return: The output tensor C and the workspace tensor + :rtype: tuple[cute.Tensor, cute.Tensor] + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): + align = 256 // self.q_dtype.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=acc_dtype) + 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 + + cute.cosize(acc_o_layout) * acc_dtype.width // 8, + dtype=acc_dtype, + ) + acc_lse = cute.make_tensor(acc_lse_iter, acc_lse_layout) + return acc_o, acc_lse + + @staticmethod + def can_implement( + B: int, + S: int, + K: int, + H: int, + L: int, + R: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + ) -> bool: + """Check if the MLA kernel can be implemented. + + :param B: The batch size of the output tensor C + :type B: int + :param S: The sequence length of the output tensor C + :type S: int + :param K: The width of the output tensor KV + :type K: int + :param H: The number of heads of the output tensor C + :type H: int + :param L: The number of latent dimensions of the tensor KV + :type L: int + :param R: The number of rope dimensions of the tensor C_rope + :type R: int + :param in_dtype: The data type of the input tensor + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: The data type of the output tensor + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: The data type of the log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: The tile shape of the query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: The tile shape of the probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: The split key-value of the output tensor C + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: The page size of the page table + :type page_size: int + + :return: Whether the MLA kernel can be implemented + :rtype: bool + """ + if L != 512 or R != 64: + return False + if in_dtype not in [cutlass.Float8E4M3FN]: + return False + if out_dtype not in [cutlass.Float8E4M3FN]: + return False + if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: + return False + # page size equals 1 is prohibited by tma specification, not 128B aligned. + if mma_qk_tiler_mn[1] % page_size != 0 or page_size == 1: + return False + if mma_qk_tiler_mn[0] != mma_pv_tiler_mn[0] or mma_qk_tiler_mn[0] != 128: + return False + if is_var_split_kv and not is_var_seq: + return False + if H > 128 or (H < 128 and split_kv != 1): + return False + if S <= 0 or S > 4: + return False + if K <= 0: + return False + return True + + +def run( + batch_size: int, + seq_len_q: int, + seq_len_k: int, + num_heads: int, + latent_dim: int, + rope_dim: int, + in_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + lse_dtype: Type[cutlass.Numeric], + mma_qk_tiler_mn: Tuple[int, int], + mma_pv_tiler_mn: Tuple[int, int], + split_kv: int, + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, + page_size: int, + softmax_scale: float, + output_scale: float, + skip_correction_threshold: float, + tolerance: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool, + **kwargs, +): + """Execute Multi-Head Latent Attention (MLA) on Blackwell architecture and validate results. + + This function creates random input tensors for query latent/rope, compressed latent/rope, and value, + then performs the complete MLA computation pipeline. It supports configurable data types, tiling parameters, + page table, variable sequence length, and variable split_kv. Results can be validated against a PyTorch reference + implementation or run multiple times for performance measurement. + + :param batch_size: Batch size + :type batch_size: int + :param seq_len_q: Sequence length of Q + :type seq_len_q: int + :param seq_len_k: Sequence length of K + :type seq_len_k: int + :param num_heads: Number of heads + :type num_heads: int + :param latent_dim: dimension of query/compressed latent + :type latent_dim: int + :param rope_dim: dimension of query/compressed rope + :type rope_dim: int + :param in_dtype: Input data type for query/compressed latent/rope tensors + :type in_dtype: Type[cutlass.Numeric] + :param out_dtype: Output data type for attention output + :type out_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for query-key matrix multiplication + :type acc_dtype: Type[cutlass.Numeric] + :param lse_dtype: Accumulator data type for log-sum-exp + :type lse_dtype: Type[cutlass.Numeric] + :param mma_qk_tiler_mn: Matrix multiply accumulate tile shape (M, N) for query-key matrix multiplication + :type mma_qk_tiler_mn: Tuple[int, int] + :param mma_pv_tiler_mn: Matrix multiply accumulate tile shape (M, N) for probability-value matrix multiplication + :type mma_pv_tiler_mn: Tuple[int, int] + :param split_kv: Split key-value + :type split_kv: int + :param is_persistent: Whether to use persistent kernel optimization + :type is_persistent: bool + :param is_var_seq: Whether to use variable sequence length + :type is_var_seq: bool + :param is_var_split_kv: Whether to use variable split_kv + :type is_var_split_kv: bool + :param page_size: Page size of the page table + :type page_size: int + :param softmax_scale: Attention score scaling factor + :type softmax_scale: float + :param output_scale: Output scaling factor + :type output_scale: float + :param skip_correction_threshold: Threshold to skip correction + :type skip_correction_threshold: float + :param tolerance: Maximum acceptable error for validation + :type tolerance: float + :param warmup_iterations: Number of warmup iterations + :type warmup_iterations: int + :param iterations: Number of iterations to run for performance testing + :type iterations: int + :param skip_ref_check: Skip validation against reference implementation + :type skip_ref_check: bool + :param use_cold_l2: Whether to use cold L2 cache + :type use_cold_l2: bool + + :raises ValueError: If input shapes are incompatible or head dimension is unsupported + :raises RuntimeError: If GPU is unavailable for computation + """ + + print("Running Blackwell MLA test with:") + print(f" batch_size: {batch_size}") + print(f" seq_len_q: {seq_len_q}") + print(f" seq_len_k: {seq_len_k}") + print(f" num_heads: {num_heads}") + print(f" latent_dim: {latent_dim}") + print(f" rope_dim: {rope_dim}") + print(f" in_dtype: {in_dtype}") + print(f" out_dtype: {out_dtype}") + print(f" acc_dtype: {acc_dtype}") + print(f" mma_qk_tiler_mn: {mma_qk_tiler_mn}") + print(f" mma_pv_tiler_mn: {mma_pv_tiler_mn}") + print(f" split_kv: {split_kv}") + print(f" is_persistent: {is_persistent}") + print(f" is_var_seq: {is_var_seq}") + print(f" is_var_split_kv: {is_var_split_kv}") + print(f" page_size: {page_size}") + print(f" softmax_scale: {softmax_scale}") + print(f" output_scale: {output_scale}") + print(f" skip_correction_threshold: {skip_correction_threshold}") + print(f" tolerance: {tolerance}") + print(f" warmup_iterations: {warmup_iterations}") + print(f" iterations: {iterations}") + print(f" skip_ref_check: {skip_ref_check}") + print(f" use_cold_l2: {use_cold_l2}") + + import cutlass.torch as cutlass_torch + import torch + + # Prepare pytorch tensors: Q, K, V (random from 0 to 2) and O (all zero) + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not BlackwellMultiHeadLatentAttentionForwardFP8.can_implement( + batch_size, + seq_len_q, + seq_len_k, + num_heads, + latent_dim, + rope_dim, + in_dtype, + out_dtype, + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + split_kv, + is_persistent, + is_var_seq, + is_var_split_kv, + page_size, + ): + raise TypeError( + f"Unsupported testcase {batch_size}, {seq_len_q}, {seq_len_k}, {num_heads}, {latent_dim}, {rope_dim}, {in_dtype}, {out_dtype}, {acc_dtype}, {lse_dtype}, {mma_qk_tiler_mn}, {mma_pv_tiler_mn}, {split_kv}, {is_persistent}, {is_var_seq}, {is_var_split_kv}, {page_size}" + ) + + torch.manual_seed(1111) + + def create_data_tensor( + B, + HK, + D, + dtype, + is_dynamic_layout=True, + page_table=None, + cache_seqs=None, + is_lse=False, + seq_len_q=None, + ): + shape = (B, HK, D) + if page_table is not None: + if cache_seqs is not None: + max_seq_len = torch.max(cache_seqs) + shape = (B * ceil_div(max_seq_len, page_size), page_size, D) + else: + shape = (B * ceil_div(HK, page_size), page_size, D) + + if seq_len_q is not None: + shape = (B, seq_len_q, HK, D) + + permute_order = (1, 2, 0) + stride_order = (2, 0, 1) + leading_dim = 1 + if is_lse: + shape = (B, seq_len_q, HK) + permute_order = (2, 1, 0) + stride_order = (2, 1, 0) + leading_dim = 0 + elif seq_len_q is not None: + permute_order = (2, 3, 1, 0) + stride_order = (3, 2, 0, 1) + leading_dim = 1 + + init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) + + torch_dtype = (cutlass_torch.dtype(dtype) + if dtype != cutlass.Float8E4M3FN else torch.int8) + + # Create dtype torch tensor (cpu) + torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch_dtype, + permute_order=permute_order, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=init_config, + ) + + # Create dtype torch tensor (gpu) + torch_tensor_gpu = torch_tensor_cpu.cuda() + + # Create f32 torch tensor (cpu) + f32_torch_tensor = torch_tensor_cpu.to(dtype=torch.float32) + + # Create dtype cute tensor (gpu) + cute_tensor = from_dlpack(torch_tensor_gpu, assumed_align=16) + cute_tensor.element_type = dtype + if is_dynamic_layout: + cute_tensor = cute_tensor.mark_layout_dynamic( + leading_dim=leading_dim) + if not is_lse: + cute_tensor = cute_tensor.mark_compact_shape_dynamic( + mode=leading_dim, + stride_order=stride_order, + divisibility=(128 // dtype.width), + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=is_dynamic_layout, + ) + + return f32_torch_tensor, cute_tensor, torch_tensor_gpu + + def create_cache_seqs(batch_size, seq_len_k, is_var_seq): + cache_seqs_ref = torch.ones(batch_size, dtype=torch.int32) * seq_len_k + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack(cache_seqs_gpu, + assumed_align=16).mark_layout_dynamic() + if is_var_seq: + max_seq_len = seq_len_k + min_seq_len = int(seq_len_k * 0.8) + cache_seqs_ref = cutlass_torch.create_and_permute_torch_tensor( + (batch_size, ), + torch.int32, + init_type=cutlass.torch.TensorInitType.RANDOM, + init_config=cutlass.torch.RandomInitConfig(min_val=min_seq_len, + max_val=max_seq_len + + 1), + ) + cache_seqs_gpu = cache_seqs_ref.cuda() + cache_seqs = from_dlpack( + cache_seqs_gpu, + assumed_align=16, + ).mark_layout_dynamic() + return cache_seqs_ref, cache_seqs, cache_seqs_gpu + + def create_page_table(batch_size, seq_len_k, is_var_seq, page_size): + max_seq_len = seq_len_k if not is_var_seq else torch.max(cache_seqs_ref) + page_count = ceil_div(max_seq_len, page_size) + page_table_ref = torch.empty([batch_size, page_count], + dtype=torch.int32) + # use transposed index for page table to make sure the value is in bound of `batch_size * seq_len_block`. In practice, the value could be any positive values. This setting is only for testing purpose. + for b in range(batch_size): + for j in range(page_count): + page_table_ref[b, j] = b + j * batch_size + page_table_gpu = page_table_ref.permute(1, 0).cuda() + page_table = from_dlpack( + page_table_gpu, assumed_align=16).mark_layout_dynamic(leading_dim=0) + return page_table_ref, page_table, page_table_gpu + + def create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ): + block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu = None, None, None + # check if split_kv is valid otherwise do auto setting of split_kv + if is_var_split_kv: + block_split_kvs_ref = torch.zeros([batch_size], dtype=torch.int32) + for b in range(batch_size): + block_split_kvs_ref[b] = ( + BlackwellMultiHeadLatentAttentionForwardFP8.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[b].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + )) + split_kv = torch.max(block_split_kvs_ref).item() + block_split_kvs_gpu = block_split_kvs_ref.cuda() + block_split_kvs = from_dlpack( + block_split_kvs_gpu, assumed_align=16).mark_layout_dynamic() + elif split_kv <= 0: + split_kv = BlackwellMultiHeadLatentAttentionForwardFP8.get_split_kv( + batch_size, + seq_len_q, + cache_seqs_ref[0].item(), + mma_qk_tiler_mn, + max_active_clusters * cluster_shape_mnk[0], + ) + return split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_gpu + + def create_workspace(num_heads, seq_len_q, latent_dim, batch_size, split_kv, + acc_dtype): + workspace_size = BlackwellMultiHeadLatentAttentionForwardFP8.get_workspace_size( + num_heads, + seq_len_q, + latent_dim, + batch_size, + split_kv, + acc_dtype, + ) + + workspace, workspace_torch = None, None + if workspace_size > 0: + workspace_torch = torch.empty([workspace_size], + dtype=torch.int8).cuda() + workspace = from_dlpack(workspace_torch, assumed_align=32) + return workspace, workspace_torch + + cache_seqs_ref, cache_seqs, cache_seqs_torch = create_cache_seqs( + batch_size, seq_len_k, is_var_seq) + page_table_ref, page_table, page_table_torch = create_page_table( + batch_size, seq_len_k, is_var_seq, page_size) + cluster_shape_mnk = (2, 1, 1) + hardware_info = utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mnk[0] * cluster_shape_mnk[1]) + split_kv, block_split_kvs_ref, block_split_kvs, block_split_kvs_torch = ( + create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + )) + + q_latent_ref, q_latent, q_latent_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + q_rope_ref, q_rope, q_rope_torch = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + c_latent_ref, c_latent, c_latent_torch = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + c_rope_ref, c_rope, c_rope_torch = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + o_ref, o, o_torch = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + lse_ref, lse, lse_torch = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace(num_heads, seq_len_q, + latent_dim, batch_size, + split_kv, acc_dtype) + + mla = BlackwellMultiHeadLatentAttentionForwardFP8( + acc_dtype, + lse_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + max_active_clusters, + page_size, + skip_correction_threshold, + is_persistent, + is_var_seq, + is_var_split_kv, + ) + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + stream = cuda.CUstream(torch_stream.cuda_stream) + + # compile mla kernel + compiled_mla = cute.compile( + mla, + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + options="--opt-level 2", + ) + + def torch_reference_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + cache_seqs, + softmax_scale=1.0, + output_scale=1.0, + ): + # expand and concat q_latent and q_rope to have the dimension of sequence length for q + q_ref = torch.cat([q_latent, q_rope], dim=1).permute(3, 2, 0, 1) + # expand and concat c_latent and c_rope to have the dimension of num_heads for k and v + page_count = page_table_ref.shape[1] + k_ref_paged = (torch.cat([c_latent, c_rope], + dim=1).permute(2, 0, 1).reshape( + batch_size * page_count, page_size, + latent_dim + rope_dim)) + v_ref_paged = c_latent.permute(2, 0, 1).reshape(batch_size * page_count, + page_size, latent_dim) + + if is_var_seq: + max_seq_len = torch.max(cache_seqs_ref) + else: + max_seq_len = seq_len_k + + k_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim + rope_dim]) + v_ref = torch.zeros([batch_size, 1, max_seq_len, latent_dim]) + k_ref = torch.index_select( + k_ref_paged, 0, torch.flatten(page_table_ref)).reshape( + batch_size, 1, -1, latent_dim + rope_dim)[:, :, :max_seq_len, :] + v_ref = torch.index_select(v_ref_paged, 0, + torch.flatten(page_table_ref)).reshape( + batch_size, 1, -1, + latent_dim)[:, :, :max_seq_len, :] + for b in range(batch_size): + k_ref[b, :, cache_seqs_ref[b]:, :] = 0 + v_ref[b, :, cache_seqs_ref[b]:, :] = 0 + import torch.nn.functional as F + + o_ref = F.scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + attn_mask=None, + dropout_p=0.0, + scale=softmax_scale, + is_causal=False, + ) + s_ref = torch.einsum("bhld,bhsd->bhls", q_ref, k_ref) + s_ref_max, s_ref_max_pos = torch.max(s_ref, dim=-1, keepdim=True) + softmax_scale_log2 = LOG2_E * softmax_scale + s_ref_sum = torch.sum(torch.exp2( + (s_ref - s_ref_max) * softmax_scale_log2), + dim=-1, + keepdim=True) + + lse_ref = s_ref_max * softmax_scale_log2 + torch.log2(s_ref_sum) + lse_ref = lse_ref.squeeze(3).permute(2, 1, 0) + o_ref = o_ref * output_scale + o_ref = o_ref.permute(2, 3, 1, 0) + + return o_ref, lse_ref + + if skip_correction_threshold > 0.0: + print( + "Skipping correction verification since skip_correction_threshold is greater than 0.0..." + ) + skip_ref_check = True + if not skip_ref_check: + # Execute kernel once for reference checking + compiled_mla( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + torch.cuda.synchronize() + + print("Verifying results...") + if in_dtype == cutlass.Float8E4M3FN: + tolerance = 0.13 + o_ref, lse_ref = torch_reference_mla( + q_latent_ref, + q_rope_ref, + c_latent_ref, + c_rope_ref, + page_table, + cache_seqs, + softmax_scale, + output_scale, + ) + + if out_dtype in [cutlass.Float8E5M2, cutlass.Float8E4M3FN]: + # convert o back to f32 for comparison + o_fp32, o_fp32_torch = cutlass_torch.cute_tensor_like( + torch.empty(*o_torch.shape, dtype=torch.float32), + cutlass.Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(o, o_fp32) + o = o_fp32_torch.cpu() + ref_fp8, _ = cutlass_torch.cute_tensor_like( + torch.empty(*o_ref.permute(3, 2, 0, 1).shape, + dtype=torch.uint8).permute(2, 3, 1, 0), + out_dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + o_ref_gpu = o_ref.cuda() + o_ref_f32 = from_dlpack(o_ref_gpu).mark_layout_dynamic( + leading_dim=1) + + # convert ref : f32 -> fp8 -> f32 + cute.testing.convert(o_ref_f32, ref_fp8) + cute.testing.convert(ref_fp8, o_ref_f32) + + o_ref = o_ref_gpu.cpu() + else: + o = o_torch.cpu().to(torch.float32) + lse = lse_torch.cpu() + lse_ref = lse_ref.to(cutlass.torch.dtype(lse_dtype)) + # Assert close results + torch.testing.assert_close(o, o_ref, atol=tolerance, rtol=1e-05) + torch.testing.assert_close(lse, lse_ref, atol=tolerance, rtol=1e-05) + print("Results verified successfully!") + + def generate_tensors(): + _, cache_seqs, _ = create_cache_seqs(batch_size, seq_len_k, is_var_seq) + _, page_table, _ = create_page_table(batch_size, seq_len_k, is_var_seq, + page_size) + _split_kv, _, block_split_kvs, _ = create_block_split_kvs( + batch_size, + split_kv, + cache_seqs_ref, + is_var_split_kv, + mma_qk_tiler_mn, + cluster_shape_mnk, + max_active_clusters, + ) + + _, q_latent, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, q_rope, _ = create_data_tensor( + batch_size, + num_heads, + rope_dim, + in_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + + _, c_latent, _ = create_data_tensor( + batch_size, + seq_len_k, + latent_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, c_rope, _ = create_data_tensor( + batch_size, + seq_len_k, + rope_dim, + in_dtype, + is_dynamic_layout=True, + page_table=page_table, + cache_seqs=cache_seqs_ref, + ) + _, o, _ = create_data_tensor( + batch_size, + num_heads, + latent_dim, + out_dtype, + is_dynamic_layout=True, + seq_len_q=seq_len_q, + ) + _, lse, _ = create_data_tensor( + batch_size, + num_heads, + 1, + lse_dtype, + is_dynamic_layout=True, + is_lse=True, + seq_len_q=seq_len_q, + ) + workspace, workspace_torch = create_workspace(num_heads, seq_len_q, + latent_dim, batch_size, + _split_kv, acc_dtype) + return testing.JitArguments( + q_latent, + q_rope, + c_latent, + c_rope, + page_table, + o, + lse, + workspace, + _split_kv, + cache_seqs, + block_split_kvs, + softmax_scale, + output_scale, + stream, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + q_latent_torch.numel() * q_latent_torch.element_size() + + q_rope_torch.numel() * q_rope_torch.element_size() + + c_latent_torch.numel() * c_latent_torch.element_size() + + c_rope_torch.numel() * c_rope_torch.element_size() + + o_torch.numel() * o_torch.element_size() + + lse_torch.numel() * lse_torch.element_size() + + cache_seqs_torch.numel() * cache_seqs_torch.element_size()) + one_workspace_bytes += (page_table_torch.numel() * + page_table_torch.element_size()) + if is_var_split_kv: + one_workspace_bytes += (block_split_kvs_torch.numel() * + block_split_kvs_torch.element_size()) + if workspace_torch is not None: + one_workspace_bytes += (workspace_torch.numel() * + workspace_torch.element_size()) + workspace_count = testing.get_workspace_count(one_workspace_bytes, + warmup_iterations, + iterations) + + avg_time_us = testing.benchmark( + compiled_mla, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + return avg_time_us # Return execution time in microseconds + + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers.") + + def parse_mma_tiler(s: str) -> Tuple[int, int, Tuple[int, int]]: + ret = parse_comma_separated_ints(s) + if len(ret) != 2: + raise argparse.ArgumentTypeError( + "Invalid format. Expected 2 comma-separated integers.") + return (ret[0], ret[1]) + + parser = argparse.ArgumentParser(description="Example of MLA on Blackwell.") + + parser.add_argument( + "--in_dtype", + type=cutlass.dtype, + default=cutlass.Float8E4M3FN, + help="Input data type", + ) + + parser.add_argument( + "--out_dtype", + type=cutlass.dtype, + default=cutlass.Float8E4M3FN, + help="Output data type", + ) + + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="Accumulator data type", + ) + + parser.add_argument( + "--lse_dtype", + type=cutlass.dtype, + default=cutlass.Float32, + help="LSE data type", + ) + parser.add_argument( + "--mma_qk_tiler_mn", + type=parse_mma_tiler, + default=(128, 128), + help="MMA tile shape (H, K)", + ) + parser.add_argument( + "--mma_pv_tiler_mn", + type=parse_mma_tiler, + default=(128, 256), + help="MMA tile shape (H, D)", + ) + + parser.add_argument( + "--is_persistent", + action="store_true", + help="Is persistent", + ) + + parser.add_argument( + "--batch_size", + type=int, + default=1, + help="Batch size", + ) + + parser.add_argument( + "--seq_len_q", + type=int, + default=1, + help="Sequence length of Q", + ) + + parser.add_argument( + "--seq_len_k", + type=int, + default=128, + help="Sequence length of K/V", + ) + + parser.add_argument( + "--num_heads", + type=int, + default=128, + help="Number of heads of Q", + ) + + parser.add_argument( + "--latent_dim", + type=int, + default=512, + help="Latent dimension of Q/C", + ) + + parser.add_argument( + "--rope_dim", + type=int, + default=64, + help="Rope dimension of Q/C", + ) + + parser.add_argument( + "--is_var_seq", + action="store_true", + help="Use variable length of sequence length or not", + ) + + parser.add_argument( + "--is_var_split_kv", + action="store_true", + help="Use variable length of split kv or not", + ) + + parser.add_argument( + "--page_size", + type=int, + default=128, + help="Page size of page table", + ) + + parser.add_argument( + "--split_kv", + type=int, + default=-1, + help="Split KV setting", + ) + + parser.add_argument( + "--softmax_scale", + type=float, + default=0.0416, + help="Scaling factor to scale softmax", + ) + + parser.add_argument( + "--output_scale", + type=float, + default=1.0, + help="Scaling factor to scale output", + ) + parser.add_argument( + "--skip_correction_threshold", + type=float, + default=0.0, + help="Threshold to skip correction", + ) + + parser.add_argument("--tolerance", + type=float, + default=1e-02, + help="Tolerance for validation") + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + help="Use cold L2 cache", + ) + + args = parser.parse_args() + + run( + args.batch_size, + args.seq_len_q, + args.seq_len_k, + args.num_heads, + args.latent_dim, + args.rope_dim, + args.in_dtype, + args.out_dtype, + args.acc_dtype, + args.lse_dtype, + args.mma_qk_tiler_mn, + args.mma_pv_tiler_mn, + args.split_kv, + args.is_persistent, + args.is_var_seq, + args.is_var_split_kv, + args.page_size, + args.softmax_scale, + args.output_scale, + args.skip_correction_threshold, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py new file mode 100644 index 000000000000..f3e1edd762f0 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.py @@ -0,0 +1,302 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import cutlass +import cutlass.cute as cute + + +class MLAStaticTileSchedulerParams: + + 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, + ): + """The static tile scheduler parameters prepared for MLA static tile scheduler. + + :param is_persistent: Whether to use persistent kernel mode + :type is_persistent: bool + :param problem_shape_b: The shape of the problem + :type problem_shape_b: cute.Int32 + :param problem_shape_s: The shape of the problem in sequence length Q dimension + :type problem_shape_s: cute.Int32 + :param cluster_shape_mnk: The shape of the cluster + :type cluster_shape_mnk: cute.Shape + :param split_kv: The scalar factor for split KV + """ + 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, + ) + + +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: + return MLAStaticTileSchedulerParams(is_persistent, problem_shape_b, + problem_shape_s, cluster_shape_mnk, + split_kv) + + +class WorkTileInfo: + + 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 self.is_valid + + @property + def tile_idx(self) -> cute.Coord: + return self.blk_coord + + +class MLAStaticTileScheduler: + + 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, + ): + """The static tile scheduler for MLA split kv kernel. + Based on `is_persistent`, it provides 2 modes for use: + - Persistent mode: Launch fixed blocks and reschedule the data blocks. + - Non-persistent mode: Launch dynamic blocks and exit when the current work is done. + + :param params: The static tile scheduler parameters + :type params: MLAStaticTileSchedulerParams + :param current_work_linear_idx: The linear index of the current work + :type current_work_linear_idx: cutlass.Int32 + :param blk_coord: The coordinate of the current work + :type blk_coord: cute.Coord + :param grid_shape: The shape of the grid + :type grid_shape: cute.Shape + :param is_valid: Whether the current work is valid + :type is_valid: bool + """ + 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) + # Used for persistent scheduling + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + else: + self.is_valid = 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: + # called by host + 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, + ) + else: + return grid_shape + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + 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 self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + if self.params.is_persistent: + self.current_work_linear_idx += advance_count * self.num_persistent_sm + else: + self.is_valid = 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)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 13 + 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:]) + return MLAStaticTileScheduler(new_params, new_current_work_linear_idx, + new_blk_coord, new_grid_shape) + + +def create_mla_static_tile_scheduler( + params: MLAStaticTileSchedulerParams, + blk_coord: cute.Coord, + grid_shape: cute.Shape, +) -> MLAStaticTileScheduler: + return MLAStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape) + + +LOG2_E = 1.4426950408889634074 +# avoid register indexing on array. +MAX_SPLITS = 256 + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 994bede176c1..18a09d54daa2 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -375,6 +375,7 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: accuracy_dict = { torch.bfloat16: (3e-2, 3e-3), + torch.float16: (3e-2, 3e-3), torch.float8_e4m3fn: (4e-1, 4e-2), } @@ -504,7 +505,9 @@ def test_attention_mla(scenario: Scenario, context_sequence_lengths: List[int], f"--------------------------------Test for scenario: {scenario} start--------------------------------" ) - _run_test_for_backend("TRTLLM", num_heads, num_kv_heads, num_layers, + import os + backend_name = os.environ.get("MLA_TEST_BACKEND", "TRTLLM") + _run_test_for_backend(backend_name, num_heads, num_kv_heads, num_layers, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, rope_config, kv_cache_tokens_per_block, device, dtype, @@ -589,13 +592,28 @@ def test_attention_mla_flashinfer(scenario: Scenario, v2_kv_cache) -def _run_test_for_backend(backend_name, num_heads, num_kv_heads, num_layers, - q_lora_rank, kv_lora_rank, qk_nope_head_dim, - qk_rope_head_dim, v_head_dim, rope_config, - kv_cache_tokens_per_block, device, dtype, - kv_cache_dtype, context_sequence_lengths, - generation_seq_len_q, num_generation_steps, - v2_kv_cache): +def _run_test_for_backend(backend_name, + num_heads, + num_kv_heads, + num_layers, + q_lora_rank, + kv_lora_rank, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + rope_config, + kv_cache_tokens_per_block, + device, + dtype, + kv_cache_dtype, + context_sequence_lengths, + generation_seq_len_q, + num_generation_steps, + v2_kv_cache, + skip_context_assert=False): + # When ``skip_context_assert`` is set, the context (step 0) result is not + # checked; the context phase only runs to populate the KV cache and build + # the reference latent cache. Used by the decode-only CuTe DSL MLA test. AttentionCls = get_attention_backend(backend_name) qk_head_dim = qk_nope_head_dim + qk_rope_head_dim @@ -1064,7 +1082,12 @@ def yarn_get_mscale(scale=1, mscale=1): f"Difference mean: {(result - ref_result).abs().mean().item()}, max: {(result - ref_result).abs().max().item()}" ) - # Assert results are close + # Assert results are close (skip context/step-0 when requested: + # the decode-only test treats the context phase as cache setup). + if skip_context_assert and step == 0: + print(f"Skipping context (step 0) assertion for {backend_name} " + f"backend at layer {layer_idx} (decode-only mode)") + continue atol, rtol = accuracy_dict[kv_cache_dtype] assert torch.allclose(result, ref_result, atol=atol, rtol=rtol), \ f"Results for MLA in {backend_name} backend don't match reference implementation at layer {layer_idx} in step {step}" diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py new file mode 100644 index 000000000000..243c28ed989c --- /dev/null +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py @@ -0,0 +1,165 @@ +# 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. +"""Decode-only MLA test for the Blackwell CuTe DSL MLA decode kernels. + +This test validates the CuTe DSL MLA *decode* kernels added under +``tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla`` and dispatched +through the ``CUTEDSL`` attention backend +(``tensorrt_llm/_torch/attention_backend/cute_dsl.py``): + +- FP8 path → ``torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell`` +- FP16 path → ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` + +Only the generation (decode) steps are asserted for numerical correctness. +The context phase runs solely to populate the paged KV cache and to build the +reference latent cache (``skip_context_assert=True``). + +Crucially, the test monkeypatches ``CuteDslAttention._dispatch_cute_dsl_mla_decode`` +to count invocations and asserts the CuTe DSL decode path was actually taken on +every decode step. Without this guard the backend would silently fall back to +TRTLLM on any kernel error and the test could "pass" without ever exercising +the CuTe DSL kernel under test. + +Platform: Blackwell SM100 / SM103 only. +""" + +import pytest +import torch + +# Reuse the proven setup + reference machinery from the full MLA test. +# The attention test directory is added to sys.path by pytest (prepend import +# mode, no package __init__), so the sibling module is imported by bare name. +from test_attention_mla import RopeConfig, Scenario, _run_test_for_backend + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + +# DeepSeek-V3-like MLA geometry the CuTe DSL kernel targets (num_heads=128, +# latent_dim=512, rope_dim=64). Kept small along the batch/step axes so the +# decode-only test stays fast. +_DECODE_CONTEXT_LENGTHS = [ + [10, 12, 5], + [100, 300, 20, 10], +] +_DECODE_NUM_STEPS = 4 + + +def _is_blackwell_sm100() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) + + +pytestmark = [ + pytest.mark.skipif( + not _is_blackwell_sm100(), + reason="CuTe DSL MLA decode kernels require Blackwell SM100/SM103.", + ), + pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="nvidia-cutlass-dsl is not available."), +] + + +# kernel name -> (activation dtype, kv cache dtype) +# +# NOTE: only the FP8 decode kernel is currently validated here. The FP16 path +# (``cute_dsl_mla_decode_fp16_blackwell`` with dtype=float16 / fp16 KV cache) +# aborts the process (SIGABRT) on SM100 in this environment, so it is excluded +# until that crash is root-caused. Re-add "fp16" below once it is fixed. +_KERNEL_DTYPES = { + "fp8": (torch.bfloat16, torch.float8_e4m3fn), +} + + +def _build_rope_config(scenario: Scenario) -> RopeConfig: + return RopeConfig( + hidden_size=scenario.hidden_size, + num_attention_heads=scenario.num_heads, + rope_scaling={ + "beta_fast": scenario.rope_beta_fast, + "beta_slow": scenario.rope_beta_slow, + "factor": scenario.rope_factor, + "mscale": scenario.rope_mscale, + "mscale_all_dim": scenario.rope_mscale_all_dim, + "original_max_position_embeddings": scenario.rope_original_max_position_embeddings, + "type": scenario.rope_type, + }, + max_position_embeddings=scenario.max_position_embeddings, + rope_theta=scenario.rope_theta, + qk_rope_head_dim=scenario.qk_rope_head_dim, + model_type=scenario.model_type, + ) + + +@pytest.fixture +def cute_dsl_decode_counter(monkeypatch): + """Count CuTe DSL MLA decode dispatches so the test fails on silent + fallback to the TRTLLM backend.""" + from tensorrt_llm._torch.attention_backend.cute_dsl import CuteDslAttention + + # Surface (rather than swallow) kernel errors during the test so a broken + # kernel fails loudly instead of falling back. + monkeypatch.setenv("TLLM_CUTE_DSL_ATTN_DEBUG_FALLBACK", "1") + + original = CuteDslAttention._dispatch_cute_dsl_mla_decode + counter = {"calls": 0} + + def _counting_dispatch(self, *args, **kwargs): + counter["calls"] += 1 + return original(self, *args, **kwargs) + + monkeypatch.setattr(CuteDslAttention, "_dispatch_cute_dsl_mla_decode", _counting_dispatch) + return counter + + +@pytest.mark.parametrize("kernel", list(_KERNEL_DTYPES)) +@pytest.mark.parametrize( + "context_sequence_lengths", _DECODE_CONTEXT_LENGTHS, ids=lambda x: f"ctx_lens={x}" +) +@pytest.mark.parametrize("generation_seq_len_q", [1, 4], ids=lambda x: f"gen_seq_len_q={x}") +def test_cute_dsl_mla_decode( + kernel, context_sequence_lengths, generation_seq_len_q, cute_dsl_decode_counter +): + """Decode-only MLA validation for the Blackwell CuTe DSL kernels.""" + dtype, kv_cache_dtype = _KERNEL_DTYPES[kernel] + + scenario = Scenario(dtype=dtype, kv_cache_dtype=kv_cache_dtype, num_layers=1) + rope_config = _build_rope_config(scenario) + + _run_test_for_backend( + "CUTEDSL", + num_heads=scenario.num_heads, + num_kv_heads=scenario.num_kv_heads, + num_layers=scenario.num_layers, + q_lora_rank=scenario.q_lora_rank, + kv_lora_rank=scenario.kv_lora_rank, + qk_nope_head_dim=scenario.qk_nope_head_dim, + qk_rope_head_dim=scenario.qk_rope_head_dim, + v_head_dim=scenario.v_head_dim, + rope_config=rope_config, + kv_cache_tokens_per_block=scenario.kv_cache_tokens_per_block, + device=torch.device("cuda"), + dtype=scenario.dtype, + kv_cache_dtype=scenario.kv_cache_dtype, + context_sequence_lengths=context_sequence_lengths, + generation_seq_len_q=generation_seq_len_q, + num_generation_steps=_DECODE_NUM_STEPS, + v2_kv_cache=True, + skip_context_assert=True, + ) + + # The decode path must have actually run the CuTe DSL kernel (1 dispatch + # per layer per decode step), not silently fallen back to TRTLLM. + expected = scenario.num_layers * _DECODE_NUM_STEPS + assert cute_dsl_decode_counter["calls"] == expected, ( + f"Expected {expected} CuTe DSL MLA decode dispatches, got " + f"{cute_dsl_decode_counter['calls']} (silent TRTLLM fallback?)" + ) From e274758123ad072b02d3a413fc44104cf2853893 Mon Sep 17 00:00:00 2001 From: haow Date: Fri, 12 Jun 2026 00:39:11 -0700 Subject: [PATCH 2/4] [None][feat] CuteDSL MLA decode: fp8 scaling fix, BF16 support, per-layer paged-KV, multi-layer test - attention_backend/cute_dsl.py: feed the scaled quant_q_buffer (not q.to(fp8)) and fold the fp8 dequant + de-folded log2(e) softmax scale into the bmm1/bmm2 scales; resolve the per-layer KV pool / page-table for layer_idx>0; gate the non-causal fast path off CUSTOM masks and speculative decoding; add a once-per-process ENGAGED log (symmetric to the fallback warning) confirming the decode kernel actually ran. - custom_ops/cute_dsl_custom_ops.py: add BF16 to the decode runner and the fp16 op with dtype dispatch + validation; pass workspace=None for split_kv==1 to avoid a zero-sized-buffer global write. - cute_dsl_kernels/.../mla_decode_fp8.py: allow a widened BF16 attention output. - cute_dsl_kernels/.../mla_decode_fp16.py: accept BF16 input/output in can_implement. - tests/.../test_cute_dsl_mla_decode.py: parametrize over num_layers to exercise the layer_idx>0 paged-KV path (reproduces the real DeepSeek-V3 E2E case). Signed-off-by: haow --- .../_torch/attention_backend/cute_dsl.py | 188 ++++++++++++++---- .../_torch/custom_ops/cute_dsl_custom_ops.py | 43 +++- .../attention/mla/mla_decode_fp16.py | 9 +- .../blackwell/attention/mla/mla_decode_fp8.py | 6 +- .../attention/test_cute_dsl_mla_decode.py | 44 ++-- 5 files changed, 223 insertions(+), 67 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/cute_dsl.py b/tensorrt_llm/_torch/attention_backend/cute_dsl.py index abb9eea3ddf6..8eb9792c1495 100644 --- a/tensorrt_llm/_torch/attention_backend/cute_dsl.py +++ b/tensorrt_llm/_torch/attention_backend/cute_dsl.py @@ -9,12 +9,13 @@ - ``kernel_dtype == torch.float8_e4m3fn`` → ``torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell`` -- ``kernel_dtype == torch.float16`` +- ``kernel_dtype in (torch.float16, torch.bfloat16)`` → ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` ``forward`` picks the kernel dtype directly from runtime state -(``has_fp8_kv_cache`` → FP8; otherwise ``q.dtype == torch.float16`` → FP16; -neither match → TRTLLM fallback). Every other code path (context / +(``has_fp8_kv_cache`` → FP8; otherwise ``q.dtype`` when it is fp16/bf16 and +matches the KV cache dtype → FP16/BF16 via the FP16 op; neither match → +TRTLLM fallback). Every other code path (context / chunked prefill / cached-KV MLA context / non-MLA / mixed batches / unsupported SM) goes through ``super().forward`` unchanged. @@ -27,7 +28,6 @@ """ import math -import os from typing import Optional import torch @@ -36,10 +36,14 @@ from tensorrt_llm.logger import logger from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from .interface import AttentionForwardArgs +from .interface import (AttentionForwardArgs, CustomAttentionMask, + PredefinedAttentionMask) from .trtllm import TrtllmAttention, TrtllmAttentionMetadata -_DEBUG_FALLBACK = os.environ.get("TLLM_CUTE_DSL_ATTN_DEBUG_FALLBACK", "0") == "1" +# log2(e): the CuTe DSL MLA kernels compute softmax via exp2 and fold this +# factor into the scale internally; the TRTLLM-produced mla_bmm1_scale already +# has it folded, so we divide it back out before handing the scale to the kernel. +_LOG2_E = math.log2(math.e) # cutlass / kernel-class imports were used by the now-removed in-eligibility # ``can_implement`` checks — the dispatch now goes through @@ -101,12 +105,12 @@ def _dispatch_cute_dsl_mla_decode( output: torch.Tensor, kernel_dtype: torch.dtype, ) -> torch.Tensor: - """MLA decode dispatch shared by FP8 and FP16 paths. + """MLA decode dispatch shared by FP8 and FP16/BF16 paths. ``kernel_dtype`` is the in/out tensor dtype the chosen CuTe DSL - kernel expects — ``torch.float8_e4m3fn`` for the FP8 kernel, - ``torch.float16`` for the FP16 kernel. The op called is selected - from it. + kernel expects — ``torch.float8_e4m3fn`` for the FP8 kernel, or + ``torch.float16`` / ``torch.bfloat16`` for the FP16 kernel class. The + op called is selected from it. Assumes the MLA module has already (1) built ``q`` as the fused ``[num_tokens, H * (D_latent + D_rope)]`` tensor with RoPE applied @@ -115,12 +119,12 @@ def _dispatch_cute_dsl_mla_decode( """ if kernel_dtype == torch.float8_e4m3fn: op = torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell - elif kernel_dtype == torch.float16: + elif kernel_dtype in (torch.float16, torch.bfloat16): op = torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell else: raise ValueError( f"CuteDslAttention: unsupported kernel_dtype={kernel_dtype}; " - "expected torch.float8_e4m3fn or torch.float16" + "expected torch.float8_e4m3fn, torch.float16, or torch.bfloat16" ) num_tokens = q.shape[0] @@ -131,37 +135,67 @@ def _dispatch_cute_dsl_mla_decode( f"({num_tokens}) divisible by num_generations ({num_seqs})" ) - # Both kernels: L == 512, R == 64, in/out dtype == kernel_dtype. + # Both kernels: L == 512, R == 64. FP16/BF16 keeps in/out dtype equal + # to kernel_dtype; FP8 widens the attention output to bf16 below. d_latent = self.kv_lora_rank d_rope = self.qk_rope_head_dim h = self.num_heads page_size = metadata.tokens_per_block - # q → [H, D, S_q, B]. Cast to kernel dtype if upstream q is in - # something else (e.g. bf16 model + FP8 KV — lossy but defined). - q_kernel = q if q.dtype == kernel_dtype else q.to(kernel_dtype) + # q → [H, D, S_q, B]. + # FP8: use the properly-scaled fp8 query that ``mla_rope_generation`` + # already produced (``quant_q_buffer``) instead of a raw ``q.to(fp8)``. + # The raw cast quantizes the bf16 q with scale 1, which is ~10x less + # accurate than the scaled quantization (whose dequant is folded into + # ``mla_bmm1_scale`` below) and blows past the fp8 test tolerance. + # FP16/BF16: q is already the kernel dtype (or cast losslessly). + if kernel_dtype == torch.float8_e4m3fn and forward_args.quant_q_buffer is not None: + q_kernel = forward_args.quant_q_buffer.view(torch.float8_e4m3fn).view_as(q) + else: + q_kernel = q if q.dtype == kernel_dtype else q.to(kernel_dtype) + # Kernel layout for the q tensors is logical [H, D, S_q, B] with the D + # axis contiguous (stride 1) — see ``mark_layout_dynamic(leading_dim=1)`` + # on the op side and the kernel reference (permute_order=(2,3,1,0)). + # ``q_view`` is [B, S_q, H, D] contiguous, so the permuted *view* + # already has stride-1 on D; calling ``.contiguous()`` here would + # re-lay-out to stride-1 on B and violate the kernel's contract. q_view = q_kernel.view(num_seqs, seq_len_q, h, d_latent + d_rope) - q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0).contiguous() - q_rope = q_view[..., d_latent:].permute(2, 3, 1, 0).contiguous() + q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0) + q_rope = q_view[..., d_latent:].permute(2, 3, 1, 0) # Paged MLA pool view as the kernel's dtype. # NOTE: pool tensor handle and block-table layout for MLA depends # on kv-cache-manager wiring; revisit if ``get_buffers`` exposes a # different layout. kv_pool = metadata.kv_cache_manager.get_buffers(self.layer_idx) + if kernel_dtype in (torch.float16, torch.bfloat16) and kv_pool.dtype != kernel_dtype: + raise ValueError( + f"CuteDslAttention MLA decode {kernel_dtype} fast path requires " + f"matching KV cache dtype, got {kv_pool.dtype}" + ) kv_pool_typed = kv_pool.view(kernel_dtype) c_pool_latent = kv_pool_typed[..., :d_latent] c_pool_rope = kv_pool_typed[..., d_latent : d_latent + d_rope] block_offsets = metadata.kv_cache_block_offsets if block_offsets.dim() == 4: - page_table_layer = block_offsets[self.layer_idx, :, 0, :] + # kv_cache_block_offsets is [num_pools, num_seqs, 2, max_blocks]. + # dim 0 is the *pool* index, not the layer — resolve this layer's + # pool via the layer→pool mapping ([num_layers, 2], col 0 = pool). + # (Indexing dim 0 by ``layer_idx`` is only valid for layer 0 / a + # single pool and goes out of bounds for every other layer.) + pool_mapping = metadata.host_kv_cache_pool_mapping + pool_idx = int(pool_mapping[self.layer_idx, 0]) + page_table_layer = block_offsets[pool_idx, :, 0, :] elif block_offsets.dim() == 3: page_table_layer = block_offsets[:, 0, :] else: page_table_layer = block_offsets - # Kernel: [max_pages, B], leading_dim=0 ⇒ pages contiguous per B. - page_table = page_table_layer.transpose(0, 1).contiguous().to(torch.int32) + # Kernel: [max_pages, B], leading_dim=0 ⇒ pages contiguous (stride 1). + # ``page_table_layer`` is [B, max_pages] with max_pages stride 1, so the + # transposed *view* already has stride 1 on dim 0; ``.contiguous()`` + # would move stride 1 to B and violate the kernel's layout contract. + page_table = page_table_layer.transpose(0, 1).to(torch.int32) cache_seqs = metadata.kv_lens_cuda_runtime.to(torch.int32) @@ -169,16 +203,30 @@ def _dispatch_cute_dsl_mla_decode( workspace = torch.empty(0, dtype=torch.float32, device=q.device) block_split_kvs = torch.empty(0, dtype=torch.int32, device=q.device) + # ``o`` mirrors the q layout: logical [H, D, S_q, B] with D contiguous + # (leading_dim=1). Allocate [B, S_q, H, D] contiguous and permute so the + # D axis keeps stride 1 (a plain contiguous [H, D, S_q, B] would put + # stride 1 on B and fail the op's layout check). + # + # For the FP8 kernel we force the output dtype to bf16 rather than fp8: + # the kernel's epilogue stores through ``o.element_type`` with a + # width-adaptive copy (so bf16 is supported — see the kernel's + # ``can_implement``), and writing the attention result straight to bf16 + # avoids an extra fp8 round-trip on the output. The FP16 kernel class + # keeps its native fp16/bf16 output. + out_kernel_dtype = torch.bfloat16 if kernel_dtype == torch.float8_e4m3fn else kernel_dtype o_kernel = torch.empty( - (h, d_latent, seq_len_q, num_seqs), - dtype=kernel_dtype, + (num_seqs, seq_len_q, h, d_latent), + dtype=out_kernel_dtype, device=q.device, - ) + ).permute(2, 3, 1, 0) + # ``lse`` is logical [H, S_q, B] with H contiguous (leading_dim=0); the + # reference builds it as [B, S_q, H] then permute_order=(2,1,0). lse = torch.empty( - (h, seq_len_q, num_seqs), + (num_seqs, seq_len_q, h), dtype=torch.float32, device=q.device, - ) + ).permute(2, 1, 0) # MLA softmax scale follows the canonical TRT-LLM formula # (see modules/attention.py:1453): @@ -194,9 +242,31 @@ def _dispatch_cute_dsl_mla_decode( qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim softmax_scale = float(1.0 / (math.sqrt(qk_head_dim) * self.q_scaling)) output_scale = 1.0 - - c_latent_kernel = c_pool_latent.unsqueeze(-1).contiguous() - c_rope_kernel = c_pool_rope.unsqueeze(-1).contiguous() + if (kernel_dtype == torch.float8_e4m3fn + and forward_args.mla_bmm1_scale is not None + and forward_args.mla_bmm2_scale is not None): + # ``mla_rope_generation`` folds the fp8 dequant (1/(q_scale*k_scale)) + # and the softmax scale into ``mla_bmm1_scale[1]`` (the FMHA + # scaleBmm1, see attentionOp.cpp bmm1_scale_offset=1), and the + # output/V dequant into ``mla_bmm2_scale[0]``. The CuTe DSL kernel + # reads q/KV as raw fp8 and applies these two scalar scales, so + # reusing them dequantizes identically to the TRTLLM FP8 MLA path. + # mla_bmm1_scale[1] is the TRTLLM FMHA scaleBmm1, which has log2(e) + # PRE-folded (its softmax uses exp2). The CuTe DSL kernel folds + # log2(e) AGAIN internally (softmax_scale_log2 = softmax_scale * + # LOG2_E), so pass the de-folded value to avoid applying it twice. + softmax_scale = float(forward_args.mla_bmm1_scale[1].item()) / _LOG2_E + output_scale = float(forward_args.mla_bmm2_scale[0].item()) + + # Paged cache layout for the kernel is logical [page_size, D, num_pages] + # with the D axis contiguous (leading_dim=1) — see the kernel reference + # (cache shape (num_pages, page_size, D), permute_order=(1,2,0)). + # ``get_buffers`` returns the pool as + # [num_pages, kv_factor=1, page_size, num_kv_heads=1, head_dim]; drop the + # two singleton axes and permute so D keeps stride 1. The latent/rope + # slices already carry stride 1 on their (last) D axis, so no copy. + c_latent_kernel = c_pool_latent.squeeze(3).squeeze(1).permute(1, 2, 0) + c_rope_kernel = c_pool_rope.squeeze(3).squeeze(1).permute(1, 2, 0) op( q_latent, @@ -245,14 +315,38 @@ def forward( # CuTe DSL fast path. Everything else falls through to TRTLLM: # - non-MLA attention → TRTLLM # - prefill / mixed batch (has context) → TRTLLM - # - decode-only MLA → try CuteDSL FP8/FP16, + # - decode-only MLA → try CuteDSL FP8/FP16/BF16, # TRTLLM on failure is_decode_only_mla = ( self.is_mla_enable and metadata.num_contexts == 0 and metadata.num_generations > 0 ) + # The CuTe DSL MLA decode kernel is non-causal: every query token + # attends to the whole cached KV with no per-query mask. That is correct + # for the plain decode case (the MLA module passes the default CAUSAL / + # FULL predefined mask with no mask tensor), but it cannot honor any + # per-query causal masking. Fall back to TRTLLM whenever masking the + # kernel cannot express is required: + # - an explicit CUSTOM mask or mask tensor, or + # - speculative decoding (MTP / Eagle / tree) — signalled by + # ``metadata.use_spec_decoding``; those steps feed >1 query token per + # request and require a causal/tree mask the non-causal kernel lacks. + attention_mask = ( + forward_args.attention_mask + if forward_args is not None else PredefinedAttentionMask.CAUSAL + ) + attention_mask_data = ( + forward_args.attention_mask_data if forward_args is not None else None + ) + mask_supported = ( + attention_mask != CustomAttentionMask.CUSTOM + and attention_mask_data is None + and not getattr(metadata, "use_spec_decoding", False) + ) + if ( is_decode_only_mla + and mask_supported and self._cute_dsl_mla_decode_common_preconditions(metadata, forward_args) and q.shape[0] % metadata.num_generations == 0 ): @@ -262,8 +356,9 @@ def forward( # try/except below. if getattr(self, "has_fp8_kv_cache", False): kernel_dtype = torch.float8_e4m3fn - elif q.dtype == torch.float16: - kernel_dtype = torch.float16 + elif q.dtype in (torch.float16, torch.bfloat16): + kv_pool_dtype = metadata.kv_cache_manager.get_buffers(self.layer_idx).dtype + kernel_dtype = q.dtype if kv_pool_dtype == q.dtype else None else: kernel_dtype = None @@ -272,17 +367,28 @@ def forward( output = q.new_empty( (q.shape[0], self.num_heads * self.kv_lora_rank), dtype=q.dtype ) - return self._dispatch_cute_dsl_mla_decode( + result = self._dispatch_cute_dsl_mla_decode( q, metadata, forward_args, output, kernel_dtype ) + # Positive confirmation (once per process) that the CuteDSL + # MLA decode kernel was actually engaged — symmetric to the + # fallback warning below so a real kernel call can never be + # confused with a silent TRTLLM fallback. + logger.warning_once( + "CuteDslAttention: MLA decode fast path ENGAGED " + "(kernel_dtype=%s); CuteDSL kernel called." % kernel_dtype, + key="cute_dsl_mla_decode_engaged", + ) + return result except Exception as exc: # noqa: BLE001 - if _DEBUG_FALLBACK: - logger.warning( - "CuteDslAttention: MLA decode fast path " - "(kernel_dtype=%s) failed (%s); falling back " - "to TRTLLM backend.", - kernel_dtype, - exc, - ) + # Surface the fallback once per process (keyed dedup keeps + # it from flooding the log on every decode step) so a broken + # kernel can never silently masquerade as a working one. + logger.warning_once( + "CuteDslAttention: MLA decode fast path " + "(kernel_dtype=%s) failed (%s); falling back " + "to TRTLLM backend." % (kernel_dtype, exc), + key="cute_dsl_mla_decode_fallback", + ) return super().forward(q, k, v, metadata, forward_args=forward_args, **kwargs) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index fa8ea384c448..3bf98374956e 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7335,7 +7335,7 @@ def _( # Used by the CUTEDSL attention backend (see attention_backend/cute_dsl.py). # # One generic Runner ``CuteDSLNVMlaDecodeBlackwellRunner`` services both - # FP8 and FP16 paths — only the cutlass ``in_dtype`` is passed at + # FP8 and FP16/BF16 paths — only the cutlass ``in_dtype`` is passed at # construction; the kernel class is derived from it via # ``CuteDSLNVMlaDecodeBlackwellRunner._KERNEL_CLASS_BY_DTYPE``. Each # dtype still has its own ``@torch.library.custom_op`` (distinct op @@ -7348,6 +7348,7 @@ def _( # # torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell # → CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.Float16) + # or CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.BFloat16) # (→ BlackwellMultiHeadLatentAttentionForwardFP16) # ========================================================================= @@ -7361,13 +7362,15 @@ def _( class CuteDSLNVMlaDecodeBlackwellRunner(TunableRunner): """Generic TunableRunner for the Blackwell CuTe DSL MLA decode kernels. - Works for both FP8 and FP16 — pass the cutlass input dtype at + Works for FP8, FP16, and BF16 — pass the cutlass input dtype at construction; the kernel class is derived from it: CuteDSLNVMlaDecodeBlackwellRunner( in_dtype=cutlass.Float8E4M3FN, ...) # → ...ForwardFP8 CuteDSLNVMlaDecodeBlackwellRunner( in_dtype=cutlass.Float16, ...) # → ...ForwardFP16 + CuteDSLNVMlaDecodeBlackwellRunner( + in_dtype=cutlass.BFloat16, ...) # → ...ForwardFP16 ``get_valid_tactics`` returns the tiler shapes as tactics (``(mma_qk_tiler_mn, mma_pv_tiler_mn)`` tuples), filtered by the @@ -7387,6 +7390,7 @@ class CuteDSLNVMlaDecodeBlackwellRunner(TunableRunner): _KERNEL_CLASS_BY_DTYPE = { cutlass.Float8E4M3FN: BlackwellMultiHeadLatentAttentionForwardFP8, cutlass.Float16: BlackwellMultiHeadLatentAttentionForwardFP16, + cutlass.BFloat16: BlackwellMultiHeadLatentAttentionForwardFP16, } def __init__( @@ -7542,8 +7546,6 @@ def forward( self.is_persistent, self.is_var_seq, self.is_var_split_kv, - num_heads=self.num_heads, - seq_len_q=self.seq_len_q, ) q_latent_ct = cute.runtime.from_dlpack( @@ -7563,8 +7565,15 @@ def forward( o, assumed_align=16).mark_layout_dynamic(leading_dim=1) lse_ct = cute.runtime.from_dlpack( lse, assumed_align=16).mark_layout_dynamic(leading_dim=0) - workspace_ct = cute.runtime.from_dlpack( + # An empty workspace means split_kv == 1: the kernel's + # initialize_workspace builds the acc_o/acc_lse accumulators iff + # ``workspace is not None`` (regardless of split_kv), so a + # non-None but zero-sized workspace makes it write the partials + # into a 0-byte buffer (illegal global write). Pass None so the + # split_kv kernel writes the final result straight into ``o``. + workspace_ct = (cute.runtime.from_dlpack( workspace, assumed_align=16).mark_layout_dynamic() + if workspace.numel() > 0 else None) cache_seqs_ct = cute.runtime.from_dlpack( cache_seqs, assumed_align=16).mark_layout_dynamic() block_split_kvs_ct = (cute.runtime.from_dlpack( @@ -7601,7 +7610,7 @@ def forward( page_table, o, lse, - workspace, + workspace if workspace.numel() > 0 else None, split_kv, cache_seqs, block_split_kvs if self.is_var_split_kv else None, @@ -7725,7 +7734,7 @@ def cute_dsl_mla_decode_fp16_blackwell( softmax_scale: float, output_scale: float, ) -> None: - """CuTe DSL FP16 MLA decode (Blackwell SM100/SM103). + """CuTe DSL FP16/BF16 MLA decode (Blackwell SM100/SM103). ``o``, ``lse``, ``workspace`` are mutated in place. Tensor layouts: see ``BlackwellMultiHeadLatentAttentionForwardFP16``. @@ -7735,8 +7744,26 @@ def cute_dsl_mla_decode_fp16_blackwell( f"trtllm::cute_dsl_mla_decode_fp16_blackwell requires SM 100 " f"or SM 103, got SM {sm_version}") + if q_latent.dtype == torch.float16: + in_dtype = cutlass.Float16 + elif q_latent.dtype == torch.bfloat16: + in_dtype = cutlass.BFloat16 + else: + raise ValueError( + "trtllm::cute_dsl_mla_decode_fp16_blackwell supports " + "torch.float16 or torch.bfloat16 inputs, got " + f"{q_latent.dtype}") + if not ( + q_rope.dtype == c_latent.dtype == c_rope.dtype == o.dtype + == q_latent.dtype): + raise ValueError( + "trtllm::cute_dsl_mla_decode_fp16_blackwell requires q, KV, " + f"and output dtypes to match; got q_latent={q_latent.dtype}, " + f"q_rope={q_rope.dtype}, c_latent={c_latent.dtype}, " + f"c_rope={c_rope.dtype}, o={o.dtype}") + runner = CuteDSLNVMlaDecodeBlackwellRunner( - in_dtype=cutlass.Float16, + in_dtype=in_dtype, num_heads=num_heads, seq_len_q=seq_len_q, page_size=page_size, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py index f1cc55fb629b..a8124e93af00 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py @@ -90,7 +90,8 @@ - Latent dimension: 512 - RoPE dimension: 64 - Number of heads: 128 -- Data types: Float16 (input), Float16 (output), Float32 (accumulation and LSE) +- Data types: Float16/BFloat16 (input), Float16/BFloat16 (output), + Float32 (accumulation and LSE) It utilizes page table storage for the KV cache and enables both variable-length KV cache sequences and variable split KV processing with persistent scheduling. @@ -110,7 +111,7 @@ Constraints for this example: * Data type requirements: - - Input/output: Float16 + - Input/output: Float16 or BFloat16 - Accumulation and LSE: Float32 * Fixed architecture parameters: - Number of attention heads: 128 @@ -3340,9 +3341,9 @@ def can_implement( """ if L != 512 or R != 64: return False - if in_dtype not in [cutlass.Float16]: + if in_dtype not in [cutlass.Float16, cutlass.BFloat16]: return False - if out_dtype not in [cutlass.Float16]: + if out_dtype not in [cutlass.Float16, cutlass.BFloat16]: return False if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: return False diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py index 560284cf17cf..3604a83d669d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py @@ -3289,7 +3289,11 @@ def can_implement( return False if in_dtype not in [cutlass.Float8E4M3FN]: return False - if out_dtype not in [cutlass.Float8E4M3FN]: + # The epilogue stores through ``self.o_dtype`` (= o.element_type) with a + # width-adaptive autovec copy, so a bf16 output is supported alongside + # fp8 (matches flashinfer's mla_decode_fp8 can_implement). Inputs stay + # fp8; only the attention output dtype is widened. + if out_dtype not in [cutlass.Float8E4M3FN, cutlass.BFloat16]: return False if acc_dtype != cutlass.Float32 or lse_dtype != cutlass.Float32: return False diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py index 243c28ed989c..e281e3f06da8 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py @@ -20,7 +20,7 @@ (``tensorrt_llm/_torch/attention_backend/cute_dsl.py``): - FP8 path → ``torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell`` -- FP16 path → ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` +- FP16/BF16 path → ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` Only the generation (decode) steps are asserted for numerical correctness. The context phase runs solely to populate the paged KV cache and to build the @@ -54,6 +54,18 @@ ] _DECODE_NUM_STEPS = 4 +# Multi-layer is the structural difference between this single-step test and the +# real DeepSeek-V3 E2E run (61 MLA layers). With ``num_layers == 1`` the dispatch +# only ever sees ``layer_idx == 0``, so the per-layer paged-KV resolution in +# ``CuteDslAttention._dispatch_cute_dsl_mla_decode`` (the +# ``host_kv_cache_pool_mapping[layer_idx]`` / per-layer ``get_buffers`` / +# block-offset path) is never exercised. The E2E run produces correct output on +# the first generated token (which comes from the TRTLLM prefill) and then +# degenerates on every subsequent CuteDSL decode step — consistent with the +# decode kernel reading the wrong blocks for ``layer_idx > 0``. Parametrize over +# >1 layers so the unit test reproduces that real case. +_DECODE_NUM_LAYERS = [1, 2] + def _is_blackwell_sm100() -> bool: return torch.cuda.is_available() and torch.cuda.get_device_capability() in ((10, 0), (10, 3)) @@ -70,12 +82,14 @@ def _is_blackwell_sm100() -> bool: # kernel name -> (activation dtype, kv cache dtype) # -# NOTE: only the FP8 decode kernel is currently validated here. The FP16 path +# NOTE: the float16 instance of the FP16 decode op # (``cute_dsl_mla_decode_fp16_blackwell`` with dtype=float16 / fp16 KV cache) -# aborts the process (SIGABRT) on SM100 in this environment, so it is excluded -# until that crash is root-caused. Re-add "fp16" below once it is fixed. +# aborts the process (SIGABRT) on SM100 in this environment, so that exact +# dtype is excluded until the crash is root-caused. The bf16 instance uses the +# same op and is covered below for DeepSeek-V3 bf16 runs. _KERNEL_DTYPES = { "fp8": (torch.bfloat16, torch.float8_e4m3fn), + "bf16": (torch.bfloat16, torch.bfloat16), } @@ -101,20 +115,23 @@ def _build_rope_config(scenario: Scenario) -> RopeConfig: @pytest.fixture def cute_dsl_decode_counter(monkeypatch): - """Count CuTe DSL MLA decode dispatches so the test fails on silent - fallback to the TRTLLM backend.""" + """Count *successful* CuTe DSL MLA decode dispatches so the test fails on + silent fallback to the TRTLLM backend. + + The increment happens only after the real dispatch returns: if the kernel + raises, ``CuteDslAttention.forward`` catches it and falls back to TRTLLM, + the counter does not advance, and the per-test assertion on the expected + dispatch count fails loudly instead of the broken kernel masquerading as a + working one.""" from tensorrt_llm._torch.attention_backend.cute_dsl import CuteDslAttention - # Surface (rather than swallow) kernel errors during the test so a broken - # kernel fails loudly instead of falling back. - monkeypatch.setenv("TLLM_CUTE_DSL_ATTN_DEBUG_FALLBACK", "1") - original = CuteDslAttention._dispatch_cute_dsl_mla_decode counter = {"calls": 0} def _counting_dispatch(self, *args, **kwargs): + result = original(self, *args, **kwargs) counter["calls"] += 1 - return original(self, *args, **kwargs) + return result monkeypatch.setattr(CuteDslAttention, "_dispatch_cute_dsl_mla_decode", _counting_dispatch) return counter @@ -125,13 +142,14 @@ def _counting_dispatch(self, *args, **kwargs): "context_sequence_lengths", _DECODE_CONTEXT_LENGTHS, ids=lambda x: f"ctx_lens={x}" ) @pytest.mark.parametrize("generation_seq_len_q", [1, 4], ids=lambda x: f"gen_seq_len_q={x}") +@pytest.mark.parametrize("num_layers", _DECODE_NUM_LAYERS, ids=lambda x: f"num_layers={x}") def test_cute_dsl_mla_decode( - kernel, context_sequence_lengths, generation_seq_len_q, cute_dsl_decode_counter + kernel, context_sequence_lengths, generation_seq_len_q, num_layers, cute_dsl_decode_counter ): """Decode-only MLA validation for the Blackwell CuTe DSL kernels.""" dtype, kv_cache_dtype = _KERNEL_DTYPES[kernel] - scenario = Scenario(dtype=dtype, kv_cache_dtype=kv_cache_dtype, num_layers=1) + scenario = Scenario(dtype=dtype, kv_cache_dtype=kv_cache_dtype, num_layers=num_layers) rope_config = _build_rope_config(scenario) _run_test_for_backend( From 214ebd6775239748da0cae1d9a56c5fa78b80a14 Mon Sep 17 00:00:00 2001 From: haow Date: Mon, 15 Jun 2026 22:57:56 -0700 Subject: [PATCH 3/4] [None][fix] CuteDSL MLA decode: v1 paged-KV layout + CUDA-graph-safe fp8 scale Normalize the per-layer paged-KV pool view for the v1 KVCacheManager (interleaved single-pool layout) to a packed combined-slot view and fold the layer offset into the page table, so the CuTe DSL paged TMA addresses the correct layer's memory. No-op for the v2 manager. Cache the static fp8 dequant scales on the per-layer backend instance and read them eagerly during warmup, avoiding an illegal .item() device->host sync under CUDA graph capture. Add a long decode-only test that crosses paged-KV block boundaries mid-decode (short prompt, 64 steps), parametrized over v1/v2 KV cache managers, reproducing the DeepSeek-V3 E2E degeneration path. Signed-off-by: haow --- .../_torch/attention_backend/cute_dsl.py | 113 +++++++++++++++++- .../attention/test_cute_dsl_mla_decode.py | 63 ++++++++++ 2 files changed, 174 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/cute_dsl.py b/tensorrt_llm/_torch/attention_backend/cute_dsl.py index 8eb9792c1495..d9c63a3500cd 100644 --- a/tensorrt_llm/_torch/attention_backend/cute_dsl.py +++ b/tensorrt_llm/_torch/attention_backend/cute_dsl.py @@ -28,6 +28,7 @@ """ import math +import os from typing import Optional import torch @@ -173,6 +174,37 @@ def _dispatch_cute_dsl_mla_decode( f"CuteDslAttention MLA decode {kernel_dtype} fast path requires " f"matching KV cache dtype, got {kv_pool.dtype}" ) + + # Paged-pool layout normalization (v1 vs v2 KV cache manager). + # ``get_buffers`` returns a per-layer view of shape + # [num_blocks, kv_factor, page_size, num_kv_heads, head_dim]. The v2 + # manager gives each layer its own densely-packed pool, so the block + # (dim-0) stride equals the packed block size. The v1 manager stores all + # layers interleaved in ONE pool, so the per-layer view's block stride is + # ``layers_in_pool * packed_block`` (non-packed) with a per-layer + # ``storage_offset``. The CuTe DSL paged TMA addresses pages with the + # packed block stride and does not honor a non-packed one, so for v1 it + # would read the wrong (interleaved) layer's memory. Re-expose the + # underlying pool as a packed ``[num_blocks * layers_in_pool, ...]`` view + # (copy-free) and fold the layer offset into the page table below: + # combined_block = block * layers_in_pool + layer_in_pool + # For v2 ``layers_in_pool == 1`` and this is a no-op. + packed_block = 1 + for s in kv_pool.shape[1:]: + packed_block *= s + block_stride = kv_pool.stride(0) + layers_in_pool = block_stride // packed_block if packed_block else 1 + layer_in_pool = 0 + if layers_in_pool > 1 and block_stride == layers_in_pool * packed_block: + layer_in_pool = kv_pool.storage_offset() // packed_block + kv_pool = kv_pool.as_strided( + (kv_pool.shape[0] * layers_in_pool, ) + tuple(kv_pool.shape[1:]), + (packed_block, ) + tuple(kv_pool.stride()[1:]), + 0, + ) + else: + layers_in_pool = 1 + kv_pool_typed = kv_pool.view(kernel_dtype) c_pool_latent = kv_pool_typed[..., :d_latent] c_pool_rope = kv_pool_typed[..., d_latent : d_latent + d_rope] @@ -197,8 +229,45 @@ def _dispatch_cute_dsl_mla_decode( # would move stride 1 to B and violate the kernel's layout contract. page_table = page_table_layer.transpose(0, 1).to(torch.int32) + # Fold the interleaved-pool layer offset into the page table so it + # indexes the packed combined-slot view built above. The v1 block + # offsets are already in combined-slot units (block * layers_in_pool), + # so only the per-layer offset has to be added. No-op for v2 + # (layers_in_pool == 1, layer_in_pool == 0). + if layers_in_pool > 1: + page_table = page_table + layer_in_pool + cache_seqs = metadata.kv_lens_cuda_runtime.to(torch.int32) + if os.environ.get("TLLM_CUTE_DSL_DUMP"): + print( + "[CUTEDSL_DUMP] layer=%d kv_pool.shape=%s kv_pool.stride=%s " + "contiguous=%s block_offsets.shape=%s page_table.shape=%s " + "page_table=%s cache_seqs=%s" % ( + self.layer_idx, + tuple(kv_pool.shape), + tuple(kv_pool.stride()), + kv_pool.is_contiguous(), + tuple(block_offsets.shape), + tuple(page_table.shape), + page_table.t().tolist() if page_table.numel() < 64 else "(big)", + cache_seqs[:8].tolist(), + ), + flush=True, + ) + _pm = getattr(metadata, "host_kv_cache_pool_mapping", None) + print( + "[CUTEDSL_DUMP2] layer=%d pool_mapping=%s " + "raw_block_offsets[...,0,:]=%s" % ( + self.layer_idx, + _pm.tolist() if _pm is not None else None, + block_offsets[..., 0, :].tolist() + if block_offsets.dim() == 4 and block_offsets.numel() < 128 + else "(big/other)", + ), + flush=True, + ) + split_kv = 1 workspace = torch.empty(0, dtype=torch.float32, device=q.device) block_split_kvs = torch.empty(0, dtype=torch.int32, device=q.device) @@ -255,8 +324,48 @@ def _dispatch_cute_dsl_mla_decode( # PRE-folded (its softmax uses exp2). The CuTe DSL kernel folds # log2(e) AGAIN internally (softmax_scale_log2 = softmax_scale * # LOG2_E), so pass the de-folded value to avoid applying it twice. - softmax_scale = float(forward_args.mla_bmm1_scale[1].item()) / _LOG2_E - output_scale = float(forward_args.mla_bmm2_scale[0].item()) + # + # CUDA-graph safety: ``.item()`` is a device->host copy + sync, which + # is illegal while a CUDA graph stream is capturing. These fp8 + # dequant scales are STATIC per layer — ``mla_rope_generation`` + # derives them from static quantization state (kv_scale_quant_orig, + # q_scaling, quant_mode), not from per-step activations — so we read + # them once (eagerly) and cache the host float on this per-layer + # backend instance. The generation CUDA-graph capture runs + # ``WARMUP_STEPS`` eager decode passes over every layer first + # (cuda_graph_runner.capture), so the cache is always populated + # before capture and the ``.item()`` never runs under capture. + cached = getattr(self, "_cute_dsl_fp8_scale", None) + if cached is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CuteDslAttention: fp8 MLA decode scale not cached for " + f"layer {self.layer_idx} before CUDA graph capture; " + "reading it via .item() during capture is illegal. The " + "eager warmup should have populated the cache.") + softmax_scale = float( + forward_args.mla_bmm1_scale[1].item()) / _LOG2_E + output_scale = float(forward_args.mla_bmm2_scale[0].item()) + self._cute_dsl_fp8_scale = (softmax_scale, output_scale) + else: + softmax_scale, output_scale = cached + # Optional staticness check: re-read the live scale eagerly and + # warn if it ever drifts from the cached value (would mean the + # static-scale assumption — and thus the cache — is wrong). + if (os.environ.get("TLLM_CUTE_DSL_SCALE_DUMP") + and not torch.cuda.is_current_stream_capturing()): + live_sm = float( + forward_args.mla_bmm1_scale[1].item()) / _LOG2_E + live_os = float(forward_args.mla_bmm2_scale[0].item()) + print( + "[CUTEDSL_SCALE] layer=%d cached=(%.8f,%.8f) " + "live=(%.8f,%.8f) drift=%s" % ( + self.layer_idx, softmax_scale, output_scale, + live_sm, live_os, + abs(live_sm - softmax_scale) > 1e-6 + or abs(live_os - output_scale) > 1e-6), + flush=True, + ) # Paged cache layout for the kernel is logical [page_size, D, num_pages] # with the D axis contiguous (leading_dim=1) — see the kernel reference diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py index e281e3f06da8..8e21cafaf778 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py @@ -181,3 +181,66 @@ def test_cute_dsl_mla_decode( f"Expected {expected} CuTe DSL MLA decode dispatches, got " f"{cute_dsl_decode_counter['calls']} (silent TRTLLM fallback?)" ) + + +# E2E-reproduction case: a SHORT prompt decoded for MANY steps. The real +# DeepSeek-V3 run feeds a ~6-token prompt and generates ~64 tokens; its output +# is correct on the first (prefill) token and then degenerates on every CuteDSL +# decode step. The parametrized ``test_cute_dsl_mla_decode`` above keeps decode +# to 4 steps and so, even with a long context, never allocates a fresh paged-KV +# block *during* generation. Here the KV length grows from a short context +# across the page_size==32 block boundaries (at lengths 32 and 64) *mid-decode*, +# allocating new blocks and extending the per-request page_table on the fly — +# the exact paged-KV path the short-decode test never exercises. fp8 + +# seq_len_q==1 only (the configuration that degenerates E2E). +_LONG_DECODE_CONTEXT_LENGTHS = [5, 8, 3, 11] +_LONG_DECODE_NUM_STEPS = 64 + + +@pytest.mark.parametrize("kernel", list(_KERNEL_DTYPES)) +@pytest.mark.parametrize("num_layers", _DECODE_NUM_LAYERS, ids=lambda x: f"num_layers={x}") +@pytest.mark.parametrize("v2_kv_cache", [True, False], ids=lambda x: f"v2_kv_cache={x}") +def test_cute_dsl_mla_decode_long_decode( + v2_kv_cache, num_layers, kernel, cute_dsl_decode_counter +): + """Long decode-only MLA run that crosses paged-KV block boundaries mid-decode. + + Reproduction for the DeepSeek-V3 E2E degeneration: short prompt, long + generation, ``seq_len_q == 1``. The real run uses the v1 ``KVCacheManager`` + (``use_kv_cache_manager_v2=False``); the rest of this file only exercised the + v2 manager, so ``v2_kv_cache`` is parametrized here to cover the v1 paged-KV + block-offset layout that the dispatch resolves in + ``_dispatch_cute_dsl_mla_decode``. + """ + dtype, kv_cache_dtype = _KERNEL_DTYPES[kernel] + + scenario = Scenario(dtype=dtype, kv_cache_dtype=kv_cache_dtype, num_layers=num_layers) + rope_config = _build_rope_config(scenario) + + _run_test_for_backend( + "CUTEDSL", + num_heads=scenario.num_heads, + num_kv_heads=scenario.num_kv_heads, + num_layers=scenario.num_layers, + q_lora_rank=scenario.q_lora_rank, + kv_lora_rank=scenario.kv_lora_rank, + qk_nope_head_dim=scenario.qk_nope_head_dim, + qk_rope_head_dim=scenario.qk_rope_head_dim, + v_head_dim=scenario.v_head_dim, + rope_config=rope_config, + kv_cache_tokens_per_block=scenario.kv_cache_tokens_per_block, + device=torch.device("cuda"), + dtype=scenario.dtype, + kv_cache_dtype=scenario.kv_cache_dtype, + context_sequence_lengths=_LONG_DECODE_CONTEXT_LENGTHS, + generation_seq_len_q=1, + num_generation_steps=_LONG_DECODE_NUM_STEPS, + v2_kv_cache=v2_kv_cache, + skip_context_assert=True, + ) + + expected = scenario.num_layers * _LONG_DECODE_NUM_STEPS + assert cute_dsl_decode_counter["calls"] == expected, ( + f"Expected {expected} CuTe DSL MLA decode dispatches, got " + f"{cute_dsl_decode_counter['calls']} (silent TRTLLM fallback?)" + ) From b893d837ecdec48f9098d7e55cdeb0ebf2c44fb8 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sat, 13 Jun 2026 08:29:55 +0000 Subject: [PATCH 4/4] [TRTLLM-12807][feat] Wrap CuteDSL MLA decode as FMHA lib Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention_backend/__init__.py | 2 - .../_torch/attention_backend/cute_dsl.py | 503 ------------------ .../_torch/attention_backend/fmha/__init__.py | 2 + .../_torch/attention_backend/fmha/cute_dsl.py | 465 ++++++++++++++++ .../_torch/attention_backend/fmha/registry.py | 2 + .../_torch/attention_backend/utils.py | 3 - .../_torch/custom_ops/cute_dsl_custom_ops.py | 70 ++- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 16 +- .../attention/test_cute_dsl_mla_decode.py | 48 +- 9 files changed, 544 insertions(+), 567 deletions(-) delete mode 100644 tensorrt_llm/_torch/attention_backend/cute_dsl.py create mode 100644 tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py diff --git a/tensorrt_llm/_torch/attention_backend/__init__.py b/tensorrt_llm/_torch/attention_backend/__init__.py index 0895418af85d..ae0d8b87cd85 100644 --- a/tensorrt_llm/_torch/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/__init__.py @@ -1,5 +1,4 @@ from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE -from .cute_dsl import CuteDslAttention from .interface import AttentionBackend, AttentionForwardArgs, AttentionMetadata from .sparse import get_sparse_attn_kv_cache_manager from .trtllm import AttentionInputType, TrtllmAttention, TrtllmAttentionMetadata @@ -10,7 +9,6 @@ "AttentionBackend", "AttentionForwardArgs", "AttentionInputType", - "CuteDslAttention", "TrtllmAttention", "TrtllmAttentionMetadata", "VanillaAttention", diff --git a/tensorrt_llm/_torch/attention_backend/cute_dsl.py b/tensorrt_llm/_torch/attention_backend/cute_dsl.py deleted file mode 100644 index d9c63a3500cd..000000000000 --- a/tensorrt_llm/_torch/attention_backend/cute_dsl.py +++ /dev/null @@ -1,503 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""CuTeDSL attention backend. - -Subclasses ``TrtllmAttention`` and overrides ``forward`` to dispatch the -MLA decode-only path into one of two Blackwell CuTe DSL kernels via a -single shared dispatcher (``_dispatch_cute_dsl_mla_decode``) parameterised -by the kernel's input dtype: - -- ``kernel_dtype == torch.float8_e4m3fn`` - → ``torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell`` -- ``kernel_dtype in (torch.float16, torch.bfloat16)`` - → ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` - -``forward`` picks the kernel dtype directly from runtime state -(``has_fp8_kv_cache`` → FP8; otherwise ``q.dtype`` when it is fp16/bf16 and -matches the KV cache dtype → FP16/BF16 via the FP16 op; neither match → -TRTLLM fallback). Every other code path (context / -chunked prefill / cached-KV MLA context / non-MLA / mixed batches / -unsupported SM) goes through ``super().forward`` unchanged. - -Subclassing ``TrtllmAttention`` matters: ``modules/attention.py`` selects -the MLA chunked-prefill / cached-context fast paths via -``isinstance(self.mha, TrtllmAttention)``, so the CUTEDSL backend must -satisfy that check or those paths would silently fall back to the slow -default context path. This mirrors the MoE CuTeDSL integration -(``CuteDslFusedMoE(CutlassFusedMoE)``). -""" - -import math -import os -from typing import Optional - -import torch - -from tensorrt_llm._utils import get_sm_version -from tensorrt_llm.logger import logger - -from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from .interface import (AttentionForwardArgs, CustomAttentionMask, - PredefinedAttentionMask) -from .trtllm import TrtllmAttention, TrtllmAttentionMetadata - -# log2(e): the CuTe DSL MLA kernels compute softmax via exp2 and fold this -# factor into the scale internally; the TRTLLM-produced mla_bmm1_scale already -# has it folded, so we divide it back out before handing the scale to the kernel. -_LOG2_E = math.log2(math.e) - -# cutlass / kernel-class imports were used by the now-removed in-eligibility -# ``can_implement`` checks — the dispatch now goes through -# ``torch.ops.trtllm.cute_dsl_mla_decode_*_blackwell`` and the kernel-level -# Runner runs ``can_implement`` itself. ``IS_CUTLASS_DSL_AVAILABLE`` is -# still consulted in the eligibility preconditions below to short-circuit -# environments without the cutlass package. - - -class CuteDslAttention(TrtllmAttention): - """CuteDSL attention backend. - - Inherits the full ``TrtllmAttention`` machinery (metadata, KV cache, - quant flags, RoPE buffers, MLA helpers) and overrides ``forward`` to - intercept the MLA decode-only batch. Anything that fails eligibility - falls through to ``TrtllmAttention.forward``. - """ - - Metadata = TrtllmAttentionMetadata - - # ------------------------------------------------------------------ - # Shared preconditions (everything except dtype-specific gates). - # ------------------------------------------------------------------ - def _cute_dsl_mla_decode_common_preconditions( - self, - metadata: TrtllmAttentionMetadata, - forward_args: Optional[AttentionForwardArgs], - ) -> bool: - """Checks that are identical for the FP8 and FP16 paths. - - Note: phase routing (is_mla_enable / num_contexts / num_generations) - is already handled in ``forward`` before either dtype's eligibility - method runs, so it isn't repeated here. - - The dtype-specific eligibility methods own everything else - (KV-cache dtype, ``q.dtype``, kernel-level ``can_implement``). - """ - if not IS_CUTLASS_DSL_AVAILABLE: - return False - if get_sm_version() not in (100, 103): - return False - if self.predicted_tokens_per_seq is None or not (1 <= self.predicted_tokens_per_seq <= 4): - return False - if metadata.kv_cache_block_offsets is None: - return False - if forward_args is None or forward_args.latent_cache is None: - return False - return True - - # ================================================================== - # MLA decode dispatch (FP8 / FP16) - # ================================================================== - - def _dispatch_cute_dsl_mla_decode( - self, - q: torch.Tensor, - metadata: TrtllmAttentionMetadata, - forward_args: AttentionForwardArgs, - output: torch.Tensor, - kernel_dtype: torch.dtype, - ) -> torch.Tensor: - """MLA decode dispatch shared by FP8 and FP16/BF16 paths. - - ``kernel_dtype`` is the in/out tensor dtype the chosen CuTe DSL - kernel expects — ``torch.float8_e4m3fn`` for the FP8 kernel, or - ``torch.float16`` / ``torch.bfloat16`` for the FP16 kernel class. The - op called is selected from it. - - Assumes the MLA module has already (1) built ``q`` as the fused - ``[num_tokens, H * (D_latent + D_rope)]`` tensor with RoPE applied - to the rope half, and (2) appended the new token to the paged - latent cache. - """ - if kernel_dtype == torch.float8_e4m3fn: - op = torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell - elif kernel_dtype in (torch.float16, torch.bfloat16): - op = torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell - else: - raise ValueError( - f"CuteDslAttention: unsupported kernel_dtype={kernel_dtype}; " - "expected torch.float8_e4m3fn, torch.float16, or torch.bfloat16" - ) - - num_tokens = q.shape[0] - num_seqs = metadata.num_generations - seq_len_q = num_tokens // num_seqs - assert seq_len_q * num_seqs == num_tokens, ( - f"CuteDslAttention MLA decode expects num_tokens " - f"({num_tokens}) divisible by num_generations ({num_seqs})" - ) - - # Both kernels: L == 512, R == 64. FP16/BF16 keeps in/out dtype equal - # to kernel_dtype; FP8 widens the attention output to bf16 below. - d_latent = self.kv_lora_rank - d_rope = self.qk_rope_head_dim - h = self.num_heads - page_size = metadata.tokens_per_block - - # q → [H, D, S_q, B]. - # FP8: use the properly-scaled fp8 query that ``mla_rope_generation`` - # already produced (``quant_q_buffer``) instead of a raw ``q.to(fp8)``. - # The raw cast quantizes the bf16 q with scale 1, which is ~10x less - # accurate than the scaled quantization (whose dequant is folded into - # ``mla_bmm1_scale`` below) and blows past the fp8 test tolerance. - # FP16/BF16: q is already the kernel dtype (or cast losslessly). - if kernel_dtype == torch.float8_e4m3fn and forward_args.quant_q_buffer is not None: - q_kernel = forward_args.quant_q_buffer.view(torch.float8_e4m3fn).view_as(q) - else: - q_kernel = q if q.dtype == kernel_dtype else q.to(kernel_dtype) - # Kernel layout for the q tensors is logical [H, D, S_q, B] with the D - # axis contiguous (stride 1) — see ``mark_layout_dynamic(leading_dim=1)`` - # on the op side and the kernel reference (permute_order=(2,3,1,0)). - # ``q_view`` is [B, S_q, H, D] contiguous, so the permuted *view* - # already has stride-1 on D; calling ``.contiguous()`` here would - # re-lay-out to stride-1 on B and violate the kernel's contract. - q_view = q_kernel.view(num_seqs, seq_len_q, h, d_latent + d_rope) - q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0) - q_rope = q_view[..., d_latent:].permute(2, 3, 1, 0) - - # Paged MLA pool view as the kernel's dtype. - # NOTE: pool tensor handle and block-table layout for MLA depends - # on kv-cache-manager wiring; revisit if ``get_buffers`` exposes a - # different layout. - kv_pool = metadata.kv_cache_manager.get_buffers(self.layer_idx) - if kernel_dtype in (torch.float16, torch.bfloat16) and kv_pool.dtype != kernel_dtype: - raise ValueError( - f"CuteDslAttention MLA decode {kernel_dtype} fast path requires " - f"matching KV cache dtype, got {kv_pool.dtype}" - ) - - # Paged-pool layout normalization (v1 vs v2 KV cache manager). - # ``get_buffers`` returns a per-layer view of shape - # [num_blocks, kv_factor, page_size, num_kv_heads, head_dim]. The v2 - # manager gives each layer its own densely-packed pool, so the block - # (dim-0) stride equals the packed block size. The v1 manager stores all - # layers interleaved in ONE pool, so the per-layer view's block stride is - # ``layers_in_pool * packed_block`` (non-packed) with a per-layer - # ``storage_offset``. The CuTe DSL paged TMA addresses pages with the - # packed block stride and does not honor a non-packed one, so for v1 it - # would read the wrong (interleaved) layer's memory. Re-expose the - # underlying pool as a packed ``[num_blocks * layers_in_pool, ...]`` view - # (copy-free) and fold the layer offset into the page table below: - # combined_block = block * layers_in_pool + layer_in_pool - # For v2 ``layers_in_pool == 1`` and this is a no-op. - packed_block = 1 - for s in kv_pool.shape[1:]: - packed_block *= s - block_stride = kv_pool.stride(0) - layers_in_pool = block_stride // packed_block if packed_block else 1 - layer_in_pool = 0 - if layers_in_pool > 1 and block_stride == layers_in_pool * packed_block: - layer_in_pool = kv_pool.storage_offset() // packed_block - kv_pool = kv_pool.as_strided( - (kv_pool.shape[0] * layers_in_pool, ) + tuple(kv_pool.shape[1:]), - (packed_block, ) + tuple(kv_pool.stride()[1:]), - 0, - ) - else: - layers_in_pool = 1 - - kv_pool_typed = kv_pool.view(kernel_dtype) - c_pool_latent = kv_pool_typed[..., :d_latent] - c_pool_rope = kv_pool_typed[..., d_latent : d_latent + d_rope] - - block_offsets = metadata.kv_cache_block_offsets - if block_offsets.dim() == 4: - # kv_cache_block_offsets is [num_pools, num_seqs, 2, max_blocks]. - # dim 0 is the *pool* index, not the layer — resolve this layer's - # pool via the layer→pool mapping ([num_layers, 2], col 0 = pool). - # (Indexing dim 0 by ``layer_idx`` is only valid for layer 0 / a - # single pool and goes out of bounds for every other layer.) - pool_mapping = metadata.host_kv_cache_pool_mapping - pool_idx = int(pool_mapping[self.layer_idx, 0]) - page_table_layer = block_offsets[pool_idx, :, 0, :] - elif block_offsets.dim() == 3: - page_table_layer = block_offsets[:, 0, :] - else: - page_table_layer = block_offsets - # Kernel: [max_pages, B], leading_dim=0 ⇒ pages contiguous (stride 1). - # ``page_table_layer`` is [B, max_pages] with max_pages stride 1, so the - # transposed *view* already has stride 1 on dim 0; ``.contiguous()`` - # would move stride 1 to B and violate the kernel's layout contract. - page_table = page_table_layer.transpose(0, 1).to(torch.int32) - - # Fold the interleaved-pool layer offset into the page table so it - # indexes the packed combined-slot view built above. The v1 block - # offsets are already in combined-slot units (block * layers_in_pool), - # so only the per-layer offset has to be added. No-op for v2 - # (layers_in_pool == 1, layer_in_pool == 0). - if layers_in_pool > 1: - page_table = page_table + layer_in_pool - - cache_seqs = metadata.kv_lens_cuda_runtime.to(torch.int32) - - if os.environ.get("TLLM_CUTE_DSL_DUMP"): - print( - "[CUTEDSL_DUMP] layer=%d kv_pool.shape=%s kv_pool.stride=%s " - "contiguous=%s block_offsets.shape=%s page_table.shape=%s " - "page_table=%s cache_seqs=%s" % ( - self.layer_idx, - tuple(kv_pool.shape), - tuple(kv_pool.stride()), - kv_pool.is_contiguous(), - tuple(block_offsets.shape), - tuple(page_table.shape), - page_table.t().tolist() if page_table.numel() < 64 else "(big)", - cache_seqs[:8].tolist(), - ), - flush=True, - ) - _pm = getattr(metadata, "host_kv_cache_pool_mapping", None) - print( - "[CUTEDSL_DUMP2] layer=%d pool_mapping=%s " - "raw_block_offsets[...,0,:]=%s" % ( - self.layer_idx, - _pm.tolist() if _pm is not None else None, - block_offsets[..., 0, :].tolist() - if block_offsets.dim() == 4 and block_offsets.numel() < 128 - else "(big/other)", - ), - flush=True, - ) - - split_kv = 1 - workspace = torch.empty(0, dtype=torch.float32, device=q.device) - block_split_kvs = torch.empty(0, dtype=torch.int32, device=q.device) - - # ``o`` mirrors the q layout: logical [H, D, S_q, B] with D contiguous - # (leading_dim=1). Allocate [B, S_q, H, D] contiguous and permute so the - # D axis keeps stride 1 (a plain contiguous [H, D, S_q, B] would put - # stride 1 on B and fail the op's layout check). - # - # For the FP8 kernel we force the output dtype to bf16 rather than fp8: - # the kernel's epilogue stores through ``o.element_type`` with a - # width-adaptive copy (so bf16 is supported — see the kernel's - # ``can_implement``), and writing the attention result straight to bf16 - # avoids an extra fp8 round-trip on the output. The FP16 kernel class - # keeps its native fp16/bf16 output. - out_kernel_dtype = torch.bfloat16 if kernel_dtype == torch.float8_e4m3fn else kernel_dtype - o_kernel = torch.empty( - (num_seqs, seq_len_q, h, d_latent), - dtype=out_kernel_dtype, - device=q.device, - ).permute(2, 3, 1, 0) - # ``lse`` is logical [H, S_q, B] with H contiguous (leading_dim=0); the - # reference builds it as [B, S_q, H] then permute_order=(2,1,0). - lse = torch.empty( - (num_seqs, seq_len_q, h), - dtype=torch.float32, - device=q.device, - ).permute(2, 1, 0) - - # MLA softmax scale follows the canonical TRT-LLM formula - # (see modules/attention.py:1453): - # softmax_scale = 1 / (sqrt(qk_head_dim) * q_scaling) - # where ``qk_head_dim`` is the *unabsorbed* Q head dim - # ``qk_nope_head_dim + qk_rope_head_dim``. We deliberately do NOT - # use ``(kv_lora_rank + qk_rope_head_dim)`` even though that is - # the absorbed attention's inner dimension — the scale stays bound - # to the original head dim so attention scores match the - # unabsorbed reference. ``self.q_scaling`` carries any YaRN - # ``mscale`` adjustment (set by the MLA module via - # ``q_scaling = 1 / (mscale * mscale)`` — see attention.py:1428). - qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - softmax_scale = float(1.0 / (math.sqrt(qk_head_dim) * self.q_scaling)) - output_scale = 1.0 - if (kernel_dtype == torch.float8_e4m3fn - and forward_args.mla_bmm1_scale is not None - and forward_args.mla_bmm2_scale is not None): - # ``mla_rope_generation`` folds the fp8 dequant (1/(q_scale*k_scale)) - # and the softmax scale into ``mla_bmm1_scale[1]`` (the FMHA - # scaleBmm1, see attentionOp.cpp bmm1_scale_offset=1), and the - # output/V dequant into ``mla_bmm2_scale[0]``. The CuTe DSL kernel - # reads q/KV as raw fp8 and applies these two scalar scales, so - # reusing them dequantizes identically to the TRTLLM FP8 MLA path. - # mla_bmm1_scale[1] is the TRTLLM FMHA scaleBmm1, which has log2(e) - # PRE-folded (its softmax uses exp2). The CuTe DSL kernel folds - # log2(e) AGAIN internally (softmax_scale_log2 = softmax_scale * - # LOG2_E), so pass the de-folded value to avoid applying it twice. - # - # CUDA-graph safety: ``.item()`` is a device->host copy + sync, which - # is illegal while a CUDA graph stream is capturing. These fp8 - # dequant scales are STATIC per layer — ``mla_rope_generation`` - # derives them from static quantization state (kv_scale_quant_orig, - # q_scaling, quant_mode), not from per-step activations — so we read - # them once (eagerly) and cache the host float on this per-layer - # backend instance. The generation CUDA-graph capture runs - # ``WARMUP_STEPS`` eager decode passes over every layer first - # (cuda_graph_runner.capture), so the cache is always populated - # before capture and the ``.item()`` never runs under capture. - cached = getattr(self, "_cute_dsl_fp8_scale", None) - if cached is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "CuteDslAttention: fp8 MLA decode scale not cached for " - f"layer {self.layer_idx} before CUDA graph capture; " - "reading it via .item() during capture is illegal. The " - "eager warmup should have populated the cache.") - softmax_scale = float( - forward_args.mla_bmm1_scale[1].item()) / _LOG2_E - output_scale = float(forward_args.mla_bmm2_scale[0].item()) - self._cute_dsl_fp8_scale = (softmax_scale, output_scale) - else: - softmax_scale, output_scale = cached - # Optional staticness check: re-read the live scale eagerly and - # warn if it ever drifts from the cached value (would mean the - # static-scale assumption — and thus the cache — is wrong). - if (os.environ.get("TLLM_CUTE_DSL_SCALE_DUMP") - and not torch.cuda.is_current_stream_capturing()): - live_sm = float( - forward_args.mla_bmm1_scale[1].item()) / _LOG2_E - live_os = float(forward_args.mla_bmm2_scale[0].item()) - print( - "[CUTEDSL_SCALE] layer=%d cached=(%.8f,%.8f) " - "live=(%.8f,%.8f) drift=%s" % ( - self.layer_idx, softmax_scale, output_scale, - live_sm, live_os, - abs(live_sm - softmax_scale) > 1e-6 - or abs(live_os - output_scale) > 1e-6), - flush=True, - ) - - # Paged cache layout for the kernel is logical [page_size, D, num_pages] - # with the D axis contiguous (leading_dim=1) — see the kernel reference - # (cache shape (num_pages, page_size, D), permute_order=(1,2,0)). - # ``get_buffers`` returns the pool as - # [num_pages, kv_factor=1, page_size, num_kv_heads=1, head_dim]; drop the - # two singleton axes and permute so D keeps stride 1. The latent/rope - # slices already carry stride 1 on their (last) D axis, so no copy. - c_latent_kernel = c_pool_latent.squeeze(3).squeeze(1).permute(1, 2, 0) - c_rope_kernel = c_pool_rope.squeeze(3).squeeze(1).permute(1, 2, 0) - - op( - q_latent, - q_rope, - c_latent_kernel, - c_rope_kernel, - page_table, - cache_seqs, - block_split_kvs, - o_kernel, - lse, - workspace, - self.num_heads, - seq_len_q, - page_size, - True, # is_persistent - True, # is_var_seq - False, # is_var_split_kv - split_kv, - softmax_scale, - output_scale, - ) - - attn_out = o_kernel.permute(3, 2, 0, 1).reshape(num_tokens, h * d_latent) - output.copy_(attn_out.to(output.dtype)) - return output - - # ================================================================== - # forward dispatch - # ================================================================== - - def forward( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata: TrtllmAttentionMetadata, - forward_args: Optional[AttentionForwardArgs] = None, - **kwargs, - ) -> torch.Tensor: - if forward_args is None and kwargs: - forward_args = AttentionForwardArgs(**kwargs) - kwargs = {} - - # Phase routing — only the MLA decode-only batch is eligible for the - # CuTe DSL fast path. Everything else falls through to TRTLLM: - # - non-MLA attention → TRTLLM - # - prefill / mixed batch (has context) → TRTLLM - # - decode-only MLA → try CuteDSL FP8/FP16/BF16, - # TRTLLM on failure - is_decode_only_mla = ( - self.is_mla_enable and metadata.num_contexts == 0 and metadata.num_generations > 0 - ) - - # The CuTe DSL MLA decode kernel is non-causal: every query token - # attends to the whole cached KV with no per-query mask. That is correct - # for the plain decode case (the MLA module passes the default CAUSAL / - # FULL predefined mask with no mask tensor), but it cannot honor any - # per-query causal masking. Fall back to TRTLLM whenever masking the - # kernel cannot express is required: - # - an explicit CUSTOM mask or mask tensor, or - # - speculative decoding (MTP / Eagle / tree) — signalled by - # ``metadata.use_spec_decoding``; those steps feed >1 query token per - # request and require a causal/tree mask the non-causal kernel lacks. - attention_mask = ( - forward_args.attention_mask - if forward_args is not None else PredefinedAttentionMask.CAUSAL - ) - attention_mask_data = ( - forward_args.attention_mask_data if forward_args is not None else None - ) - mask_supported = ( - attention_mask != CustomAttentionMask.CUSTOM - and attention_mask_data is None - and not getattr(metadata, "use_spec_decoding", False) - ) - - if ( - is_decode_only_mla - and mask_supported - and self._cute_dsl_mla_decode_common_preconditions(metadata, forward_args) - and q.shape[0] % metadata.num_generations == 0 - ): - # Direct dtype-based dispatch (no per-dtype eligibility helper). - # The kernel-level Runner runs ``can_implement`` for the chosen - # dtype; anything it rejects falls through to TRTLLM via the - # try/except below. - if getattr(self, "has_fp8_kv_cache", False): - kernel_dtype = torch.float8_e4m3fn - elif q.dtype in (torch.float16, torch.bfloat16): - kv_pool_dtype = metadata.kv_cache_manager.get_buffers(self.layer_idx).dtype - kernel_dtype = q.dtype if kv_pool_dtype == q.dtype else None - else: - kernel_dtype = None - - if kernel_dtype is not None: - try: - output = q.new_empty( - (q.shape[0], self.num_heads * self.kv_lora_rank), dtype=q.dtype - ) - result = self._dispatch_cute_dsl_mla_decode( - q, metadata, forward_args, output, kernel_dtype - ) - # Positive confirmation (once per process) that the CuteDSL - # MLA decode kernel was actually engaged — symmetric to the - # fallback warning below so a real kernel call can never be - # confused with a silent TRTLLM fallback. - logger.warning_once( - "CuteDslAttention: MLA decode fast path ENGAGED " - "(kernel_dtype=%s); CuteDSL kernel called." % kernel_dtype, - key="cute_dsl_mla_decode_engaged", - ) - return result - except Exception as exc: # noqa: BLE001 - # Surface the fallback once per process (keyed dedup keeps - # it from flooding the log on every decode step) so a broken - # kernel can never silently masquerade as a working one. - logger.warning_once( - "CuteDslAttention: MLA decode fast path " - "(kernel_dtype=%s) failed (%s); falling back " - "to TRTLLM backend." % (kernel_dtype, exc), - key="cute_dsl_mla_decode_fallback", - ) - - return super().forward(q, k, v, metadata, forward_args=forward_args, **kwargs) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py index 1c3981abcf91..580b8dfa9138 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/__init__.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .cute_dsl import CuteDslMlaFmha from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha @@ -22,6 +23,7 @@ __all__ = [ "DEFAULT_FMHA_LIBS", "FMHA_LIBS", + "CuteDslMlaFmha", "FallbackFmha", "FlashInferTrtllmGenFmha", "Fmha", diff --git a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py new file mode 100644 index 000000000000..fc6ecde7a0e9 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py @@ -0,0 +1,465 @@ +# 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. +"""CuTe DSL MLA decode FMHA library.""" + +import math +import os +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + AttentionInputType, + CustomAttentionMask, +) +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger + +from .phased import FmhaParams, PhasedFmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + +_LOG2_E = math.log2(math.e) + + +class CuteDslMlaFmha(PhasedFmha): + """Blackwell CuTe DSL FMHA library for decode-only MLA.""" + + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + if not IS_CUTLASS_DSL_AVAILABLE: + logger.debug("CuTe DSL MLA FMHA is unavailable: nvidia-cutlass-dsl is not installed.") + return False + + sm = get_sm_version() + if sm not in (100, 103): + logger.debug(f"CuTe DSL MLA FMHA is unavailable: requires SM100 or SM103, got SM{sm}.") + return False + + if not attn.is_mla_enable: + logger.debug("CuTe DSL MLA FMHA is unavailable: only MLA is supported.") + return False + if attn.predicted_tokens_per_seq is None or not (1 <= attn.predicted_tokens_per_seq <= 4): + logger.debug( + "CuTe DSL MLA FMHA is unavailable: predicted_tokens_per_seq " + f"must be in [1, 4], got {attn.predicted_tokens_per_seq}." + ) + return False + if attn.kv_lora_rank is None or attn.kv_lora_rank <= 0: + logger.debug("CuTe DSL MLA FMHA is unavailable: kv_lora_rank must be positive.") + return False + if attn.qk_rope_head_dim is None or attn.qk_rope_head_dim <= 0: + logger.debug("CuTe DSL MLA FMHA is unavailable: qk_rope_head_dim must be positive.") + return False + if attn.qk_nope_head_dim is None or attn.qk_nope_head_dim <= 0: + logger.debug("CuTe DSL MLA FMHA is unavailable: qk_nope_head_dim must be positive.") + return False + if attn.kv_lora_rank != 512 or attn.qk_rope_head_dim != 64: + logger.debug( + "CuTe DSL MLA FMHA is unavailable: kernels require kv_lora_rank=512 and " + f"qk_rope_head_dim=64, got kv_lora_rank={attn.kv_lora_rank}, " + f"qk_rope_head_dim={attn.qk_rope_head_dim}." + ) + return False + if attn.num_heads > 128: + logger.debug( + f"CuTe DSL MLA FMHA is unavailable: num_heads must be <= 128, got {attn.num_heads}." + ) + return False + + return True + + @staticmethod + def _get_kernel_dtype(attn: "TrtllmAttention", q: torch.Tensor) -> Optional[torch.dtype]: + if getattr(attn, "has_fp8_kv_cache", False): + return torch.float8_e4m3fn + if q.dtype in (torch.float16, torch.bfloat16): + return q.dtype + return None + + @staticmethod + def _select_page_table_layer( + block_offsets: torch.Tensor, + layer_idx: int, + host_kv_cache_pool_mapping: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + if block_offsets.dim() == 4: + if block_offsets.shape[2] < 1: + return None + if host_kv_cache_pool_mapping is not None: + if layer_idx >= host_kv_cache_pool_mapping.shape[0]: + return None + pool_idx = int(host_kv_cache_pool_mapping[layer_idx, 0]) + else: + pool_idx = layer_idx if block_offsets.shape[0] > 1 else 0 + if pool_idx >= block_offsets.shape[0]: + return None + return block_offsets[pool_idx, :, 0, :] + if block_offsets.dim() == 3: + if block_offsets.shape[1] < 1: + return None + return block_offsets[:, 0, :] + if block_offsets.dim() == 2: + return block_offsets + return None + + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> bool: + supported, reason = self._is_supported_with_reason( + q, + self.attn, + metadata, + forward_args, + ) + if not supported: + logger.debug(f"CuTe DSL MLA FMHA does not support request: {reason}") + return supported + + def _is_supported_with_reason( + self, + q: torch.Tensor, + attn: "TrtllmAttention", + meta: "TrtllmAttentionMetadata", + fwd: AttentionForwardArgs, + ) -> tuple[bool, str]: + if fwd.attention_input_type != AttentionInputType.generation_only: + return False, "CuTe DSL MLA FMHA only supports generation-only attention." + if meta.num_contexts != 0 or meta.num_generations <= 0: + return False, "CuTe DSL MLA FMHA only supports decode-only batches." + if meta.beam_width != 1: + return False, f"Beam search is not supported, got beam_width={meta.beam_width}." + if ( + fwd.attention_mask == CustomAttentionMask.CUSTOM + or fwd.attention_mask_data is not None + or getattr(meta, "use_spec_decoding", False) + or getattr(meta, "is_spec_decoding_enabled", False) + ): + return False, "CuTe DSL MLA FMHA does not support custom/speculative masks." + if q.shape[0] % meta.num_generations != 0: + return ( + False, + f"num_tokens ({q.shape[0]}) must be divisible by " + f"num_generations ({meta.num_generations}).", + ) + seq_len_q = q.shape[0] // meta.num_generations + if not (1 <= seq_len_q <= 4): + return False, f"Only query lengths in [1, 4] are supported, got {seq_len_q}." + if meta.kv_cache_block_offsets is None: + return False, "Paged KV block offsets are required." + page_table_layer = self._select_page_table_layer( + meta.kv_cache_block_offsets, + attn.layer_idx, + meta.host_kv_cache_pool_mapping, + ) + if page_table_layer is None: + return ( + False, + "Unsupported KV block offsets shape " + f"{tuple(meta.kv_cache_block_offsets.shape)} for layer_idx={attn.layer_idx}.", + ) + if meta.kv_cache_manager is None: + return False, "KV cache manager is required." + if fwd.latent_cache is None: + return False, "latent_cache is required." + if fwd.output is None: + return False, "output is required." + + tokens_per_block = meta.tokens_per_block + if tokens_per_block is None: + tokens_per_block = getattr(meta.kv_cache_manager, "tokens_per_block", 0) + if tokens_per_block <= 1 or 128 % tokens_per_block != 0: + return ( + False, + f"tokens_per_block must divide 128 and be greater than 1, got {tokens_per_block}.", + ) + + kernel_dtype = self._get_kernel_dtype(attn, q) + if kernel_dtype is None: + return ( + False, + f"Unsupported dtype combination: q={q.dtype}, " + f"has_fp8_kv_cache={getattr(attn, 'has_fp8_kv_cache', False)}.", + ) + if kernel_dtype == torch.float8_e4m3fn and ( + fwd.quant_q_buffer is None or fwd.mla_bmm1_scale is None or fwd.mla_bmm2_scale is None + ): + return ( + False, + "FP8 CuTe DSL MLA decode requires quant_q_buffer, " + "mla_bmm1_scale, and mla_bmm2_scale from MLA RoPE generation.", + ) + if kernel_dtype in (torch.float16, torch.bfloat16): + kv_pool_dtype = meta.kv_cache_manager.get_buffers(attn.layer_idx).dtype + if kv_pool_dtype != kernel_dtype: + return ( + False, + f"CuTe DSL MLA {kernel_dtype} fast path requires matching " + f"KV cache dtype, got {kv_pool_dtype}.", + ) + + return True, "" + + def _run_mla_decode( + self, + q: torch.Tensor, + output: torch.Tensor, + params: FmhaParams, + kernel_dtype: torch.dtype, + ) -> None: + attn = params.attn + meta = params.meta + + if kernel_dtype == torch.float8_e4m3fn: + op = torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell + elif kernel_dtype in (torch.float16, torch.bfloat16): + op = torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell + else: + raise ValueError( + f"CuTe DSL MLA FMHA got unsupported kernel_dtype={kernel_dtype}; " + "expected torch.float8_e4m3fn, torch.float16, or torch.bfloat16." + ) + + num_tokens = q.shape[0] + batch_size = params.num_requests + seq_len_q = num_tokens // batch_size + if seq_len_q * batch_size != num_tokens: + raise RuntimeError( + f"CuTe DSL MLA decode expects num_tokens ({num_tokens}) divisible by " + f"batch_size ({batch_size})." + ) + + d_latent = attn.kv_lora_rank + d_rope = attn.qk_rope_head_dim + qk_nope_head_dim = attn.qk_nope_head_dim + if d_latent is None or d_rope is None or qk_nope_head_dim is None: + raise RuntimeError("CuTe DSL MLA decode requires complete MLA dimensions.") + + num_heads = attn.num_heads + page_size = params.tokens_per_block + + if kernel_dtype == torch.float8_e4m3fn and params.fwd.quant_q_buffer is not None: + q_kernel = params.fwd.quant_q_buffer.view(torch.float8_e4m3fn).view_as(q) + else: + q_kernel = q if q.dtype == kernel_dtype else q.to(kernel_dtype) + q_view = q_kernel.view(batch_size, seq_len_q, num_heads, d_latent + d_rope) + + kv_pool = meta.kv_cache_manager.get_buffers(attn.layer_idx) + if kernel_dtype in (torch.float16, torch.bfloat16) and kv_pool.dtype != kernel_dtype: + raise RuntimeError( + f"CuTe DSL MLA {kernel_dtype} fast path requires matching " + f"KV cache dtype, got {kv_pool.dtype}." + ) + # Paged-pool layout normalization for both KV cache managers. + # KVCacheManagerV2 exposes each layer as a densely-packed page pool. + # KVCacheManagerV1 exposes a per-layer view over one interleaved pool, + # where dim-0 strides by all local layers. The CuTe DSL kernel addresses + # pages with the packed block stride, so represent v1 as a packed + # combined-slot view and fold this layer's slot offset into the page + # table below. + packed_block = 1 + for size in kv_pool.shape[1:]: + packed_block *= size + block_stride = kv_pool.stride(0) + layers_in_pool = block_stride // packed_block if packed_block else 1 + layer_in_pool = 0 + if layers_in_pool > 1 and block_stride == layers_in_pool * packed_block: + layer_in_pool = kv_pool.storage_offset() // packed_block + kv_pool = kv_pool.as_strided( + (kv_pool.shape[0] * layers_in_pool, *kv_pool.shape[1:]), + (packed_block, *kv_pool.stride()[1:]), + 0, + ) + else: + layers_in_pool = 1 + + kv_pool_typed = kv_pool.view(kernel_dtype) + if kv_pool_typed.dim() != 5 or kv_pool_typed.shape[1] != 1 or kv_pool_typed.shape[3] != 1: + raise RuntimeError( + "CuTe DSL MLA decode expects KV cache layout " + f"[num_pages, 1, page_size, 1, head_dim], got {tuple(kv_pool_typed.shape)}." + ) + if kv_pool_typed.shape[2] != page_size or kv_pool_typed.shape[-1] < d_latent + d_rope: + raise RuntimeError( + "CuTe DSL MLA decode got incompatible KV cache shape " + f"{tuple(kv_pool_typed.shape)} for page_size={page_size}, " + f"kv_lora_rank={d_latent}, qk_rope_head_dim={d_rope}." + ) + + block_offsets = meta.kv_cache_block_offsets + page_table_layer = self._select_page_table_layer( + block_offsets, + attn.layer_idx, + meta.host_kv_cache_pool_mapping, + ) + if page_table_layer is None: + raise RuntimeError( + "CuTe DSL MLA decode got unsupported KV block offsets shape " + f"{tuple(block_offsets.shape)} for layer_idx={attn.layer_idx}." + ) + cache_seqs_base = params.sequence_lengths.to(torch.int32) + page_table = page_table_layer.transpose(0, 1).to(torch.int32) + if layers_in_pool > 1: + page_table = page_table + layer_in_pool + + if os.environ.get("TLLM_CUTE_DSL_DUMP"): + print( + "[CUTEDSL_DUMP] layer=%d kv_pool.shape=%s kv_pool.stride=%s " + "contiguous=%s block_offsets.shape=%s page_table.shape=%s " + "page_table=%s cache_seqs=%s" + % ( + attn.layer_idx, + tuple(kv_pool.shape), + tuple(kv_pool.stride()), + kv_pool.is_contiguous(), + tuple(block_offsets.shape), + tuple(page_table.shape), + page_table.t().tolist() if page_table.numel() < 64 else "(big)", + cache_seqs_base[:8].tolist(), + ), + flush=True, + ) + + # KVCacheManager exposes NHD pages as [num_pages, 1, page_size, 1, head_dim]. + # The CuTe DSL kernel consumes a paged [page_size, dim, num_pages] view + # with the dim axis contiguous. + kv_pages = kv_pool_typed[:, 0, :, 0, : d_latent + d_rope] + c_pool_latent = kv_pages[..., :d_latent].permute(1, 2, 0) + c_pool_rope = kv_pages[..., d_latent:].permute(1, 2, 0) + + block_split_kvs = torch.empty(0, dtype=torch.int32, device=q.device) + split_kv = 1 + workspace = torch.empty(0, dtype=torch.float32, device=q.device) + + softmax_scale = float(1.0 / (math.sqrt(qk_nope_head_dim + d_rope) * attn.q_scaling)) + output_scale = 1.0 + if kernel_dtype == torch.float8_e4m3fn: + if params.fwd.mla_bmm1_scale is None or params.fwd.mla_bmm2_scale is None: + raise RuntimeError("FP8 CuTe DSL MLA decode requires MLA FP8 scales.") + cached = getattr(self, "_cute_dsl_fp8_scale", None) + if cached is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CuTe DSL MLA FMHA: fp8 decode scale was not cached for " + f"layer {attn.layer_idx} before CUDA graph capture." + ) + softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E + output_scale = float(params.fwd.mla_bmm2_scale[0].item()) + self._cute_dsl_fp8_scale = (softmax_scale, output_scale) + else: + softmax_scale, output_scale = cached + if ( + os.environ.get("TLLM_CUTE_DSL_SCALE_DUMP") + and not torch.cuda.is_current_stream_capturing() + ): + live_softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E + live_output_scale = float(params.fwd.mla_bmm2_scale[0].item()) + print( + "[CUTEDSL_SCALE] layer=%d cached=(%.8f,%.8f) " + "live=(%.8f,%.8f) drift=%s" + % ( + attn.layer_idx, + softmax_scale, + output_scale, + live_softmax_scale, + live_output_scale, + abs(live_softmax_scale - softmax_scale) > 1e-6 + or abs(live_output_scale - output_scale) > 1e-6, + ), + flush=True, + ) + + out_kernel_dtype = torch.bfloat16 if kernel_dtype == torch.float8_e4m3fn else kernel_dtype + output_view = output.view(batch_size, seq_len_q, num_heads, d_latent) + for query_idx in range(seq_len_q): + q_step = q_view[:, query_idx : query_idx + 1, :, :] + q_latent = q_step[..., :d_latent].permute(2, 3, 1, 0) + q_rope = q_step[..., d_latent:].permute(2, 3, 1, 0) + + o_storage = torch.empty( + (batch_size, 1, num_heads, d_latent), + dtype=out_kernel_dtype, + device=q.device, + ) + o_kernel = o_storage.permute(2, 3, 1, 0) + lse_storage = torch.empty( + (batch_size, 1, num_heads), + dtype=torch.float32, + device=q.device, + ) + lse = lse_storage.permute(2, 1, 0) + + # MLA RoPE generation has already appended all query tokens in this + # step. For multi-query decode, trim the effective KV length so each + # query attends only through its own generated token. + cache_seqs = cache_seqs_base - (seq_len_q - query_idx - 1) + + op( + q_latent, + q_rope, + c_pool_latent, + c_pool_rope, + page_table, + cache_seqs, + block_split_kvs, + o_kernel, + lse, + workspace, + num_heads, + 1, # seq_len_q + page_size, + True, # is_persistent + True, # is_var_seq + False, # is_var_split_kv + split_kv, + softmax_scale, + output_scale, + ) + + attn_out = o_kernel.permute(3, 2, 0, 1).reshape(batch_size, num_heads, d_latent) + output_view[:, query_idx, :, :].copy_(attn_out.to(output.dtype)) + + def run_mla_generation( + self, + params: FmhaParams, + ) -> None: + if params.qkv_input is None: + raise RuntimeError("CuTe DSL MLA generation requires qkv_input.") + if params.context_buf is None: + raise RuntimeError("CuTe DSL MLA generation requires context_buf.") + if params.sequence_lengths is None: + raise RuntimeError("CuTe DSL MLA generation requires sequence lengths.") + + kernel_dtype = self._get_kernel_dtype(params.attn, params.qkv_input) + if kernel_dtype is None: + raise RuntimeError("CuTe DSL MLA generation was selected for an unsupported dtype.") + + self._run_mla_decode( + params.qkv_input, + params.context_buf, + params, + kernel_dtype, + ) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/registry.py b/tensorrt_llm/_torch/attention_backend/fmha/registry.py index 97467657e9b5..efacb3fadb8e 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/registry.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/registry.py @@ -16,6 +16,7 @@ import os from typing import TypeAlias +from .cute_dsl import CuteDslMlaFmha from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha from .interface import Fmha @@ -23,6 +24,7 @@ FmhaCls: TypeAlias = type[Fmha] FMHA_LIBS: dict[str, FmhaCls] = { + "cute_dsl_mla": CuteDslMlaFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, } diff --git a/tensorrt_llm/_torch/attention_backend/utils.py b/tensorrt_llm/_torch/attention_backend/utils.py index fa699c5d536c..2c4bc574bb58 100644 --- a/tensorrt_llm/_torch/attention_backend/utils.py +++ b/tensorrt_llm/_torch/attention_backend/utils.py @@ -48,9 +48,6 @@ def get_attention_backend( elif backend_name == "FLASHINFER_STAR_ATTENTION" and IS_FLASHINFER_AVAILABLE: from .star_flashinfer import StarAttention return StarAttention - elif backend_name == "CUTEDSL": - from .cute_dsl import CuteDslAttention - return CuteDslAttention logger.warning("Falling back to TRTLLM attention backend") return TrtllmAttention diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 3bf98374956e..da3c670da804 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -7330,12 +7330,12 @@ def _( device=q.device) # ========================================================================= - # MLA decode (Blackwell) — wraps the CuTe DSL kernels that live at + # MLA decode (Blackwell) - wraps the CuTe DSL kernels that live at # tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/. - # Used by the CUTEDSL attention backend (see attention_backend/cute_dsl.py). + # Used by the cute_dsl_mla FMHA library (see attention_backend/fmha/cute_dsl.py). # # One generic Runner ``CuteDSLNVMlaDecodeBlackwellRunner`` services both - # FP8 and FP16/BF16 paths — only the cutlass ``in_dtype`` is passed at + # FP8 and FP16/BF16 paths - only the cutlass ``in_dtype`` is passed at # construction; the kernel class is derived from it via # ``CuteDSLNVMlaDecodeBlackwellRunner._KERNEL_CLASS_BY_DTYPE``. Each # dtype still has its own ``@torch.library.custom_op`` (distinct op @@ -7343,13 +7343,13 @@ def _( # they hand the generic Runner. # # torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell - # → CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.Float8E4M3FN) - # (→ BlackwellMultiHeadLatentAttentionForwardFP8) + # -> CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.Float8E4M3FN) + # (-> BlackwellMultiHeadLatentAttentionForwardFP8) # # torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell - # → CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.Float16) + # -> CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.Float16) # or CuteDSLNVMlaDecodeBlackwellRunner(in_dtype=cutlass.BFloat16) - # (→ BlackwellMultiHeadLatentAttentionForwardFP16) + # (-> BlackwellMultiHeadLatentAttentionForwardFP16) # ========================================================================= from ..cute_dsl_kernels.blackwell.attention.mla.mla_decode_fp8 import \ @@ -7362,29 +7362,29 @@ def _( class CuteDSLNVMlaDecodeBlackwellRunner(TunableRunner): """Generic TunableRunner for the Blackwell CuTe DSL MLA decode kernels. - Works for FP8, FP16, and BF16 — pass the cutlass input dtype at + Works for FP8, FP16, and BF16 - pass the cutlass input dtype at construction; the kernel class is derived from it: CuteDSLNVMlaDecodeBlackwellRunner( - in_dtype=cutlass.Float8E4M3FN, ...) # → ...ForwardFP8 + in_dtype=cutlass.Float8E4M3FN, ...) # -> ...ForwardFP8 CuteDSLNVMlaDecodeBlackwellRunner( - in_dtype=cutlass.Float16, ...) # → ...ForwardFP16 + in_dtype=cutlass.Float16, ...) # -> ...ForwardFP16 CuteDSLNVMlaDecodeBlackwellRunner( - in_dtype=cutlass.BFloat16, ...) # → ...ForwardFP16 + in_dtype=cutlass.BFloat16, ...) # -> ...ForwardFP16 ``get_valid_tactics`` returns the tiler shapes as tactics (``(mma_qk_tiler_mn, mma_pv_tiler_mn)`` tuples), filtered by the kernel's static ``can_implement``. The current candidate list - carries only ``((128, 128), (128, 256))`` — the lone combination - both kernels accept today — but more can be added without + carries only ``((128, 128), (128, 256))`` - the lone combination + both kernels accept today - but more can be added without touching ``forward``. ``kernel_cache`` is class-level and keyed - by ``(in_dtype, ..., mma_qk_tiler_mn, mma_pv_tiler_mn)``, so FP8 - and FP16 compilations, plus future tilers, coexist without - collisions. + by ``(in_dtype, ..., out_dtype, mma_qk_tiler_mn, + mma_pv_tiler_mn)``, so FP8/FP16/BF16 output variants and future + tilers coexist without collisions. """ kernel_cache = dict() - # in_dtype → kernel class. The kernels' own ``can_implement`` is + # in_dtype -> kernel class. The kernels' own ``can_implement`` is # what ultimately rejects unsupported dtypes, but this lookup # picks which kernel we even try to compile. _KERNEL_CLASS_BY_DTYPE = { @@ -7424,7 +7424,7 @@ def __init__( def unique_id(self): # `kernel_class` is derived from `in_dtype`, so dropping it # from the key keeps cache slots 1-to-1 with the in_dtype. - # The tilers are NOT here — they're part of the tactic and + # The tilers are NOT here - they're part of the tactic and # appended into the cache key inside ``forward``. return ( self.in_dtype, @@ -7451,12 +7451,18 @@ def get_valid_tactics( if get_sm_version() not in (100, 103): return [] q_latent, q_rope, _c_latent, _c_rope, _page_table, cache_seqs, \ - *_rest = inputs + _block_split_kvs, o, *_rest = inputs h, latent_dim, seq_len_q, _ = q_latent.shape rope_dim = q_rope.shape[1] batch_size = cache_seqs.shape[0] + if o.dtype == torch.float16: + out_dtype = cutlass.Float16 + elif o.dtype == torch.bfloat16: + out_dtype = cutlass.BFloat16 + else: + out_dtype = self.in_dtype - # Candidate tilers — widen this list to enable AutoTuner + # Candidate tilers - widen this list to enable AutoTuner # exploration over tile shapes. Each entry is # ``(mma_qk_tiler_mn, mma_pv_tiler_mn)``. candidate_tiler_tactics = [ @@ -7473,7 +7479,7 @@ def get_valid_tactics( latent_dim, rope_dim, self.in_dtype, # in_dtype - self.in_dtype, # out_dtype + out_dtype, cutlass.Float32, # acc_dtype cutlass.Float32, # lse_dtype mma_qk_tiler_mn, @@ -7514,8 +7520,8 @@ def forward( output_scale = float(kwargs.get("output_scale", 1.0)) # Unpack the tactic produced by ``get_valid_tactics``. When - # AutoTuner isn't engaged (e.g. the attention backend calls - # the op without ``choose_one``), tactic may be ``None`` — + # AutoTuner isn't engaged (e.g. the FMHA library calls the op + # without ``choose_one``), tactic may be ``None`` - # fall back to the default (128,128)/(128,256) shape. if isinstance(tactic, tuple) and len(tactic) == 2: mma_qk_tiler_mn, mma_pv_tiler_mn = tactic @@ -7527,7 +7533,18 @@ def forward( torch_stream = torch.cuda.current_stream() stream = cuda.CUstream(torch_stream.cuda_stream) - cache_key = self.unique_id() + (mma_qk_tiler_mn, mma_pv_tiler_mn) + if o.dtype == torch.float16: + out_dtype = cutlass.Float16 + elif o.dtype == torch.bfloat16: + out_dtype = cutlass.BFloat16 + else: + out_dtype = self.in_dtype + + cache_key = self.unique_id() + ( + out_dtype, + mma_qk_tiler_mn, + mma_pv_tiler_mn, + ) if cache_key not in CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache: hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -7753,9 +7770,8 @@ def cute_dsl_mla_decode_fp16_blackwell( "trtllm::cute_dsl_mla_decode_fp16_blackwell supports " "torch.float16 or torch.bfloat16 inputs, got " f"{q_latent.dtype}") - if not ( - q_rope.dtype == c_latent.dtype == c_rope.dtype == o.dtype - == q_latent.dtype): + if not (q_rope.dtype == c_latent.dtype == c_rope.dtype == o.dtype == + q_latent.dtype): raise ValueError( "trtllm::cute_dsl_mla_decode_fp16_blackwell requires q, KV, " f"and output dtypes to match; got q_latent={q_latent.dtype}, " diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 4e98b8d688a4..2743a5c4149f 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -246,21 +246,23 @@ The main differences across backends: #### 3.2.2 `TRTLLM` internal FMHA libraries `TrtllmAttention` dispatches attention through an ordered list of internal FMHA -libraries. `FlashInferTrtllmGenFmha` integrates trtllm-gen kernels from -FlashInfer into the `TRTLLM` backend, and `FallbackFmha` calls the regular -`thop.attention` runtime path. These are not separate attention backends. +libraries. `CuteDslMlaFmha` integrates Blackwell CuTe DSL MLA decode kernels, +`FlashInferTrtllmGenFmha` integrates trtllm-gen kernels from FlashInfer into +the `TRTLLM` backend, and `FallbackFmha` calls the regular `thop.attention` +runtime path. These are not separate attention backends. `TLLM_FMHA_LIBS` controls the ordered list. Unset means -`flashinfer_trtllm_gen,fallback`; use `TLLM_FMHA_LIBS=fallback` or -`TLLM_FMHA_LIBS=-flashinfer_trtllm_gen` to force the fallback path. Each FMHA -library exposes `is_available()` for module/static environment checks and -`is_supported()` for per-forward request checks. +`cute_dsl_mla,flashinfer_trtllm_gen,fallback`; use `TLLM_FMHA_LIBS=fallback` +or `TLLM_FMHA_LIBS=-cute_dsl_mla,-flashinfer_trtllm_gen` to force the fallback +path. Each FMHA library exposes `is_available()` for module/static environment +checks and `is_supported()` for per-forward request checks. The FMHA package is split by role: - `fmha/interface.py` defines the `Fmha` runtime contract. - `fmha/phased.py` defines `PhasedFmha`, which handles mixed context/generation requests and dispatches them to phase-specific hooks. +- `fmha/cute_dsl.py` implements the CuTe DSL MLA decode FMHA library. - `fmha/flashinfer_trtllm_gen.py` implements the FlashInfer trtllm-gen FMHA library. - `fmha/fallback.py` implements the regular `thop.attention` fallback library. diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py index 8e21cafaf778..431ddd39f131 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py @@ -16,21 +16,19 @@ This test validates the CuTe DSL MLA *decode* kernels added under ``tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla`` and dispatched -through the ``CUTEDSL`` attention backend -(``tensorrt_llm/_torch/attention_backend/cute_dsl.py``): +through the ``cute_dsl_mla`` TRTLLM FMHA library +(``tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py``): -- FP8 path → ``torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell`` -- FP16/BF16 path → ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` +- FP8 path: ``torch.ops.trtllm.cute_dsl_mla_decode_fp8_blackwell`` +- FP16/BF16 path: ``torch.ops.trtllm.cute_dsl_mla_decode_fp16_blackwell`` Only the generation (decode) steps are asserted for numerical correctness. The context phase runs solely to populate the paged KV cache and to build the reference latent cache (``skip_context_assert=True``). -Crucially, the test monkeypatches ``CuteDslAttention._dispatch_cute_dsl_mla_decode`` -to count invocations and asserts the CuTe DSL decode path was actually taken on -every decode step. Without this guard the backend would silently fall back to -TRTLLM on any kernel error and the test could "pass" without ever exercising -the CuTe DSL kernel under test. +Crucially, the test monkeypatches ``CuteDslMlaFmha._run_mla_decode`` to count +invocations and asserts the CuTe DSL decode path was actually taken on every +decode step. Platform: Blackwell SM100 / SM103 only. """ @@ -47,7 +45,8 @@ # DeepSeek-V3-like MLA geometry the CuTe DSL kernel targets (num_heads=128, # latent_dim=512, rope_dim=64). Kept small along the batch/step axes so the -# decode-only test stays fast. +# decode-only test stays fast. Multi-token/MTP decode is handled by replaying +# the single-token CuTe DSL kernel once per intra-step query. _DECODE_CONTEXT_LENGTHS = [ [10, 12, 5], [100, 300, 20, 10], @@ -57,11 +56,11 @@ # Multi-layer is the structural difference between this single-step test and the # real DeepSeek-V3 E2E run (61 MLA layers). With ``num_layers == 1`` the dispatch # only ever sees ``layer_idx == 0``, so the per-layer paged-KV resolution in -# ``CuteDslAttention._dispatch_cute_dsl_mla_decode`` (the +# ``CuteDslMlaFmha._run_mla_decode`` (the # ``host_kv_cache_pool_mapping[layer_idx]`` / per-layer ``get_buffers`` / # block-offset path) is never exercised. The E2E run produces correct output on # the first generated token (which comes from the TRTLLM prefill) and then -# degenerates on every subsequent CuteDSL decode step — consistent with the +# degenerates on every subsequent CuteDSL decode step, consistent with the # decode kernel reading the wrong blocks for ``layer_idx > 0``. Parametrize over # >1 layers so the unit test reproduces that real case. _DECODE_NUM_LAYERS = [1, 2] @@ -115,17 +114,18 @@ def _build_rope_config(scenario: Scenario) -> RopeConfig: @pytest.fixture def cute_dsl_decode_counter(monkeypatch): - """Count *successful* CuTe DSL MLA decode dispatches so the test fails on + """Count successful CuTe DSL MLA decode dispatches so the test fails on silent fallback to the TRTLLM backend. - The increment happens only after the real dispatch returns: if the kernel - raises, ``CuteDslAttention.forward`` catches it and falls back to TRTLLM, - the counter does not advance, and the per-test assertion on the expected - dispatch count fails loudly instead of the broken kernel masquerading as a - working one.""" - from tensorrt_llm._torch.attention_backend.cute_dsl import CuteDslAttention + The increment happens only after the real dispatch returns. If the kernel + raises or the registry selects the fallback FMHA library, the counter does + not advance and the per-test assertion fails loudly instead of the broken + kernel masquerading as a working one.""" + from tensorrt_llm._torch.attention_backend.fmha.cute_dsl import CuteDslMlaFmha - original = CuteDslAttention._dispatch_cute_dsl_mla_decode + monkeypatch.setenv("TLLM_FMHA_LIBS", "cute_dsl_mla,fallback") + + original = CuteDslMlaFmha._run_mla_decode counter = {"calls": 0} def _counting_dispatch(self, *args, **kwargs): @@ -133,7 +133,7 @@ def _counting_dispatch(self, *args, **kwargs): counter["calls"] += 1 return result - monkeypatch.setattr(CuteDslAttention, "_dispatch_cute_dsl_mla_decode", _counting_dispatch) + monkeypatch.setattr(CuteDslMlaFmha, "_run_mla_decode", _counting_dispatch) return counter @@ -153,7 +153,7 @@ def test_cute_dsl_mla_decode( rope_config = _build_rope_config(scenario) _run_test_for_backend( - "CUTEDSL", + "TRTLLM", num_heads=scenario.num_heads, num_kv_heads=scenario.num_kv_heads, num_layers=scenario.num_layers, @@ -200,9 +200,7 @@ def test_cute_dsl_mla_decode( @pytest.mark.parametrize("kernel", list(_KERNEL_DTYPES)) @pytest.mark.parametrize("num_layers", _DECODE_NUM_LAYERS, ids=lambda x: f"num_layers={x}") @pytest.mark.parametrize("v2_kv_cache", [True, False], ids=lambda x: f"v2_kv_cache={x}") -def test_cute_dsl_mla_decode_long_decode( - v2_kv_cache, num_layers, kernel, cute_dsl_decode_counter -): +def test_cute_dsl_mla_decode_long_decode(v2_kv_cache, num_layers, kernel, cute_dsl_decode_counter): """Long decode-only MLA run that crosses paged-KV block boundaries mid-decode. Reproduction for the DeepSeek-V3 E2E degeneration: short prompt, long