diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 48f3254537d8..f68f66608f7f 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,3 +1,198 @@ +diff --git a/_compile_cache.py b/_compile_cache.py +index b374a65311402982d818d80e9145ab40ed9727a1..cdfa76294f2d854117db62ad36724254006f8cb8 100644 +--- a/_compile_cache.py ++++ b/_compile_cache.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -15,10 +16,13 @@ + """Persistent compilation shared by the PrimTS attention adapters.""" + + from collections.abc import Callable ++import contextlib ++import contextvars + import dataclasses + import enum + import functools + import hashlib ++import os + from pathlib import Path + from typing import Any + +@@ -64,6 +68,128 @@ def _compile_cache_key(spec: object, options: str, extra: tuple = ()) -> str | N + return hashlib.sha256(encoded.encode()).hexdigest() + + ++# Warm-up precompile modes. Under ``precompile_only()`` the attention kernels ++# are compiled through ``JitSpecCuteDsl.compile_and_persist()`` (the ++# ``cute.compile`` runs outside the module-level file lock that ++# ``build_and_load()`` holds for a whole build, so sibling processes compile ++# different kernels concurrently), committed to disk, and not loaded: the ++# serving process loads them afterwards. Under ``probe_cached()`` nothing is ++# compiled; a kernel missing from the cache raises ``KernelNotCached``. Both ++# modes return a placeholder in place of the kernel. ++_PRECOMPILE_MODE: contextvars.ContextVar[str | None] = contextvars.ContextVar( ++ "prims_ts_precompile_mode", default=None ++) ++_SKIP_BUSY: contextvars.ContextVar[bool] = contextvars.ContextVar( ++ "prims_ts_precompile_skip_busy", default=False ++) ++ ++ ++class KernelNotCached(Exception): ++ """Raised under ``probe_cached()`` for a kernel that is not on disk.""" ++ ++ ++class KernelBusy(Exception): ++ """Raised under ``precompile_only(skip_busy=True)`` while another process compiles the kernel.""" ++ ++ ++@contextlib.contextmanager ++def precompile_only(*, skip_busy: bool = False): ++ """Compile and persist kernels without loading them (see module note). ++ ++ With ``skip_busy`` a kernel whose per-kernel lock another process holds ++ raises ``KernelBusy`` instead of waiting, so a worker can defer it and ++ compile something else first. ++ """ ++ token = _PRECOMPILE_MODE.set("persist") ++ busy_token = _SKIP_BUSY.set(skip_busy) ++ try: ++ yield ++ finally: ++ _SKIP_BUSY.reset(busy_token) ++ _PRECOMPILE_MODE.reset(token) ++ ++ ++@contextlib.contextmanager ++def probe_cached(): ++ """Report, without compiling, whether every kernel a plan needs is on disk.""" ++ token = _PRECOMPILE_MODE.set("probe") ++ try: ++ yield ++ finally: ++ _PRECOMPILE_MODE.reset(token) ++ ++ ++def expected_cache_meta() -> dict: ++ """The on-disk module metadata this process would write (arch, DSL version, source hash).""" ++ from flashinfer.jit.cute_dsl_core import ( ++ _get_compile_arch, ++ _get_cute_dsl_version, ++ _hash_source_files, ++ ) ++ ++ return { ++ "arch": _get_compile_arch(), ++ "cute_dsl_version": _get_cute_dsl_version(), ++ "source_sha256": _hash_source_files(tuple(_kernel_source_files())), ++ } ++ ++ ++def _unloaded_kernel(*args, **kwargs): ++ raise RuntimeError( ++ "PrimTS kernel was compiled in precompile-only mode and never loaded; " ++ "load it through a normal plan() outside precompile_only()." ++ ) ++ ++ ++def is_placeholder(compiled) -> bool: ++ """True for the stand-in returned under the precompile modes.""" ++ return compiled is _unloaded_kernel ++ ++ ++def _make_spec(module_name: str, kernel_name: str, compile_fn, key_files): ++ from flashinfer.jit.cute_dsl_core import JitSpecCuteDsl, _hash_source_files ++ ++ return JitSpecCuteDsl(module_name, kernel_name, compile_fn, _hash_source_files(tuple(key_files))) ++ ++ ++def _precompile_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: ++ """Compile one kernel into the on-disk cache unless it is already there. ++ ++ One compile per kernel per workspace: a sibling that wants the same kernel ++ waits on the per-kernel lock (or raises ``KernelBusy`` under ``skip_busy``) ++ and then finds it on disk. The lock lives in FlashInfer's tmp directory, ++ not in the module directory: creating that directory before the first ++ export would let ``compile_and_persist()`` of a sibling read a missing ++ meta.json as a stale module and wipe the exports already there. ++ """ ++ from filelock import FileLock, Timeout ++ from flashinfer.jit.core import get_tmpdir ++ ++ spec = _make_spec(module_name, kernel_name, compile_fn, key_files) ++ if spec.is_compiled: ++ return ++ lock = FileLock(get_tmpdir() / f"{spec.module_dir_name}.{spec.kernel_name}.lock", thread_local=False) ++ if _SKIP_BUSY.get(): ++ try: ++ lock.acquire(timeout=0) ++ except Timeout: ++ raise KernelBusy(spec.name) from None ++ else: ++ lock.acquire() ++ try: ++ if spec.is_compiled: ++ return ++ spec.compile_and_persist() ++ finally: ++ lock.release() ++ ++ ++def _probe_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: ++ spec = _make_spec(module_name, kernel_name, compile_fn, key_files) ++ if not spec.is_compiled: ++ raise KernelNotCached(spec.name) ++ ++ + @functools.cache + def _kernel_source_files() -> tuple[str, ...]: + # Family adapters share masks, reductions, sparse/Sage preparation and +@@ -82,16 +208,26 @@ def build_and_load_attention_kernel( + extra: tuple = (), + ) -> Any: + """Build or load under the caller's selected CUDA device context.""" +- from ...jit.cute_dsl_core import ( ++ from flashinfer.jit.cute_dsl_core import ( + build_and_load_cute_dsl_kernel, + cute_dsl_cache_disabled, + ) + ++ mode = _PRECOMPILE_MODE.get() ++ if mode is not None and os.environ.get("FLASHINFER_DISABLE_JIT"): ++ raise RuntimeError("PrimTS warm-up precompile cannot compile with FLASHINFER_DISABLE_JIT set.") + if cute_dsl_cache_disabled(): ++ if mode is not None: ++ raise RuntimeError( ++ "PrimTS warm-up precompile needs the on-disk CuTe DSL cache " ++ "(FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 is set)." ++ ) + return compile_fn() + + key = _compile_cache_key(spec, options, extra) + if key is None: ++ if mode is not None: ++ raise RuntimeError("PrimTS warm-up precompile cannot key this kernel for the on-disk cache.") + return compile_fn() + import torch + +@@ -99,6 +235,14 @@ def build_and_load_attention_kernel( + # host-selected tuning defaults. The module metadata keys the target; + # the specialization must also distinguish the physical capability. + major, minor = torch.cuda.get_device_capability() ++ if mode == "persist": ++ _precompile_kernel( ++ f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files() ++ ) ++ return _unloaded_kernel ++ if mode == "probe": ++ _probe_kernel(f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files()) ++ return _unloaded_kernel + return build_and_load_cute_dsl_kernel( + f"prims_ts_{family}", + f"{kind}_sm{major}{minor}_{key}", diff --git a/block_sparse.py b/block_sparse.py index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe2532a1c7 100644 --- a/block_sparse.py @@ -61,10 +256,10 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d1bcc2020 100644 +index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9592697ad 100644 --- a/context.py +++ b/context.py -@@ -45,7 +45,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional +@@ -48,7 +48,7 @@ from ._hardware import get_max_active_clusters import torch @@ -73,7 +268,7 @@ index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d if TYPE_CHECKING: -@@ -490,7 +490,7 @@ def _validate_device(device: torch.device) -> int: +@@ -493,7 +493,7 @@ def _validate_device(device: torch.device) -> int: # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. if capability == (10, 7): @@ -83,7 +278,7 @@ index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d require_cute_dsl_arch(device_index) return device_index diff --git a/decode.py b/decode.py -index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329fc4a3c11 100644 +index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac722fed527 100644 --- a/decode.py +++ b/decode.py @@ -1,3 +1,4 @@ @@ -91,9 +286,9 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -24,11 +25,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast +@@ -26,11 +27,7 @@ import torch - import torch + from ._compile_cache import build_and_load_attention_kernel -from flashinfer.api_logging import flashinfer_experimental_api -from flashinfer.trace.templates.attention import ( @@ -104,7 +299,7 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 from . import _q_token_kv_block_sparse_policy as _sparse_policy from ._block_sparse.common import _num_sparse_pattern_heads -@@ -789,7 +786,7 @@ def _validate_runtime_device(device: torch.device) -> None: +@@ -780,7 +777,7 @@ def _validate_runtime_device(device: torch.device) -> None: # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. if capability == (10, 7): @@ -113,7 +308,7 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 require_cute_dsl_arch(device_index) -@@ -4048,7 +4045,7 @@ class BatchDecodePagedTSWrapper: +@@ -4075,7 +4072,7 @@ class BatchDecodePagedTSWrapper: # previous complete plan revision usable. self._plan_state = candidate @@ -122,7 +317,7 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 def run( self, q: torch.Tensor, -@@ -4230,7 +4227,7 @@ class BatchDecodePagedTSWrapper: +@@ -4257,7 +4254,7 @@ class BatchDecodePagedTSWrapper: ) @@ -132,10 +327,15 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py -index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633428e32c3 100644 +index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2fe30a922 100644 --- a/kernels/fmha_context/fmha_resources.py +++ b/kernels/fmha_context/fmha_resources.py -@@ -1199,8 +1199,6 @@ class GmemQKVResource(MemoryResource): +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -1199,8 +1200,6 @@ class GmemQKVResource(MemoryResource): if cutlass.const_expr(self.cfg.has_uniform_varlen): seqlen_k = Int32(self.cfg.uniform_seq_len_k) else: @@ -144,7 +344,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 seqlen_k = _load_runtime_seq_len_kv( self.seqlens_kv, self.max_seq_len_kv, batch_coord ) -@@ -1223,8 +1221,6 @@ class GmemQKVResource(MemoryResource): +@@ -1223,8 +1222,6 @@ class GmemQKVResource(MemoryResource): kv_request_begin = batch_coord * Int32(self.block_table_row_stride) kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1) else: @@ -153,7 +353,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1300,13 +1296,9 @@ class GmemQKVResource(MemoryResource): +@@ -1300,13 +1297,9 @@ class GmemQKVResource(MemoryResource): kv_request_begin = batch_coord * Int32(self.block_table_row_stride) kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1) else: @@ -167,7 +367,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1743,8 +1735,6 @@ class SmemPageOffsetsKvResource(MemoryResource): +@@ -1743,8 +1736,6 @@ class SmemPageOffsetsKvResource(MemoryResource): kv_request_begin: Int32, kv_page_idx_ub: Int32, ) -> None: @@ -176,7 +376,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 cfg = self.cfg # Context's K/V tile index is kv_tile_start + loop_offset; for V the # producer runs one tile ahead in TAIL, so reuse the same expression -@@ -1835,8 +1825,6 @@ class SmemPageOffsetsKvResource(MemoryResource): +@@ -1835,8 +1826,6 @@ class SmemPageOffsetsKvResource(MemoryResource): tile_offset: cutlass.Constexpr[int] = 0, ) -> cutlass.Array: """Retain one tile's page IDs after its SMEM window is released.""" @@ -185,7 +385,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 tile_idx = _resolve_kv_tile_idx_context( stage_info, kv_tile_start, tile_offset=tile_offset ) -@@ -5770,3 +5758,14 @@ class GmemOResource(MemoryResource): +@@ -5770,3 +5759,14 @@ class GmemOResource(MemoryResource): prims.cp_async_bulk_commit_group() if cutlass.const_expr(self.cfg.gmem_o_store_wait_after_write): prims.cp_async_bulk_wait_group(0, read=True) @@ -200,8 +400,26 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 + _load_runtime_seq_len_kv, + _resolve_kv_tile_idx_context, +) +diff --git a/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py b/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py +index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a32c5289e 100644 +--- a/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py ++++ b/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -36,7 +37,7 @@ from cutlass.experimental.task_scheduling.resources import ( + producer_work, + ) + +-from ......cute_dsl.availability import is_rubin_cute_dsl_available ++from flashinfer.cute_dsl.availability import is_rubin_cute_dsl_available + from ..fmha_decode_config import FmhaDecodeConfig + from .helpers_common import ( + _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2eef74a1a4b 100644 +index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e231c98a52 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -209,19 +427,79 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -21,10 +22,9 @@ from typing import Any, Literal, Optional, cast +@@ -27,11 +28,7 @@ from ._hardware import get_max_active_clusters import torch -from flashinfer.api_logging import flashinfer_experimental_api -+from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api - from flashinfer.trace.templates.attention import ( - prims_ts_decode_mla_one_shot_trace_dispatch, +-from flashinfer.trace.templates.attention import ( +- prims_ts_decode_mla_one_shot_trace_dispatch, - prims_ts_decode_mla_wrapper_trace_dispatch, - ) +-) ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api from .decode import ( -@@ -1576,7 +1576,7 @@ class BatchMLADecodePagedTSWrapper: + _WorkspaceSection, +@@ -1047,11 +1044,35 @@ def _make_mla_decode_compile_spec( + ) + + +-@functools.cache +-def _get_compiled_mla_decode( +- compile_spec: _MLADecodeCompileSpec, +-): +- """Compile and cache one batch-dynamic TS MLA topology.""" ++_COMPILED_MLA_DECODE: dict = {} ++ ++ ++def _get_compiled_mla_decode(compile_spec: _MLADecodeCompileSpec): ++ """Compile and cache one batch-dynamic TS MLA topology. ++ ++ A manual memo rather than ``functools.cache``: under the warm-up precompile ++ modes of ``_compile_cache`` the build returns a placeholder (probe / persist ++ only), which must not be remembered for the serving launches that follow. ++ """ ++ compiled = _COMPILED_MLA_DECODE.get(compile_spec) ++ if compiled is None: ++ compiled = _compile_mla_decode(compile_spec) ++ if not _compile_cache_is_placeholder(compiled): ++ _COMPILED_MLA_DECODE[compile_spec] = compiled ++ return compiled ++ ++ ++_get_compiled_mla_decode.cache_clear = _COMPILED_MLA_DECODE.clear # type: ignore[attr-defined] ++ ++ ++def _compile_cache_is_placeholder(compiled) -> bool: ++ from ._compile_cache import is_placeholder ++ ++ return is_placeholder(compiled) ++ ++ ++def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): ++ """Build (or load from the on-disk cache) one batch-dynamic TS MLA topology.""" + + import cutlass + import cutlass.cute as cute +@@ -1459,11 +1480,12 @@ def _launch_mla_decode( + runtime.sparse_inputs, + runtime.bmm1_scale, + runtime.bmm2_scale, +- **( +- {"output_quant_args": runtime.output_quant_args} +- if runtime.output_quant_args is not None +- else {} +- ), ++ # Passed positionally and also when None. The kernel was compiled with ++ # ``output_quant_args`` in its signature (``None`` for unquantized ++ # output), and a kernel reloaded from the persistent compile cache is ++ # a positional TVM-FFI entry with no Python defaults: it rejects both ++ # a missing parameter and a keyword argument. ++ runtime.output_quant_args, + ) + return runtime.out + +@@ -1637,7 +1659,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) @@ -230,8 +508,241 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee def run( self, query: torch.Tensor, +@@ -1742,7 +1764,7 @@ class BatchMLADecodePagedTSWrapper: + ) + + +-@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) ++@flashinfer_experimental_api + def batch_mla_decode_with_paged_kv_cache( + query: torch.Tensor, + kv_cache: torch.Tensor, +@@ -2099,6 +2121,7 @@ def _prepare_sparse_output_quant( + ) + + ++@functools.lru_cache(maxsize=1024) + def _resolve_sparse_mla_plan( + device, + batch_size, +@@ -2119,17 +2142,17 @@ def _resolve_sparse_mla_plan( + ): + """Resolve sparse geometry and scratch sizing without allocation/compilation.""" + import cutlass +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_latency_1cta.config import ( ++ from .kernels.mla_decode.throughput_latency_1cta.config import ( + MlaProfile, + compute_workspace_size, + ) +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_latency_1cta.kernel import ( ++ from .kernels.mla_decode.throughput_latency_1cta.kernel import ( + ThroughputLatencyMlaDecodeTs, + ) +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_2cta.config import ( ++ from .kernels.mla_decode.throughput_2cta.config import ( + compute_workspace_size as workspace_2cta, + ) +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_2cta.kernel import ( ++ from .kernels.mla_decode.throughput_2cta.kernel import ( + MlaDecodeTs, + ) + +@@ -2264,6 +2287,9 @@ def _resolve_sparse_mla_plan( + return device, profile, spec, sections, byte_end, capacity, output_quant + + ++# Prepared launches kept per plan (one per layer sharing the plan; DeepSeek-V4 has 61). ++_PREPARED_LAUNCH_LIMIT = 256 ++ + class BatchSparseMLADecodePagedTSWrapper: + """Plan native BF16/E4M3 attention, then bind one or two source index lists. + +@@ -2297,6 +2323,11 @@ class BatchSparseMLADecodePagedTSWrapper: + self._workspace_buffer = workspace_buffer + self._state = None + self._scalar_cache = {} ++ # Prepared launches for the unfused validate=False path: everything that ++ # is static for one (plan, pools, scales, sinks) tuple, keyed by the pools' ++ # storage and the sinks address. One plan serves every layer of a model, ++ # and each layer brings its own sinks, so one entry per layer. ++ self._prepared: dict[tuple, dict] = {} + self.workspace_size_bytes = 0 + + def plan( +@@ -2386,6 +2417,7 @@ class BatchSparseMLADecodePagedTSWrapper: + kernel_signature=_mla_kernel_compile_signature(static_kernel), + ) + ) ++ self._prepared.clear() + self._state = dict( + device=device, + batch=batch_size, +@@ -2608,6 +2640,38 @@ class BatchSparseMLADecodePagedTSWrapper: + state = self._state + if state is None: + raise RuntimeError("plan() must be called before run()") ++ lse_arg = lse ++ prepared_key = ( ++ kv_cache.data_ptr(), ++ kv_cache.shape, ++ kv_cache.stride(), ++ kv_cache.dtype, ++ None ++ if extra_kv_cache is None ++ else ( ++ extra_kv_cache.data_ptr(), ++ extra_kv_cache.shape, ++ extra_kv_cache.stride(), ++ extra_kv_cache.dtype, ++ ), ++ None if sinks is None else sinks.data_ptr(), ++ softmax_scale, ++ output_scale, ++ ) ++ if ( ++ not validate ++ and out is not None ++ and lse is None ++ and qo_indptr is None ++ and token_positions is None ++ and cos_sin_cache is None ++ and out_scale is None ++ ): ++ prepared = self._prepared.get(prepared_key) ++ if prepared is not None: ++ return self._run_prepared( ++ prepared, query, indices, lengths, extra_indices, extra_lengths, out ++ ) + _validate_query( + query, + packed_query=state["packed"], +@@ -2781,10 +2845,124 @@ class BatchSparseMLADecodePagedTSWrapper: + if static_scales + else state["compiled"], + ) ++ if ( ++ rows ++ and quant is None ++ and static_scales ++ and not state["packed"] ++ and not state["return_lse"] ++ and state["max_q"] == 1 ++ and lse_arg is None ++ ): ++ # Static for this (plan, pools, scales): pool views, scalars, sinks, split, ++ # workspace. The next validate=False launch with the same pools re-forms ++ # only the query / index / output views (see _run_prepared). ++ if len(self._prepared) >= _PREPARED_LAUNCH_LIMIT: ++ self._prepared.clear() ++ self._prepared[prepared_key] = dict( ++ primary=primary, ++ extra=extra, ++ primary_rows=primary.shape[0], ++ c_latent=primary.permute(1, 2, 0), ++ c_rope=extra.permute(1, 2, 0), ++ scalars=(scale_tensors[0], scale_tensors[1], sinks), ++ bmm=(bmm1_scale, bmm2_scale), ++ ks=state["ks"], ++ kc=state["kc"], ++ dummy_indices=state["dummy_indices"], ++ default_cl=state["default_cl"], ++ ) + if quant is not None: + return (out, out_scale, lse) if state["return_lse"] else (out, out_scale) + return (out, lse) if state["return_lse"] else out + ++ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_lengths, out): ++ """Launch with the static arguments recorded by the last full run(). ++ ++ Only the per-call tensors are checked and viewed: query and output ++ geometry, index/length shapes, dtypes and strides. Pool geometry, ++ scales and sinks were validated by the full path that prepared this ++ launch, and the caller passed validate=False. Byte-identical to the ++ full path for the same arguments. ++ """ ++ state = self._state ++ rows, heads, device = state["batch"], state["heads"], state["device"] ++ if ( ++ query.shape != (rows, 1, heads, 512) ++ or query.dtype != state["dtype"] ++ or query.device != device ++ or not query.is_contiguous() ++ or query.data_ptr() % 16 ++ ): ++ raise ValueError( ++ "prepared launch: query must be a contiguous, 16-byte aligned " ++ "[batch, 1, heads, 512] tensor of the planned dtype on the planned device" ++ ) ++ if ( ++ out.numel() != rows * heads * 512 ++ or out.dtype != torch.bfloat16 ++ or out.device != device ++ or not out.is_contiguous() ++ or out.data_ptr() % 16 ++ ): ++ raise ValueError( ++ "prepared launch: out must be a contiguous, 16-byte aligned BF16 tensor " ++ "holding one 512-wide row per query head" ++ ) ++ if ( ++ indices.shape != (rows, prepared["ks"]) ++ or indices.dtype != torch.int32 ++ or indices.device != device ++ or indices.stride(-1) != 1 ++ or lengths.shape != (rows,) ++ or lengths.dtype != torch.int32 ++ or lengths.device != device ++ or not lengths.is_contiguous() ++ ): ++ raise ValueError("prepared launch: invalid primary indices/lengths") ++ if prepared["kc"]: ++ if ( ++ not isinstance(extra_indices, torch.Tensor) ++ or not isinstance(extra_lengths, torch.Tensor) ++ or extra_indices.shape != (rows, prepared["kc"]) ++ or extra_indices.dtype != torch.int32 ++ or extra_indices.device != device ++ or extra_indices.stride(-1) != 1 ++ or extra_lengths.shape != (rows,) ++ or extra_lengths.dtype != torch.int32 ++ or extra_lengths.device != device ++ or not extra_lengths.is_contiguous() ++ ): ++ raise ValueError("prepared launch: invalid extra indices/lengths") ++ ci, cl = extra_indices, extra_lengths ++ else: ++ if extra_indices is not None or extra_lengths is not None: ++ raise ValueError("prepared launch: extra indices without a planned extra pool") ++ ci, cl = prepared["dummy_indices"][:rows], prepared["default_cl"][:rows] ++ buffers = state["buffers"] ++ s0, s1, sinks = prepared["scalars"] ++ q_kernel = query.permute(2, 3, 1, 0) ++ lse_kernel = buffers["public_lse"][:rows].view(rows, 1, heads).permute(2, 1, 0) ++ bmm1, bmm2 = prepared["bmm"] ++ state["compiled_static"]( ++ q_kernel, ++ q_kernel, ++ prepared["c_latent"], ++ prepared["c_rope"], ++ indices.transpose(0, 1), ++ out.view(rows, 1, heads, 512).permute(2, 3, 1, 0), ++ lse_kernel, ++ buffers["core"] if buffers["core"].numel() else None, ++ state["splits"], ++ lengths, ++ None, ++ (indices, ci, lengths, cl, s0, s1, sinks), ++ bmm1, ++ bmm2, ++ None, ++ ) ++ return out ++ + + def get_prims_ts_sparse_mla_decode_workspace_size(*plan_args, **plan_kwargs): + """Return workspace bytes using the same arguments as wrapper.plan(). diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py -index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c0788f9f3a4 100644 +index e2bbef13c0f552030875e6b433e0fcdb20686cf6..dff16046cf7cf9ff03950043cf462ac28de6c670 100644 --- a/q_token_kv_block_sparse_metadata.py +++ b/q_token_kv_block_sparse_metadata.py @@ -1,3 +1,4 @@ @@ -239,7 +750,7 @@ index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c07 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -40,9 +41,9 @@ from dataclasses import dataclass +@@ -41,9 +42,9 @@ from dataclasses import dataclass from typing import Any, Literal, Optional import torch @@ -251,7 +762,7 @@ index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c07 from ._block_sparse.common import _num_sparse_pattern_heads _Q_TOKEN_KV_BLOCK_SPARSE_MIN_BLOCK_SIZE = 4 -@@ -62,7 +63,7 @@ _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT = 256 +@@ -63,7 +64,7 @@ _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT = 256 def _get_prims_ts_q_token_kv_block_sparse_metadata_module() -> Any: """Build and cache the CUDA C++ QToken-KvBlock-Sparse-Attention metadata implementation.""" diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index ea34d65b41be..e19632edea75 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -13,16 +13,16 @@ vendors: patch_digest: sha256:bc9b96a69873e0ddfcba66ab899c666cf52766976f24d850048d2a0bf8d967ff digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac flashinfer-prims-ts: - url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev-20261012 - commit: 0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf + url: https://github.com/yunruis/flashinfer.git + branch: trtllm-prims-ts-dev-20261023 + commit: aa709061230a58858375c513056e7733fa74294a source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:1ec340aac4fea9c77b9e9d956b1676c09cff5081053ef1e41f962ed38a6f728c - digest: sha256-tree-v1:e4129294b2b26295581cd9cbf1ac5ac4d7eec2ec65f0108626692ddf8a036f96 + patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07 + digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 8f949a1aeb75..4aa228406737 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -8,6 +8,11 @@ All published functionality in the Release Notes has been fully tested and verif ### Key Features and Enhancements +- **DeepSeek-V4 generation on B200/GB200 (SM100) now runs on the PrimTS native sparse MLA FMHA by default.** The new `prims_ts_sparse_mla` FMHA library (a CuTe DSL kernel vendored from FlashInfer) serves the DeepSeek-V4 sparse MLA generation phase ahead of the TRTLLM-Gen path on SM100; it covers BF16 and FP8 KV caches, the fused inverse-RoPE + FP8 output epilogue, speculative decoding (MTP) rows and CUDA graph replay. SM103 (GB300), `fp8_ds_mla` and NVFP4 KV caches keep the TRTLLM-Gen path. Opt out with `TLLM_FMHA_LIBS=-prims_ts_sparse_mla`. + - **First-launch compile cost.** The kernels are JIT-compiled by the CuTe DSL and cached on disk under the FlashInfer workspace. The first start on a host (or after a FlashInfer / CuTe DSL upgrade) compiles the kernel variants the deployment's generation row counts select during `LLM()` startup, about 5 to 10 minutes on B200 for DeepSeek-V4; later starts load the cache. Under `trtllm-llmapi-launch` and the default `LLM()` MPI session every rank has its own FlashInfer workspace (`TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=1`), so each rank compiles its own copy; set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0` to share one cache between the ranks of a host. + - **Warm-up controls.** `TLLM_PRIMTS_WARM_WORKERS` sets the compile subprocesses per rank (default 4, reduced so that the ranks on a host, counted from the model's parallel mapping, start at most 16 together; `0` compiles in-process), `TLLM_PRIMTS_WARM_TIMEOUT` bounds the wait for them (default 1800 s), and `TLLM_PRIMTS_WARM_STRICT=1` turns a worker failure into an error instead of an in-process retry. Each worker imports torch and the CuTe DSL and opens a CUDA context, so lower the worker count on nodes with little host memory. + - **`FLASHINFER_DISABLE_JIT=1`** now also refuses to compile these kernels (with the on-disk CuTe DSL cache enabled, its default): startup then fails on a variant missing from the cache instead of falling back to TRTLLM-Gen, so a deployment that sets it must pre-populate the cache (one start without the variable) or opt out of the library. + ### Infrastructure Changes ### API Changes diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index f9b52aef8a9d..1eff1545c3e2 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -479,7 +479,11 @@ the selected library. it may add host overhead; use `TLLM_FMHA_LIBS=+prims_ts` to add it to the defaults or `TLLM_FMHA_LIBS=fallback` to force the fallback path. Generic block-sparse PrimTS remains enabled by default because a dense fallback cannot -preserve its routing semantics. Delta entries update the +preserve its routing semantics. The PrimTS sparse MLA library +(`prims_ts_sparse_mla`) is also enabled by default: it serves only DeepSeek-V4 +generation on SM100 and sits ahead of the TRTLLM-Gen libraries, so DeepSeek-V4 +decode on B200/GB200 runs on it unless `TLLM_FMHA_LIBS=-prims_ts_sparse_mla` +removes it. Delta entries update the default membership and follow canonical registry order, while an exact list preserves the user-specified order. Each FMHA library exposes `is_available()` for module/static environment checks and `is_supported()` for per-forward @@ -552,6 +556,10 @@ The FMHA package is split by role: With an FP8 KV cache it passes `mla_bmm1_scale` / `mla_bmm2_scale` to the kernel as device tensors (element 0 of each), so the scales need no host read and stay valid under CUDA graph replay. +- `fmha/prims_ts_sparse_mla.py` runs DeepSeek-V4 generation on the vendored + PrimTS two-source sparse MLA kernel (SM100 only; SM103 stays on TRTLLM-Gen); + it is in the default library set and `TLLM_FMHA_LIBS=-prims_ts_sparse_mla` + removes it. - `fmha/prims_ts_block_sparse.py` adapts generic block-sparse requests to the vendored PrimTS contiguous and paged wrappers. Paged generation passes a live, zero-copy 2D K-page-table view with its TRT-LLM padded row stride; it diff --git a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py index 0fc278ce8c1c..2e1ee8497bc9 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py @@ -24,6 +24,7 @@ from .msa_prefill import MsaPrefillFmha from .phased import FmhaParams, PhasedFmha from .prims_ts import PrimsTSFmha +from .prims_ts_sparse_mla import PrimsTSSparseMlaFmha from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes from .triton_custom_mask import TritonCustomMaskFmha @@ -44,6 +45,7 @@ "MsaPrefillFmha", "PhasedFmha", "PrimsTSFmha", + "PrimsTSSparseMlaFmha", "TritonCustomMaskFmha", "get_enabled_fmha_lib_classes", ] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index c9afff66e1d8..5f86c100f8c7 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -311,6 +311,20 @@ def __init__(self, attn: TrtllmAttention) -> None: for fmha_cls in get_enabled_fmha_lib_classes(): if fmha_cls.is_available(attn): self.fmha_libs.append(fmha_cls(attn)) + self._layer_idx = getattr(attn, "layer_idx", None) + + def _announce(self, fmha: Fmha, cache_key: _FmhaCacheKey) -> None: + # One line per library per layer is enough to tell which kernel + # serves a model; the logger dedups on the key. + logger.info_once( + f"FMHA layer {self._layer_idx}: {type(fmha).__name__} serves " + f"{cache_key.attention_mask_type.name} requests " + f"(context_batch<={cache_key.context_batch_size}, " + f"generation_batch<={cache_key.generation_batch_size}, " + f"generation_seq_len_q<={cache_key.generation_seq_len_q}, " + f"output_dtype={cache_key.output_dtype}).", + key=f"fmha_serves:{self._layer_idx}:{type(fmha).__name__}", + ) def _make_cache_key( self, @@ -422,6 +436,7 @@ def select( if fmha is None: return None self._cache[cache_key] = fmha + self._announce(fmha, cache_key) if self._cache_sanity_check_enabled: self._cache_inputs[cache_key] = _snapshot_fmha_cache_inputs( q, k, v, metadata, forward_args diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py new file mode 100644 index 000000000000..f293388f6a1f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parallel precompilation of PrimTS sparse MLA kernel variants. + +``PrimsTSSparseMlaFmha._warm_kernel_variants`` plans every kernel variant a +deployment can select. Each variant is one ``cute.compile`` (tens of seconds), +and FlashInfer's on-disk CuTe DSL cache serializes the builds of one op family +behind a file lock, so N ranks warming the same variants finish no sooner than +one rank compiling them all in sequence. This module fans the variants out to +worker subprocesses that compile through ``_compile_cache.precompile_only()`` +(``cute.compile`` outside the lock, the artifact committed under it); the +serving process then loads every variant from disk. + +Workers are plain interpreter subprocesses rather than ``multiprocessing`` +children: importing ``tensorrt_llm`` initializes MPI, so a worker must not +inherit the launcher's MPI environment, and a subprocess gets a fully +controlled one. Worker ``i`` of ``N`` starts with payload slots ``i, i+N, ...`` +and then goes through the rest of the list, deferring kernels another process +is compiling (per-kernel lock) and finishing with a blocking pass, so every +worker ends with every variant on disk whether the workspace is private to the +rank (TRT-LLM's default) or shared between ranks. Failures are best-effort: +the caller falls back to compiling in-process. + +Host load. Every worker is an interpreter that imports torch and the CuTe DSL, +opens a CUDA context on its rank's GPU and runs ``cute.compile`` (one to a few +GB of host memory and one busy core each). Under the default launcher each +rank owns a private FlashInfer workspace, so the ranks of one host do not +share compiles and every rank fans out on its own. The default worker count +therefore follows the number of ranks on the host (from the model's +``Mapping``, or the launcher's environment): at most 4 per rank and at most +``_HOST_COMPILE_BUDGET`` (16) workers per host in total, so a TP=8 node runs 2 +workers per rank, not 32 subprocesses at once. Set +``TLLM_PRIMTS_WARM_WORKERS`` explicitly to trade startup time against host +memory and CPU on a given node. + +Environment: + TLLM_PRIMTS_WARM_WORKERS subprocess count per rank; default 4 capped so + that all ranks on the host start at most 16; + 0 or 1 disables. + TLLM_PRIMTS_WARM_TIMEOUT seconds to wait for the workers, default 1800. + TLLM_PRIMTS_WARM_STRICT 1 turns a worker failure into an error (tests). +""" + +from __future__ import annotations + +import json +import os +import subprocess # nosec B404 - runs sys.executable workers with a fixed argv, no shell +import sys +import tempfile +import time +from typing import Any, Sequence + +from tensorrt_llm.logger import logger + +_WORKER_BOOTSTRAP = """ +import json, os, sys +_path = os.environ.get("TLLM_PRIMTS_PRECOMPILE_SYS_PATH") +if _path: + sys.path[:] = json.loads(_path) +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile import _main +sys.exit(_main(sys.argv)) +""" + + +_MAX_DEFAULT_WORKERS = 4 +_HOST_COMPILE_BUDGET = 16 # default workers of all ranks on one host, together +_LOCAL_RANK_ENV = ("OMPI_COMM_WORLD_LOCAL_SIZE", "RAY_LOCAL_WORLD_SIZE", "SLURM_NTASKS_PER_NODE") + + +def local_rank_count(mapping) -> int: + """Ranks of this job on this host, without any MPI call. + + From the model's ``Mapping`` (``min(world_size, gpus_per_node)``) when the + caller has one, else from the launcher's environment, else 1. No + communicator query: this runs inside the first forward, where a collective + would hang if the ranks reached it at different points. + """ + if mapping is not None: + world = int(getattr(mapping, "world_size", 0) or 0) + per_node = int(getattr(mapping, "gpus_per_node", 0) or 0) + if world > 0 and per_node > 0: + return max(1, min(world, per_node)) + if world > 0: + return world + for name in _LOCAL_RANK_ENV: + raw = os.environ.get(name) + if raw: + try: + return max(1, int(raw)) + except ValueError: + continue + return 1 + + +def default_worker_count(local_ranks: int = 1) -> int: + """Per-rank default: 4, lowered so the ranks on the host start at most 16 workers.""" + return max(1, min(_MAX_DEFAULT_WORKERS, _HOST_COMPILE_BUDGET // max(1, int(local_ranks)))) + + +def worker_count(num_payloads: int, *, local_ranks: int = 1) -> int: + """Subprocess count from ``TLLM_PRIMTS_WARM_WORKERS``; ``0``/``1`` disables. + + Without the variable the default comes from ``default_worker_count()`` + for ``local_ranks`` ranks on this host (see the module note on host load). + Zero when the on-disk cache or JIT compilation is disabled: a worker could + then compile but keep nothing. + """ + if os.environ.get("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "0") == "1" or os.environ.get( + "FLASHINFER_DISABLE_JIT" + ): + return 0 + raw = os.environ.get("TLLM_PRIMTS_WARM_WORKERS") + if raw is None: + workers = default_worker_count(local_ranks) + else: + try: + workers = int(raw) + except ValueError: + workers = default_worker_count(local_ranks) + logger.warning( + f"Ignoring TLLM_PRIMTS_WARM_WORKERS={raw!r}; expected an integer " + f"(using the default of {workers})." + ) + return max(0, min(workers, num_payloads)) + + +def strict() -> bool: + return os.environ.get("TLLM_PRIMTS_WARM_STRICT", "0") == "1" + + +def encode_payload(device_index: int, rows: int, heads: int, plan_kwargs: dict) -> dict[str, Any]: + """JSON-safe description of one ``BatchSparseMLADecodePagedTSWrapper.plan`` call.""" + kwargs = dict(plan_kwargs) + dtype = kwargs.pop("q_data_type") + kwargs["q_data_type"] = str(dtype).removeprefix("torch.") + return {"device_index": device_index, "rows": rows, "heads": heads, "plan_kwargs": kwargs} + + +def shard_order(payloads: Sequence[Any], slot: int, slots: int) -> list: + """Own shard first (``payloads[slot::slots]``), then the rest in list order.""" + mine = list(payloads[slot::slots]) + return mine + [p for p in payloads if not any(p is m for m in mine)] + + +def _worker_env() -> dict[str, str]: + from tensorrt_llm.llmapi.mpi_session import split_mpi_env + + env, _ = split_mpi_env() + env["TLLM_PRIMTS_WARM_WORKERS"] = "0" # workers never fan out again + env.setdefault("TLLM_LOG_LEVEL", "warning") + # The on-disk cache key includes the installed CuTe DSL distributions as + # enumerated along sys.path; a worker must see the parent's sys.path + # exactly (entries, order and duplicates) to write the same meta.json. + env["TLLM_PRIMTS_PRECOMPILE_SYS_PATH"] = json.dumps(sys.path) + return env + + +def run_precompile_workers(payloads: Sequence[dict[str, Any]], workers: int) -> bool: + """Compile ``payloads`` with ``workers`` subprocesses; True iff every worker succeeded. + + Worker output goes to temporary files (a pipe could fill while the worker + holds a kernel lock another worker waits on). The wait is bounded by + ``TLLM_PRIMTS_WARM_TIMEOUT``; on timeout or any error every worker is + killed and the caller compiles in-process. + """ + if workers <= 1 or not payloads: + return False + timeout_s = float(os.environ.get("TLLM_PRIMTS_WARM_TIMEOUT", "1800")) + env = _worker_env() + payload_json = json.dumps(list(payloads)) + procs: list[tuple[int, subprocess.Popen, Any]] = [] + ok = True + try: + for slot in range(workers): + out = tempfile.TemporaryFile(mode="w+", prefix="prims_ts_precompile_", suffix=".log") + proc = subprocess.Popen( + [sys.executable, "-c", _WORKER_BOOTSTRAP, payload_json, str(slot), str(workers)], + env=env, + stdout=out, + stderr=subprocess.STDOUT, + text=True, + ) + procs.append((slot, proc, out)) + deadline = time.monotonic() + timeout_s + for slot, proc, out in procs: + remaining = deadline - time.monotonic() + try: + proc.wait(timeout=max(remaining, 0.0)) + except subprocess.TimeoutExpired: + ok = False + logger.warning( + f"PrimTS precompile worker (slot {slot}) exceeded TLLM_PRIMTS_WARM_TIMEOUT=" + f"{timeout_s:g}s; the serving process compiles what is missing." + ) + continue + out.seek(0) + text = out.read() + if proc.returncode != 0: + ok = False + tail = "\n".join(text.strip().splitlines()[-15:]) + logger.warning( + f"PrimTS precompile worker (slot {slot}) exited with {proc.returncode}; " + f"the serving process compiles what is missing.\n{tail}" + ) + continue + meta_line = next( + ( + line + for line in text.splitlines() + if line.startswith("[prims-ts precompile] meta=") + ), + "", + ) + expected = str(_expected_cache_meta()) + if meta_line and meta_line[len("[prims-ts precompile] meta=") :] != expected: + ok = False + logger.warning( + f"PrimTS precompile worker (slot {slot}) wrote a different cache key than this " + f"process expects; its kernels will be recompiled.\n worker: {meta_line}\n parent: {expected}" + ) + else: + logger.debug(f"PrimTS precompile worker (slot {slot}) finished") + finally: + for _, proc, out in procs: + if proc.poll() is None: + proc.kill() + proc.wait() + out.close() + return ok + + +def _expected_cache_meta() -> dict: + from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache + + return _compile_cache.expected_cache_meta() + + +def _plan(mla_decode, payload: dict[str, Any], torch) -> None: + device = torch.device("cuda", int(payload["device_index"])) + torch.cuda.set_device(device) + kwargs = dict(payload["plan_kwargs"]) + kwargs["q_data_type"] = getattr(torch, kwargs["q_data_type"]) + mla_decode.BatchSparseMLADecodePagedTSWrapper().plan( + device, int(payload["rows"]), int(payload["heads"]), **kwargs + ) + + +def _main(argv: list[str]) -> int: + """Worker entry: compile ``payloads[slot::slots]`` first, then everything else.""" + payloads = json.loads(argv[1]) + slot, slots = (int(argv[2]), int(argv[3])) if len(argv) > 3 else (0, 1) + import torch + + from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache, mla_decode + + deferred = [] + for payload in shard_order(payloads, slot, slots): + try: + with _compile_cache.precompile_only(skip_busy=True): + _plan(mla_decode, payload, torch) + except _compile_cache.KernelBusy: + deferred.append(payload) + continue + print( + f"[prims-ts precompile] rows={payload['rows']} extra={payload['plan_kwargs'].get('max_extra_topk')} done", + flush=True, + ) + for payload in deferred: + with _compile_cache.precompile_only(): + _plan(mla_decode, payload, torch) + extra = payload["plan_kwargs"].get("max_extra_topk") + print( + f"[prims-ts precompile] rows={payload['rows']} extra={extra} done (waited)", flush=True + ) + print(f"[prims-ts precompile] meta={_compile_cache.expected_cache_meta()}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(_main(sys.argv)) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py new file mode 100644 index 000000000000..2cb531e21c3d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -0,0 +1,888 @@ +# 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. + +"""TRT-LLM FMHA adapter for the vendored PrimTS native sparse MLA kernels. + +Serves the DeepSeek-V4 generation phase on SM100 with the Blackwell +task-scheduled two-source sparse MLA decode kernel +(``BatchSparseMLADecodePagedTSWrapper``, flashinfer-ai/flashinfer#5434). The +sliding-window pool is the primary source and, for compressing layers, the +compressed pool is the extra source. Both lists feed one softmax, exactly like +the TRTLLM-Gen sparse MLA kernel the fallback library runs. + +The adapter consumes the dual-pool global indices that +``DeepseekV4TrtllmAttention.sparse_attn_predict`` already produces: SWA slots +are storage rows of the SWA pool and compressed slots are storage rows of the +layer's compress pool, so no page-table conversion happens here. The kernel +takes the two index lists and their live lengths directly; ``-1`` holes before a +live length are masked in the kernel only on the configurations that can produce +them (KV-cache offload, heuristic top-k), otherwise the plan promises valid +prefixes. The fused DSv4 epilogue (conjugate RoPE on ``D[448:512]`` + +1x128 block-scaled E4M3 output in the group-major ``[G, T, (H/G)*512]`` layout +with FP32 scales ``[G, (H/G)*4, pad4(T)]``) maps one-to-one onto the kernel's +``output_quant_block_size=128`` / ``out_scale_format="fp32"`` mode, so the +``o_a_proj`` BMM consumes the attention output without a separate quant +kernel. + +The library is enabled by default on SM100 and precedes the TRTLLM-Gen +fallback for the DeepSeek-V4 generation phase; ``TLLM_FMHA_LIBS=-prims_ts_sparse_mla`` +restores the previous library. The context phase stays on the fallback library. +SM103 (B300/GB300) is admitted by the kernels' architecture guard but not yet +signoff-qualified upstream, so it keeps the fallback for now. +""" + +from __future__ import annotations + +import math +import time +import weakref +from dataclasses import dataclass +from importlib import import_module +from importlib.metadata import PackageNotFoundError, version +from typing import TYPE_CHECKING, Optional + +import torch +from packaging.version import InvalidVersion, Version + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_sm_version +from tensorrt_llm.logger import logger + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( + BatchSparseMLADecodePagedTSWrapper, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import ( + DeepseekV4CacheManager, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +_MIN_CUTLASS_DSL_VERSION = Version("4.7.0") +# Same compiler-target gate as fmha/prims_ts.py: the kernels share one CuTe DSL +# task-scheduling stack, so a toolchain the dense libraries refuse is refused here too. +_MIN_CUTLASS_COMPILER_VERSION = "13.3" +# Process-level state shared by every library instance and metadata object: +# the static profiles whose kernel variants were warmed (compiled kernels are +# cached per process by the vendored wrapper, so one warm-up per profile is +# enough), and the index contracts that were validated with one eager launch. +_WARMED_PROFILES: set = set() +_VALIDATED_CONTRACTS: set = set() +_LATENT_DIM = 512 +# Prepared-sinks cache entries (one per layer; DeepSeek-V4 has 61). +_SINKS_CACHE_LIMIT = 512 +_ROPE_DIM = 64 +_FUSED_OUTPUT_BLOCK = 128 + + +def get_prims_ts_toolchain_unsupported_reason() -> Optional[str]: + """Return why the CuTe DSL task-scheduling toolchain cannot run PrimTS kernels.""" + try: + installed_version = Version(version("nvidia-cutlass-dsl")) + except (PackageNotFoundError, InvalidVersion): + return "nvidia-cutlass-dsl>=4.7 is required." + if installed_version < _MIN_CUTLASS_DSL_VERSION: + return ( + f"nvidia-cutlass-dsl>={_MIN_CUTLASS_DSL_VERSION} is required, got {installed_version}." + ) + try: + cutlass = import_module("cutlass") + compiler_version_supported = cutlass.target_version( + min_version=_MIN_CUTLASS_COMPILER_VERSION + ) + except Exception as error: # noqa: BLE001 - availability probes must fail closed + return f"could not query the active CUTLASS compiler version: {error}" + if not compiler_version_supported: + return f"the active CUTLASS compiler must target CUDA>={_MIN_CUTLASS_COMPILER_VERSION}." + try: + import_module("cutlass.experimental.task_scheduling") + except ImportError: + return "CUTLASS task scheduling is missing." + return None + + +@dataclass +class _SparsePlan: + """One planned wrapper for a static profile; every launch input comes from the metadata.""" + + wrapper: "BatchSparseMLADecodePagedTSWrapper" + + +class PrimsTSSparseMlaFmha(Fmha): + """DeepSeek-V4 generation-phase sparse MLA on the PrimTS Blackwell kernels.""" + + # The native kernel reads one KV head of 512 channels per selected token. + MAX_HEADS = 128 + # Plans live in ``TrtllmAttentionMetadata.fmha_plan_caches`` under this key, + # like the block-sparse and FlashInfer wrappers: every layer that runs with + # one metadata object plans each static profile once, and layers sharing a + # metadata object execute serially. All plans of a metadata object share + # one kernel workspace (under WORKSPACE_CACHE_KEY) that is grown to the + # largest profile planned so far (a plan keeps the buffer it was planned + # with alive until the metadata object is dropped); warming plans the + # largest row count first so the metadata that warms allocates it once. + PLAN_CACHE_KEY = "prims_ts_sparse_mla" + WORKSPACE_CACHE_KEY = "prims_ts_sparse_mla.workspace" + # Generation row counts scanned when warming a profile's kernel variants: + # every count up to WARM_ROW_SCAN (the split-KV schedule changes often + # at small counts), then a geometric tail up to the metadata's token + # budget. One plan is created per distinct kernel compile signature. + WARM_ROW_SCAN = 128 + WARM_ROW_TAIL = (192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096) + + def __init__(self, attn: "TrtllmAttention") -> None: + super().__init__(attn) + # FP8 dequantization scales are fixed per layer; the wrapper takes host + # scalars, so read them once during eager warmup, never under capture. + self._kv_dequant_scale: Optional[float] = None + # Pool-base-rooted views of the SWA and compress pools, owned by the + # manager they were built from. The owner is held weakly: a freed + # manager (or a new one at the same address) never resolves to a + # stale view of released pool memory. + self._checked_key = None + self._pool_views_owner: Optional[weakref.ReferenceType] = None + self._pool_views: dict[tuple[int, int], torch.Tensor] = {} + # One prepared FP32 sinks vector per source tensor (one per layer). + self._sinks_cache: dict[tuple[int, int, torch.dtype], torch.Tensor] = {} + + # ------------------------------------------------------------------ availability + @classmethod + def _is_available(cls, attn: "TrtllmAttention") -> bool: + if not attn.is_mla_enable: + return False + if getattr(getattr(attn, "sparse_params", None), "algorithm", None) != "deepseek_v4": + return False + reason = cls._unavailable_reason(attn) + if reason is not None: + # A DeepSeek-V4 layer that does not get this library falls back to + # TRTLLM-Gen silently otherwise; say so once per reason at info. + logger.info_once( + f"PrimTS sparse MLA FMHA is unavailable, DeepSeek-V4 generation stays on the " + f"TRTLLM-Gen fallback: {reason}", + key=f"prims_ts_sparse_mla_unavailable:{reason}", + ) + return False + return True + + @classmethod + def _unavailable_reason(cls, attn: "TrtllmAttention") -> Optional[str]: + sm = get_sm_version() + if sm == 103: + # The kernels' architecture guard admits SM103, but they are only + # signed off on SM100 (prims_ts/kernels/mla_decode/README.md); keep + # B300/GB300 on the TRTLLM-Gen fallback until they are qualified. + return "SM103 is not yet qualified for the PrimTS sparse MLA kernels." + if sm != 100: + return f"requires SM100, got SM{sm}." + if getattr(attn, "use_fp8_ds_mla", False) or getattr(attn, "kv_cache_dtype", "auto") == ( + "fp8_ds_mla" + ): + return "the footer-scale fp8_ds_mla cache layout is not a native E4M3 pool." + if getattr(attn, "_uses_nvfp4_compress", False): + return "NVFP4 compressed pools are not supported." + if ( + int(attn.kv_lora_rank) + int(attn.qk_rope_head_dim) != _LATENT_DIM + or int(attn.head_dim) != _LATENT_DIM + ): + return ( + "requires a 512-wide latent (kv_lora_rank + qk_rope_head_dim), got " + f"{attn.kv_lora_rank} + {attn.qk_rope_head_dim} with head_dim {attn.head_dim}." + ) + if int(attn.qk_rope_head_dim) != _ROPE_DIM: + return f"requires a {_ROPE_DIM}-wide RoPE segment, got {attn.qk_rope_head_dim}." + if not 1 <= int(attn.num_heads) <= cls.MAX_HEADS: + return f"supports 1..{cls.MAX_HEADS} heads, got {attn.num_heads}." + return get_prims_ts_toolchain_unsupported_reason() + + # ------------------------------------------------------------------ support + def _is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase=None, + ) -> bool: + reason = self._unsupported_reason(q, k, v, metadata, forward_args) + if reason is not None: + logger.debug(f"PrimTS sparse MLA FMHA does not support request: {reason}") + return False + return True + + def _unsupported_reason( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + fwd: AttentionForwardArgs, + ) -> Optional[str]: + attn = self.attn + if fwd.attention_input_type != AttentionInputType.generation_only: + return "only generation-only requests are served; context stays on the fallback." + if k is not None or v is not None: + return "MLA generation takes a fused query only." + if q.device.type != "cuda" or q.ndim != 2: + return "the query must be a rank-2 CUDA tensor." + if metadata.kv_cache_manager is None or not hasattr( + metadata.kv_cache_manager, "swa_pool_ptr" + ): + return "a DeepSeek-V4 KV cache manager is required." + if metadata.beam_width != 1: + return "beam search is not supported." + if getattr(metadata, "token_positions_cuda", None) is None: + return "token positions are required for the output RoPE." + if getattr(metadata, "sparse_mla_swa_lens", None) is None: + return "per-step live lengths are required (DeepSeek-V4 metadata)." + if fwd.attention_mask_data is not None or fwd.relative_attention_bias is not None: + return "custom masks and relative bias are not supported." + if metadata.helix_position_offsets is not None: + return "Helix parallelism is not supported." + rows = int(q.shape[0]) + if rows == 0: + return "the request contains no generation tokens." + if int(metadata.num_tokens) - int(metadata.num_ctx_tokens) != rows: + return "the query rows do not match the generation token count." + heads = int(attn.num_heads) + if q.shape[1] != heads * _LATENT_DIM: + return f"query width must be {heads * _LATENT_DIM}, got {q.shape[1]}." + + indices = fwd.sparse_runtime_params.sparse_attn_indices + if not isinstance(indices, torch.Tensor): + return "dual-pool sparse indices are required." + if indices.ndim != 2 or indices.dtype != torch.int32 or indices.shape[0] != rows: + return "sparse indices must be int32 [rows, window + compressed]." + if indices.stride(-1) != 1: + return "sparse index rows must be contiguous." + window = int(attn.sparse_attention_config.window_size) + extra = self._extra_capacity(metadata) + if indices.shape[1] != window + extra: + return ( + f"sparse index width {indices.shape[1]} does not match the dual-pool " + f"layout {window} + {extra}." + ) + + is_fp8 = self._is_fp8_cache(metadata) + if is_fp8: + quant_q = fwd.quant_q_buffer + if quant_q is None or quant_q.dtype not in (torch.uint8, torch.float8_e4m3fn): + return "FP8 sparse MLA requires the quantized query from MLA preprocessing." + if not quant_q.is_contiguous() or quant_q.numel() < rows * heads * _LATENT_DIM: + return "the quantized query buffer is too small or not contiguous." + if attn.kv_scale_quant_orig is None: + return "FP8 sparse MLA requires the KV-cache dequantization scale." + elif q.dtype != torch.bfloat16: + return f"BF16 KV caches require a BF16 query, got {q.dtype}." + + output = fwd.output + if output is None: + return "a preallocated output is required." + if fwd.enable_dsv4_epilogue_fusion: + groups = output.shape[0] if output.ndim == 3 else 0 + if ( + output.ndim != 3 + or output.dtype != torch.float8_e4m3fn + or groups <= 0 + or heads % groups + or output.shape[1] != rows + or output.shape[2] != heads // groups * _LATENT_DIM + or not output.is_contiguous() + ): + return "the fused epilogue output must be a contiguous E4M3 [G, T, (H/G)*512]." + sf = fwd.output_sf + if ( + sf is None + or sf.dtype != torch.float32 + or sf.ndim != 3 + or sf.shape[0] != groups + or sf.shape[1] != heads // groups * (_LATENT_DIM // _FUSED_OUTPUT_BLOCK) + or sf.shape[2] != (rows + 3) // 4 * 4 + or not sf.is_contiguous() + ): + return "the fused epilogue scales must be FP32 [G, (H/G)*4, pad4(T)]." + cos_sin = fwd.dsv4_inv_rope_cos_sin_cache + if cos_sin is None or cos_sin.dtype != torch.float32 or cos_sin.numel() % _ROPE_DIM: + return "the fused epilogue needs the FP32 inverse-RoPE cos/sin table." + else: + if output.dtype != torch.bfloat16 or output.numel() != rows * heads * _LATENT_DIM: + return "the output must be BF16 with one 512-wide row per query head." + if fwd.out_scale is not None or fwd.output_sf is not None: + return "unfused quantized attention output is not supported." + return None + + # ------------------------------------------------------------------ helpers + @staticmethod + def _is_fp8_cache(metadata: "TrtllmAttentionMetadata") -> bool: + from tensorrt_llm.bindings import DataType + + return metadata.kv_cache_manager.dtype == DataType.FP8 + + def _extra_capacity(self, metadata: "DeepseekV4TrtllmAttentionMetadata") -> int: + ratio = int(self.attn.compress_ratio) + if ratio <= 1: + return 0 + return int(metadata.max_compressed_indices[ratio]) + + def _pool_view( + self, + manager: "DeepseekV4CacheManager", + attn_type, + compress_ratio: int, + ) -> torch.Tensor: + """Return a ``[rows, 1, 512]`` view of a whole pool rooted at its base pointer. + + ``sparse_attn_predict`` emits storage rows relative to the pool base, so + the view must start there and span every layer sharing the pool. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType, + ) + + owner = self._pool_views_owner() if self._pool_views_owner is not None else None + if owner is not manager: + self._pool_views.clear() + self._pool_views_owner = weakref.ref(manager) + key = (int(attn_type.value), compress_ratio) + view = self._pool_views.get(key) + if view is not None: + return view + if attn_type == DeepseekV4AttentionType.SWA: + base_ptr = manager.swa_pool_ptr + layers = list(manager.pp_layers) + else: + base_ptr = manager.compress_pool_ptrs[compress_ratio] + layers = [ + layer + for layer in manager.pp_layers + if manager._compress_ratios[layer] == compress_ratio + ] + total_rows = 0 + for layer in layers: + buffer = manager.get_buffers(layer, attn_type) + if buffer.shape[-1] != _LATENT_DIM: + raise RuntimeError( + f"PrimTS sparse MLA expects native 512-wide pool rows, got {tuple(buffer.shape)}." + ) + offset = buffer.data_ptr() - base_ptr + row_bytes = _LATENT_DIM * buffer.element_size() + if offset < 0 or offset % row_bytes: + raise RuntimeError( + f"{attn_type.name} buffer for layer {layer} is not row-aligned to its pool base." + ) + total_rows = max(total_rows, offset // row_bytes + buffer.shape[0] * buffer.shape[1]) + view = convert_to_torch_tensor( + TensorWrapper(base_ptr, manager.dtype, (total_rows, 1, _LATENT_DIM)) + ) + self._pool_views[key] = view + return view + + def _sinks(self, fwd: AttentionForwardArgs) -> Optional[torch.Tensor]: + sinks = fwd.attention_sinks + if sinks is None: + return None + if ( + sinks.dtype == torch.float32 + and sinks.dim() == 1 + and sinks.is_contiguous() + and not sinks.requires_grad + ): + # DeepseekV4TrtllmAttention.attn_sink is already a per-head FP32 + # parameter; pass it through so its address stays stable per layer. + return sinks + key = (sinks.data_ptr(), sinks.numel(), sinks.dtype) + prepared = self._sinks_cache.get(key) + if prepared is not None: + return prepared + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "PrimTS sparse MLA attention sinks must be prepared before CUDA graph capture." + ) + prepared = sinks.detach().to(torch.float32).reshape(-1).contiguous() + if len(self._sinks_cache) >= _SINKS_CACHE_LIMIT: + self._sinks_cache.clear() + self._sinks_cache[key] = prepared + return prepared + + @staticmethod + def _softmax_scale(attn: "TrtllmAttention") -> float: + """TRT-LLM's MLA softmax scale, ``1 / (q_scaling * sqrt(qk_nope + qk_rope))``. + + The scale follows the un-absorbed QK head width (attentionOp.cpp and the + DeepSeek-V4 fallback), not the 512-wide absorbed latent the kernel + contracts over. DeepSeek-V4 checkpoints have 448 + 64 = 512, so the two + only differ for configs such as the 128 + 64 unit-test scenario. + """ + qk_head_dim = int(attn.qk_nope_head_dim) + int(attn.qk_rope_head_dim) + return 1.0 / (float(attn.q_scaling) * math.sqrt(float(qk_head_dim))) + + def _kv_dequant(self) -> float: + if self._kv_dequant_scale is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "PrimTS sparse MLA FP8 scales must be cached before CUDA graph capture." + ) + scale = float(self.attn.kv_scale_quant_orig.view(-1)[0].item()) + if not (math.isfinite(scale) and scale > 0): + raise RuntimeError("PrimTS sparse MLA FP8 scales must be finite and positive.") + self._kv_dequant_scale = scale + return self._kv_dequant_scale + + def _get_plan( + self, + *, + metadata: "TrtllmAttentionMetadata", + device: torch.device, + rows: int, + extra: int, + q_dtype: torch.dtype, + has_sinks: bool, + fused_groups: Optional[int], + ) -> _SparsePlan: + """Return the metadata-cached plan for this static profile, planning it once. + + The key holds every value that changes kernel selection or buffer + geometry: device, query rows, extra capacity, heads, window, dtype, + sinks, fusion groups. One plan serves every layer with that profile. + """ + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + attn = self.attn + heads = int(attn.num_heads) + window = int(attn.sparse_attention_config.window_size) + key = (device, rows, extra, heads, window, q_dtype, has_sinks, fused_groups) + cache = metadata.fmha_plan_caches.setdefault(self.PLAN_CACHE_KEY, {}) + plan = cache.get(key) + if plan is not None: + return plan + t_plan = time.perf_counter() + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "PrimTS sparse MLA plans must be created during eager warmup, not under " + f"CUDA graph capture (rows={rows}, extra={extra}, dtype={q_dtype})." + ) + plan_kwargs, _ = self._plan_kwargs(extra, q_dtype, has_sinks, fused_groups) + # One workspace per metadata object, shared by every plan and grown to + # the largest profile: plans run serially, and a plan's scratch is only + # live during its own launch. The size comes from the same resolver the + # wrapper's plan() uses (its fifth result is the byte extent). + required_bytes = int( + mla_decode._resolve_sparse_mla_plan(device, rows, heads, **plan_kwargs)[4] + ) + shared = metadata.fmha_plan_caches.setdefault(self.WORKSPACE_CACHE_KEY, {}) + workspace = shared.get("buffer") + if workspace is None or workspace.numel() < required_bytes: + workspace = torch.empty((required_bytes,), device=device, dtype=torch.uint8) + shared["buffer"] = workspace + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper(workspace_buffer=workspace) + wrapper.plan(device, rows, heads, **plan_kwargs) + plan = _SparsePlan(wrapper=wrapper) + cache[key] = plan + logger.debug( + f"PrimTS sparse MLA plan rows={rows} extra={extra} sinks={has_sinks} " + f"fused={fused_groups}: {(time.perf_counter() - t_plan) * 1e3:.1f} ms " + f"({len(cache)} plans on this metadata)" + ) + return plan + + def _holes_possible(self) -> bool: + """Whether -1 can appear inside a live index prefix with this layer's config + (KV-cache offload masks non-resident pages in place; the heuristic GVR top-k + pads from its own candidate count); read directly so a renamed field cannot + silently fail open.""" + sparse_cfg = self.attn.sparse_attention_config + return bool(sparse_cfg.enable_kv_cache_offload) or bool(sparse_cfg.enable_heuristic_topk) + + def _plan_kwargs( + self, extra: int, q_dtype: torch.dtype, has_sinks: bool, fused_groups: Optional[int] + ) -> tuple[dict, bool]: + """Return the wrapper plan() kwargs for a static profile and whether -1 holes + can appear inside a live index prefix. + + DeepSeek-V4's producers emit valid prefixes: the window slots are the + first min(pos + 1, window) rows, the compressed slots the first + `topk_lens - window` rows (ratio 128 by construction, ratio 4 from the + indexer's top-k), with -1 only in the tail. Promising that lets the + kernel derive its masks from the live lengths instead of loading and + checking every index. KV-cache offload breaks the promise (non-resident + pages are masked to -1 in place), and the opt-in heuristic (GVR) top-k + engines pad from their own candidate count, so both keep the generic + hole-checking path. The fields are read directly so a renamed config + cannot silently fail open. + """ + holes_possible = self._holes_possible() + plan_kwargs = dict( + max_topk=int(self.attn.sparse_attention_config.window_size), + max_extra_topk=extra, + max_seq_len_q=1, + q_data_type=q_dtype, + kv_layout="NHD", + has_sinks=has_sinks, + return_lse=False, + assume_valid_prefix=not holes_possible, + ) + if fused_groups: + plan_kwargs.update( + output_quant_block_size=_FUSED_OUTPUT_BLOCK, + out_scale_format="fp32", + num_output_groups=fused_groups, + ) + return plan_kwargs, holes_possible + + def _warm_kernel_variants( + self, + *, + metadata: "TrtllmAttentionMetadata", + device: torch.device, + extra: int, + q_dtype: torch.dtype, + has_sinks: bool, + fused_groups: Optional[int], + ) -> None: + """Load every kernel variant of a static profile during eager warmup. + + A batch's context ramp brings a new generation row count almost every + step, and mixed (context + generation) steps run eager. Planning a new + row count is cheap (~1 ms), but the first plan of each kernel variant + (the split-KV schedule changes with the row count and the extra + capacity) loads it from the persistent compile cache (~0.25 s), which + stalls the host while the GPU runs prefill. Compiled kernels are cached + per process by the wrapper, so this runs once per process and profile, + on the eager metadata only (CUDA graph metadata objects replay one row + count each and never need it), and never under capture. It scans the + row counts a step can carry, keeps one representative per distinct + compile signature, and plans them largest first so the shared + workspace is allocated once at its final size. + """ + if torch.cuda.is_current_stream_capturing() or getattr(metadata, "is_cuda_graph", False): + return + attn = self.attn + heads = int(attn.num_heads) + profile_key = ( + str(device), + heads, + extra, + q_dtype, + has_sinks, + fused_groups, + int(attn.sparse_attention_config.window_size), + ) + if profile_key in _WARMED_PROFILES: + return + _WARMED_PROFILES.add(profile_key) + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + from . import prims_ts_precompile + + plan_kwargs, _ = self._plan_kwargs(extra, q_dtype, has_sinks, fused_groups) + # Generation rows per step: one per request, plus its draft tokens with + # speculative decoding, within the step's token budget. The fields are + # declared on TrtllmAttentionMetadata; read them directly so a rename + # fails here instead of silently shrinking the warmed set and moving + # the MTP row counts' compiles to the first serving step. + requests = int(metadata.max_num_requests or 0) + draft = int(metadata.max_total_draft_tokens or 0) + tokens = int(metadata.max_num_tokens or 0) + max_rows = requests * (1 + draft) + if tokens: + max_rows = min(max_rows, tokens) if max_rows else tokens + candidates = [r for r in range(1, min(max_rows, self.WARM_ROW_SCAN) + 1)] + candidates += [r for r in self.WARM_ROW_TAIL if self.WARM_ROW_SCAN < r <= max_rows] + t_warm = time.perf_counter() + representative: dict = {} + for rows in candidates: + spec = mla_decode._resolve_sparse_mla_plan(device, rows, heads, **plan_kwargs)[2] + signature = mla_decode._mla_kernel_compile_signature(spec.kernel) + representative[signature] = rows # the largest row count of each variant + self._precompile_variants( + device, + heads, + plan_kwargs, + sorted(representative.values(), reverse=True), + local_ranks=prims_ts_precompile.local_rank_count(metadata.mapping), + ) + for rows in sorted(representative.values(), reverse=True): + self._get_plan( + metadata=metadata, + device=device, + rows=rows, + extra=extra, + q_dtype=q_dtype, + has_sinks=has_sinks, + fused_groups=fused_groups, + ) + logger.debug( + f"PrimTS sparse MLA warmed {len(representative)} kernel variants for extra={extra} " + f"(rows scanned up to {max_rows}): {(time.perf_counter() - t_warm) * 1e3:.0f} ms" + ) + + @staticmethod + def _precompile_variants( + device: torch.device, + heads: int, + plan_kwargs: dict, + rows_list: list, + *, + local_ranks: int = 1, + ) -> None: + """Populate the on-disk kernel cache for ``rows_list`` with worker subprocesses. + + Skipped when every variant is already cached (a probe plans each row + count without compiling), so a warm restart spawns nothing. Otherwise + best-effort: on any failure the serial plan loop that follows compiles + in-process, exactly as without workers. See ``prims_ts_precompile``. + """ + from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache, mla_decode + + from . import prims_ts_precompile + + workers = prims_ts_precompile.worker_count(len(rows_list), local_ranks=local_ranks) + if workers <= 1: + return + try: + with _compile_cache.probe_cached(): + for rows in rows_list: + mla_decode.BatchSparseMLADecodePagedTSWrapper().plan( + device, rows, heads, **plan_kwargs + ) + except _compile_cache.KernelNotCached: + pass + except Exception as error: # noqa: BLE001 -- a probe failure only costs the fan-out + if prims_ts_precompile.strict(): + raise + logger.warning( + f"PrimTS sparse MLA cache probe failed: {error!r}; compiling in-process." + ) + return + else: + logger.debug("PrimTS sparse MLA kernel variants are all cached; no precompile workers") + return + device_index = device.index if device.index is not None else torch.cuda.current_device() + payloads = [ + prims_ts_precompile.encode_payload(device_index, rows, heads, plan_kwargs) + for rows in rows_list + ] + t_pre = time.perf_counter() + try: + ok = prims_ts_precompile.run_precompile_workers(payloads, workers) + except Exception as error: # noqa: BLE001 -- warmup must not fail on a helper + if prims_ts_precompile.strict(): + raise + logger.warning(f"PrimTS sparse MLA parallel precompile skipped: {error!r}") + return + if not ok and prims_ts_precompile.strict(): + raise RuntimeError( + "PrimTS sparse MLA precompile workers failed (TLLM_PRIMTS_WARM_STRICT=1)" + ) + logger.info( + f"PrimTS sparse MLA precompiled {len(payloads)} kernel variants with {workers} workers " + f"in {time.perf_counter() - t_pre:.0f} s{'' if ok else ' (some workers failed)'}" + ) + + @staticmethod + def _claim_contract_validation(contract_key: tuple) -> bool: + """Return True exactly once per process for an index contract (never under capture). + + The contract (valid prefixes, indices inside the pools) is a property + of the index producers and the pools they address, so one validated + eager launch per process and contract is enough; validating every + plan's first launch re-synchronized the host with a GPU that was busy + with prefill on the context-ramp steps that created new plans. + """ + if torch.cuda.is_current_stream_capturing() or contract_key in _VALIDATED_CONTRACTS: + return False + _VALIDATED_CONTRACTS.add(contract_key) + logger.debug(f"PrimTS sparse MLA validating the index contract {contract_key}") + return True + + # ------------------------------------------------------------------ forward + def _check_request( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + fwd: AttentionForwardArgs, + ) -> None: + """Run the full request check once per request geometry. + + The manager selected this library through is_supported(); the re-check + here guards a changed request, so it runs only when the per-step + geometry (metadata object, query, indices, output, fusion mode, token + counts) differs from the last accepted call. + """ + indices = fwd.sparse_runtime_params.sparse_attn_indices + output = fwd.output + check_key = ( + id(metadata), + int(metadata.num_tokens), + int(metadata.num_ctx_tokens), + q.shape, + q.dtype, + k is None, + v is None, + None + if not isinstance(indices, torch.Tensor) + else (indices.shape, indices.dtype, indices.stride(-1)), + None if output is None else (output.shape, output.dtype, output.is_contiguous()), + fwd.enable_dsv4_epilogue_fusion, + fwd.quant_q_buffer is None, + fwd.attention_sinks is None, + fwd.output_sf is None, + fwd.out_scale is None, + fwd.dsv4_inv_rope_cos_sin_cache is None, + ) + if check_key == self._checked_key: + return + reason = self._unsupported_reason(q, k, v, metadata, fwd) + if reason is not None: + raise RuntimeError(f"PrimTS sparse MLA FMHA cannot run this request: {reason}") + self._checked_key = check_key + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType, + ) + + self._check_request(q, k, v, metadata, forward_args) + indices = forward_args.sparse_runtime_params.sparse_attn_indices + output = forward_args.output + attn = self.attn + manager = metadata.kv_cache_manager + heads = int(attn.num_heads) + rows = int(q.shape[0]) + start = int(metadata.num_ctx_tokens) + end = start + rows + ratio = int(attn.compress_ratio) + window = int(attn.sparse_attention_config.window_size) + extra = self._extra_capacity(metadata) + is_fp8 = self._is_fp8_cache(metadata) + fused = bool(forward_args.enable_dsv4_epilogue_fusion) + groups = int(output.shape[0]) if fused else None + + if is_fp8: + # Mirror the fallback's scheduler prologue, which sets both + # dequant_scale_q and dequant_scale_kv to the layer's + # kv_scale_quant_orig (DeepseekV4TrtllmAttention.sparse_attn_predict). + dequant = self._kv_dequant() + q_scale = kv_scale = dequant + query = ( + forward_args.quant_q_buffer.view(-1)[: rows * heads * _LATENT_DIM] + .view(torch.float8_e4m3fn) + .view(rows, 1, heads, _LATENT_DIM) + ) + else: + q_scale = kv_scale = 1.0 + query = q.view(rows, 1, heads, _LATENT_DIM) + # The kernel takes the combined scales: the softmax scale carries the Q + # and KV descales, the output scale carries the KV descale (both pools + # share it). They are per-layer constants, so the wrapper's static-scale + # kernel serves them without a device scalar. + softmax_scale = self._softmax_scale(attn) * q_scale * kv_scale + output_scale = kv_scale + + sinks = self._sinks(forward_args) + profile = dict( + metadata=metadata, + device=q.device, + extra=extra, + q_dtype=query.dtype, + has_sinks=sinks is not None, + fused_groups=groups, + ) + self._warm_kernel_variants(**profile) + plan = self._get_plan(rows=rows, **profile) + + # Storage-row indices straight from sparse_attn_predict: the window + # slots address the SWA pool and the compressed slots address the + # compress pool, both relative to their pool base. The live lengths + # come from the metadata too, derived once per step: the token's + # position within the window, and the compressed rows selected beyond + # the window. + positions = metadata.token_positions_cuda[start:end] + lengths = metadata.sparse_mla_swa_lens[start:end] + swa_pool = self._pool_view(manager, DeepseekV4AttentionType.SWA, 1) + extra_pool = None + run_kwargs = {} + if extra: + extra_pool = self._pool_view(manager, DeepseekV4AttentionType.COMPRESS, ratio) + run_kwargs["extra_indices"] = indices[:, window:] + run_kwargs["extra_lengths"] = metadata.sparse_mla_extra_lens[ratio][start:end] + + if fused: + # TRT-LLM's fused epilogue buffers are group-major [G, T, C] with + # FP32 scales [G, (H/G)*4, pad4(T)]; the kernel takes their + # token-major logical views. + run_kwargs["out"] = output.transpose(0, 1) + run_kwargs["out_scale"] = forward_args.output_sf.permute(2, 0, 1)[:rows] + # Positions carry the query prefix shape [rows, 1]. + run_kwargs["token_positions"] = positions.contiguous().view(rows, 1) + cos_sin = forward_args.dsv4_inv_rope_cos_sin_cache + run_kwargs["cos_sin_cache"] = cos_sin.reshape(-1, _ROPE_DIM).contiguous() + else: + run_kwargs["out"] = output.view(rows, 1, heads, _LATENT_DIM) + + # One validated eager launch per process and contract; the key names + # the producers (window, extra capacity, compress ratio, hole policy) + # and the pools they address, so the KV-cache estimation executor's + # temporary pools and the serving pools are validated separately. + validate = self._claim_contract_validation( + ( + str(q.device), + window, + extra, + ratio, + self._holes_possible(), + str(query.dtype), + swa_pool.data_ptr(), + int(swa_pool.shape[0]), + extra_pool.data_ptr() if extra_pool is not None else 0, + int(extra_pool.shape[0]) if extra_pool is not None else 0, + ) + ) + plan.wrapper.run( + query, + swa_pool, + indices[:, :window], + lengths, + extra_pool, + softmax_scale=softmax_scale, + output_scale=output_scale, + sinks=sinks, + validate=validate, + **run_kwargs, + ) + + +__all__ = ["PrimsTSSparseMlaFmha", "get_prims_ts_toolchain_unsupported_reason"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index b7b60a0de6d1..3e00148b9c6a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -38,6 +38,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: from .msa_decode import MsaDecodeFmha from .msa_prefill import MsaPrefillFmha from .prims_ts_block_sparse import PrimsTSBlockSparseFmha + from .prims_ts_sparse_mla import PrimsTSSparseMlaFmha return { "fp4_mla": Fp4MlaFmha, @@ -51,13 +52,21 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: "flashinfer_sparse_mla": FlashInferSparseMlaFmha, "prims_ts": PrimsTSFmha, "prims_ts_block_sparse": PrimsTSBlockSparseFmha, + # DeepSeek-V4 generation on the vendored PrimTS native sparse MLA + # kernel. Enabled by default: it only reports itself available for + # DeepSeek-V4 on SM100, so every other model and architecture never + # sees it; everything else keeps the TRTLLM-Gen libraries below. + "prims_ts_sparse_mla": PrimsTSSparseMlaFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, } FMHA_LIBS: dict[str, FmhaCls] = init_fmha_libs() -DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple(name for name in FMHA_LIBS if name != "prims_ts") +_OPT_IN_FMHA_LIBS: frozenset[str] = frozenset({"prims_ts"}) +DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple( + name for name in FMHA_LIBS if name not in _OPT_IN_FMHA_LIBS +) def _parse_fmha_libs_env() -> tuple[str, ...]: diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py index e40d8ce0d41c..477f52de75cf 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py @@ -37,6 +37,9 @@ batch_prefill_with_paged_kv_cache, ) from .mla_decode import ( + BatchSparseMLADecodePagedTSWrapper, + batch_sparse_mla_decode_with_paged_kv_cache, + get_prims_ts_sparse_mla_decode_workspace_size, BatchMLADecodePagedTSWrapper, batch_mla_decode_with_paged_kv_cache, get_prims_ts_batch_mla_decode_workspace_size, @@ -49,6 +52,9 @@ from .sage import SageAttentionConfig, SageAttentionParams __all__ = [ + "BatchSparseMLADecodePagedTSWrapper", + "batch_sparse_mla_decode_with_paged_kv_cache", + "get_prims_ts_sparse_mla_decode_workspace_size", "BlockSparseTSWrapper", "SageAttentionConfig", "SageAttentionParams", diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py index c66cf74d37d7..b0510cfc090b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py @@ -17,6 +17,8 @@ from collections.abc import Callable import functools +from .._compile_cache import build_and_load_attention_kernel + import torch from flashinfer.utils import ceil_div @@ -535,7 +537,7 @@ def fake_compact( Float32(1.0), ) - with torch.cuda.device(key.device_index): + def compile_kernel(): return cute.compile( tensor_adapter, *dynamic_args, @@ -549,6 +551,11 @@ def fake_compact( options=_COMPILE_OPTIONS, ) + with torch.cuda.device(key.device_index): + return build_and_load_attention_kernel( + "block_sparse", "main", key, _COMPILE_OPTIONS, compile_kernel, extra=() + ) + @functools.cache def _get_compiled_block_sparse( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py new file mode 100644 index 000000000000..cdfa76294f2d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py @@ -0,0 +1,251 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Persistent compilation shared by the PrimTS attention adapters.""" + +from collections.abc import Callable +import contextlib +import contextvars +import dataclasses +import enum +import functools +import hashlib +import os +from pathlib import Path +from typing import Any + + +def _stable_value(value: object) -> object: + if isinstance(value, enum.Enum): + return (type(value).__module__, type(value).__qualname__, value.name) + if value is None or type(value) in (bool, int, float, str): + return value + if isinstance(value, type): + return (value.__module__, value.__qualname__) + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return ( + type(value).__module__, + type(value).__qualname__, + tuple( + (field.name, _stable_value(getattr(value, field.name))) + for field in dataclasses.fields(value) + if field.compare + ), + ) + if isinstance(value, (tuple, list)): + return tuple(_stable_value(item) for item in value) + raise TypeError(f"no stable cache-key encoding for {type(value).__name__}") + + +def _compile_cache_key(spec: object, options: str, extra: tuple = ()) -> str | None: + """Hash compile identity, excluding only device placement and non-key fields. + + MLA's non-comparing kernel object is represented by its kernel_signature. + Unsupported values bypass persistence rather than risking a stale artifact. + """ + if dataclasses.is_dataclass(spec) and not isinstance(spec, type): + spec = tuple( + (field.name, getattr(spec, field.name)) + for field in dataclasses.fields(spec) + if field.compare and field.name != "device_index" + ) + try: + encoded = repr(_stable_value((spec, options, extra))) + except TypeError: + return None + return hashlib.sha256(encoded.encode()).hexdigest() + + +# Warm-up precompile modes. Under ``precompile_only()`` the attention kernels +# are compiled through ``JitSpecCuteDsl.compile_and_persist()`` (the +# ``cute.compile`` runs outside the module-level file lock that +# ``build_and_load()`` holds for a whole build, so sibling processes compile +# different kernels concurrently), committed to disk, and not loaded: the +# serving process loads them afterwards. Under ``probe_cached()`` nothing is +# compiled; a kernel missing from the cache raises ``KernelNotCached``. Both +# modes return a placeholder in place of the kernel. +_PRECOMPILE_MODE: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "prims_ts_precompile_mode", default=None +) +_SKIP_BUSY: contextvars.ContextVar[bool] = contextvars.ContextVar( + "prims_ts_precompile_skip_busy", default=False +) + + +class KernelNotCached(Exception): + """Raised under ``probe_cached()`` for a kernel that is not on disk.""" + + +class KernelBusy(Exception): + """Raised under ``precompile_only(skip_busy=True)`` while another process compiles the kernel.""" + + +@contextlib.contextmanager +def precompile_only(*, skip_busy: bool = False): + """Compile and persist kernels without loading them (see module note). + + With ``skip_busy`` a kernel whose per-kernel lock another process holds + raises ``KernelBusy`` instead of waiting, so a worker can defer it and + compile something else first. + """ + token = _PRECOMPILE_MODE.set("persist") + busy_token = _SKIP_BUSY.set(skip_busy) + try: + yield + finally: + _SKIP_BUSY.reset(busy_token) + _PRECOMPILE_MODE.reset(token) + + +@contextlib.contextmanager +def probe_cached(): + """Report, without compiling, whether every kernel a plan needs is on disk.""" + token = _PRECOMPILE_MODE.set("probe") + try: + yield + finally: + _PRECOMPILE_MODE.reset(token) + + +def expected_cache_meta() -> dict: + """The on-disk module metadata this process would write (arch, DSL version, source hash).""" + from flashinfer.jit.cute_dsl_core import ( + _get_compile_arch, + _get_cute_dsl_version, + _hash_source_files, + ) + + return { + "arch": _get_compile_arch(), + "cute_dsl_version": _get_cute_dsl_version(), + "source_sha256": _hash_source_files(tuple(_kernel_source_files())), + } + + +def _unloaded_kernel(*args, **kwargs): + raise RuntimeError( + "PrimTS kernel was compiled in precompile-only mode and never loaded; " + "load it through a normal plan() outside precompile_only()." + ) + + +def is_placeholder(compiled) -> bool: + """True for the stand-in returned under the precompile modes.""" + return compiled is _unloaded_kernel + + +def _make_spec(module_name: str, kernel_name: str, compile_fn, key_files): + from flashinfer.jit.cute_dsl_core import JitSpecCuteDsl, _hash_source_files + + return JitSpecCuteDsl(module_name, kernel_name, compile_fn, _hash_source_files(tuple(key_files))) + + +def _precompile_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: + """Compile one kernel into the on-disk cache unless it is already there. + + One compile per kernel per workspace: a sibling that wants the same kernel + waits on the per-kernel lock (or raises ``KernelBusy`` under ``skip_busy``) + and then finds it on disk. The lock lives in FlashInfer's tmp directory, + not in the module directory: creating that directory before the first + export would let ``compile_and_persist()`` of a sibling read a missing + meta.json as a stale module and wipe the exports already there. + """ + from filelock import FileLock, Timeout + from flashinfer.jit.core import get_tmpdir + + spec = _make_spec(module_name, kernel_name, compile_fn, key_files) + if spec.is_compiled: + return + lock = FileLock(get_tmpdir() / f"{spec.module_dir_name}.{spec.kernel_name}.lock", thread_local=False) + if _SKIP_BUSY.get(): + try: + lock.acquire(timeout=0) + except Timeout: + raise KernelBusy(spec.name) from None + else: + lock.acquire() + try: + if spec.is_compiled: + return + spec.compile_and_persist() + finally: + lock.release() + + +def _probe_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: + spec = _make_spec(module_name, kernel_name, compile_fn, key_files) + if not spec.is_compiled: + raise KernelNotCached(spec.name) + + +@functools.cache +def _kernel_source_files() -> tuple[str, ...]: + # Family adapters share masks, reductions, sparse/Sage preparation and + # compatibility helpers. Invalidate them together when any source changes. + root = Path(__file__).resolve().parent + return tuple(sorted(str(path) for path in root.rglob("*.py"))) + + +def build_and_load_attention_kernel( + family: str, + kind: str, + spec: object, + options: str, + compile_fn: Callable[[], Any], + *, + extra: tuple = (), +) -> Any: + """Build or load under the caller's selected CUDA device context.""" + from flashinfer.jit.cute_dsl_core import ( + build_and_load_cute_dsl_kernel, + cute_dsl_cache_disabled, + ) + + mode = _PRECOMPILE_MODE.get() + if mode is not None and os.environ.get("FLASHINFER_DISABLE_JIT"): + raise RuntimeError("PrimTS warm-up precompile cannot compile with FLASHINFER_DISABLE_JIT set.") + if cute_dsl_cache_disabled(): + if mode is not None: + raise RuntimeError( + "PrimTS warm-up precompile needs the on-disk CuTe DSL cache " + "(FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 is set)." + ) + return compile_fn() + + key = _compile_cache_key(spec, options, extra) + if key is None: + if mode is not None: + raise RuntimeError("PrimTS warm-up precompile cannot key this kernel for the on-disk cache.") + return compile_fn() + import torch + + # A family target such as sm100f can cover physical GPUs with different + # host-selected tuning defaults. The module metadata keys the target; + # the specialization must also distinguish the physical capability. + major, minor = torch.cuda.get_device_capability() + if mode == "persist": + _precompile_kernel( + f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files() + ) + return _unloaded_kernel + if mode == "probe": + _probe_kernel(f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files()) + return _unloaded_kernel + return build_and_load_cute_dsl_kernel( + f"prims_ts_{family}", + f"{kind}_sm{major}{minor}_{key}", + compile_fn, + extra_key_files=_kernel_source_files(), + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_hardware.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_hardware.py new file mode 100644 index 000000000000..f554a507e816 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_hardware.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Device-local occupancy probes shared by PrimTS attention planners.""" + +import functools + + +@functools.cache +def _query_max_active_clusters(device_index: int, cluster_size: int) -> int: + import torch + import cutlass.utils as utils + from cuda.bindings import driver as cuda_drv + + with torch.cuda.device(device_index): + # HardwareInfo must first see a live primary context. Do this only on + # cache misses; new plan shapes reuse the device's occupancy result. + torch.empty(0, device="cuda") + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return int( + utils.HardwareInfo(device_index).get_max_active_clusters( + cluster_size, stream + ) + ) + + +def get_max_active_clusters(cluster_size: int) -> int: + import torch + + return _query_max_active_clusters(torch.cuda.current_device(), cluster_size) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py index a0ea3425aab7..e1a309e07a21 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py @@ -121,7 +121,17 @@ def fanout(instances: int) -> int: tile_size_q=candidate.tile_size_q, num_insts_kv=instances, splits_kv=splits, - head_dim_per_stage_kv=128 if head_dim == 256 else 0, + # Full FP8 stages preserve the 128-byte swizzled planes while issuing + # PV at N256. BF16 and Swaps retain their existing D128 head bands. + head_dim_per_stage_kv=( + 256 + if head_dim == 256 + and q_dtype_key == "float8_e4m3fn" + and candidate.variant == "keeps_mma_ab" + else 128 + if head_dim == 256 + else 0 + ), o_stages=1 if head_dim == 256 and candidate.variant == "keeps_mma_ab" else 2, use_persistent_scheduler=( splits == 1 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index d297bc79c1e0..1ec1080b942f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -43,6 +43,9 @@ import struct from typing import TYPE_CHECKING, Callable, Literal, Optional +from ._compile_cache import build_and_load_attention_kernel +from ._hardware import get_max_active_clusters + import torch from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api @@ -1461,7 +1464,6 @@ def _make_context_scheduler_probe( """Build the common static probe and logical work count for policy.""" import cutlass - import cutlass.utils as utils dtype_map = { torch.float16: cutlass.Float16, @@ -1488,7 +1490,7 @@ def _make_context_scheduler_probe( max_kv_len=max_kv_len, ) with torch.cuda.device(geometry.device_index): - max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) + max_active_clusters = get_max_active_clusters(1) num_seq_tiles = ( geometry.max_seq_len_q + probe.cfg.cta_tiler[0] - 1 ) // probe.cfg.cta_tiler[0] @@ -1757,7 +1759,6 @@ def _get_compiled_context( import cutlass import cutlass.cute as cute from cuda.bindings import driver as cuda_drv - import cutlass.utils as utils dtype_map = { "float16": cutlass.Float16, @@ -1769,7 +1770,7 @@ def _get_compiled_context( output_dtype = dtype_map[output_dtype_key] has_variable_window = mask_type == "variable_window" with torch.cuda.device(device_index): - max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) + max_active_clusters = get_max_active_clusters(1) fmha = _make_context_kernel( device_index=device_index, input_qk_dtype=input_qk_dtype, @@ -1916,8 +1917,8 @@ def fake_compact(dtype, shape, assumed_align): # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. - with torch.cuda.device(device_index): - compiled = cute.compile[cute.FrontendNext]( + def compile_kernel(): + return cute.compile[cute.FrontendNext]( tensor_adapter, q_fake, k_fake, @@ -1937,6 +1938,16 @@ def fake_compact(dtype, shape, assumed_align): max_seq_len_k, options=_COMPILE_OPTIONS, ) + + with torch.cuda.device(device_index): + compiled = build_and_load_attention_kernel( + "context", + "main", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(max_active_clusters,), + ) policy = ( ("scheduler", scheduler), ("pairing", "head" if head_paired else "query"), @@ -1975,7 +1986,6 @@ def _get_compiled_paged_context( import cutlass import cutlass.cute as cute from cuda.bindings import driver as cuda_drv - import cutlass.utils as utils dtype_map = { "float16": cutlass.Float16, @@ -1986,7 +1996,7 @@ def _get_compiled_paged_context( input_pv_dtype = dtype_map[pv_dtype_key] output_dtype = dtype_map[output_dtype_key] with torch.cuda.device(device_index): - max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) + max_active_clusters = get_max_active_clusters(1) fmha = _make_context_kernel( device_index=device_index, input_qk_dtype=input_qk_dtype, @@ -2090,8 +2100,8 @@ def fake_compact(dtype, shape, assumed_align): seq_lens_fake = fake_compact(cutlass.Int32, (batch_size,), 4) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - with torch.cuda.device(device_index): - compiled = cute.compile[cute.FrontendNext]( + def compile_kernel(): + return cute.compile[cute.FrontendNext]( tensor_adapter, q_fake, k_fake, @@ -2108,6 +2118,16 @@ def fake_compact(dtype, shape, assumed_align): max_kv_len, options=_COMPILE_OPTIONS, ) + + with torch.cuda.device(device_index): + compiled = build_and_load_attention_kernel( + "context", + "paged", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(max_active_clusters,), + ) policy = ( ("scheduler", scheduler), ("pairing", "head" if head_paired else "query"), diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 205e875fc390..f91814650c48 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -16,7 +16,7 @@ """Task-scheduled paged decode with a FlashInfer-style plan/run lifecycle.""" from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace import functools import math import numbers @@ -25,6 +25,8 @@ import torch +from ._compile_cache import build_and_load_attention_kernel + from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api from . import _q_token_kv_block_sparse_policy as _sparse_policy @@ -694,8 +696,6 @@ def _validate_dtype_pair( k_dtype: torch.dtype, v_dtype: torch.dtype, output_dtype: torch.dtype, - *, - allow_fp8_bf16_output: bool = False, ) -> None: _dtype_key(q_dtype) _dtype_key(k_dtype) @@ -723,25 +723,16 @@ def _validate_dtype_pair( q_dtype == torch.float8_e4m3fn and matching_kv and k_dtype in (torch.float8_e4m3fn, torch.uint8) - and output_dtype - in ( - (torch.float16, torch.bfloat16, torch.float8_e4m3fn) - if allow_fp8_bf16_output - else (torch.float16, torch.float8_e4m3fn) - ) + and output_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) ) ) if not supported: raise NotImplementedError( "attention-ts decode supports FP16->FP16, BF16->BF16, " - + ( - "FP8-E4M3->FP16/BF16, and FP8-E4M3->FP8-E4M3; " - if allow_fp8_bf16_output - else "FP8-E4M3->FP16, and FP8-E4M3->FP8-E4M3; " - ) - + "BF16 Q + FP8 K/V, BF16/FP8-E4M3 Q + NVFP4 K/V, " - + "and BF16 Q/K + FP8 V; got " - + f"Q_{q_dtype}_K_{k_dtype}_V_{v_dtype}->{output_dtype}" + "FP8-E4M3->FP16/BF16/FP8-E4M3; " + "BF16 Q + FP8 K/V, BF16/FP8-E4M3 Q + NVFP4 K/V, " + "and BF16 Q/K + FP8 V; got " + f"Q_{q_dtype}_K_{k_dtype}_V_{v_dtype}->{output_dtype}" ) @@ -1357,7 +1348,7 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: head_dim=head_dim, page_size=page_size, max_kv_len=max_kv_len, - seq_len_q=seq_len_q, + group_size=seq_len_q, q_dtype_key=q_dtype_key, k_dtype_key=k_dtype_key, v_dtype_key=v_dtype_key, @@ -1487,12 +1478,16 @@ def _make_decode_compile_spec( kv_prefix_mode: Literal["dynamic", "planned_full"], kv_lengths_mode: Literal["dynamic", "planned_uniform_max"], direct_q1_spec: Optional[_DirectQ1CompileSpec] = None, + flat_native_kv_tma: bool = False, ) -> _DecodeCompileSpec: """Freeze the resolved topology while leaving batch in runtime tensors.""" + config_items = dict(launch_spec.config.compile_signature()) + config_items["use_flat_native_kv_tma"] = flat_native_kv_tma + config_items["flat_native_kv_num_heads"] = num_kv_heads if flat_native_kv_tma else 1 return _DecodeCompileSpec( device_index=device_index, - config_items=launch_spec.config.compile_signature(), + config_items=tuple(config_items.items()), num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, @@ -1905,8 +1900,8 @@ def fake_compact(dtype, shape, assumed_align): ), ) - with torch.cuda.device(device_index): - compiled_main = cute.compile( + def compile_main(): + return cute.compile( main_tensor_adapter, q_fake, k_fake, @@ -1944,28 +1939,38 @@ def fake_compact(dtype, shape, assumed_align): static_native_uniform_kv, options=_COMPILE_OPTIONS, ) + + def compile_reducer(): + assert reduction_tensor_adapter is not None + return cute.compile( + reduction_tensor_adapter, + out_fake, + seq_lens_fake, + cu_seqlens_q_fake, + partial_o_fake, + partial_stats_fake, + attention_sinks_fake, + Float32(1.0), + Float32(1.0), + direct_q1_fake, + stream_fake, + cfg, + num_qo_heads, + num_kv_heads, + head_dim, + max_kv_len, + static_full_split_prefix, + options=_COMPILE_OPTIONS, + ) + + with torch.cuda.device(device_index): + compiled_main = build_and_load_attention_kernel( + "decode", "main", compile_spec, _COMPILE_OPTIONS, compile_main + ) compiled_reducer = None if cfg.use_separate_reduction_kernel: - assert reduction_tensor_adapter is not None - compiled_reducer = cute.compile( - reduction_tensor_adapter, - out_fake, - seq_lens_fake, - cu_seqlens_q_fake, - partial_o_fake, - partial_stats_fake, - attention_sinks_fake, - Float32(1.0), - Float32(1.0), - direct_q1_fake, - stream_fake, - cfg, - num_qo_heads, - num_kv_heads, - head_dim, - max_kv_len, - static_full_split_prefix, - options=_COMPILE_OPTIONS, + compiled_reducer = build_and_load_attention_kernel( + "decode", "reducer", compile_spec, _COMPILE_OPTIONS, compile_reducer ) return compiled_main, compiled_reducer @@ -2486,6 +2491,7 @@ class PrimsTSBatchDecodePlan: _compiled_main: Callable[..., object] _compiled_reducer: Optional[Callable[..., object]] _direct_q1_inputs: tuple[torch.Tensor, ...] = () + _flat_locator_heads: int = 1 def run( self, @@ -2854,7 +2860,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( head_dim: int, page_size: int, max_kv_len: int, - seq_len_q: int, + group_size: int, q_dtype_key: str, k_dtype_key: str, v_dtype_key: str, @@ -2875,7 +2881,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( heads_q_per_kv = num_qo_heads // num_kv_heads _validate_prims_ts_q_token_kv_block_sparse_group_capacity( - seq_len_q, + group_size, num_qo_heads, num_kv_heads, ) @@ -2905,7 +2911,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( ) launch = _sparse_policy.select_sparse_launch( - group_size=seq_len_q, + group_size=group_size, heads_q_per_kv=heads_q_per_kv, head_dim=head_dim, q_dtype_key=q_dtype_key, @@ -2935,14 +2941,14 @@ def _resolve_q_token_kv_block_sparse_decode_config( "scheduler_warp_idx": 13, "clc_load_warp_idx": 16, "page_offsets_warp_idx": 14, - "page_offsets_num_warps": 2 if seq_len_q == 1 else 1, + "page_offsets_num_warps": 2 if group_size == 1 else 1, "load_warp_idx": 16, # FP8's predicated TMA helper constructs coordinates in every load # lane, so Q1 needs only one producer warpgroup. BF16 distributes page # fragments across two producer warpgroups. Grouped routes keep # their qualified two-warpgroup schedule for both dtypes. "load_num_warps": ( - 4 if seq_len_q == 1 and q_dtype_key == "float8_e4m3fn" else 8 + 4 if group_size == 1 and q_dtype_key == "float8_e4m3fn" else 8 ), } if launch.use_keeps_mma_ab: @@ -2956,7 +2962,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( ), splits_kv=launch.splits_kv, max_splits_kv=launch.splits_kv, - min_loop_iters_per_split=(1 if seq_len_q == 1 else MIN_LOOP_ITERS_PER_SPLIT), + min_loop_iters_per_split=(1 if group_size == 1 else MIN_LOOP_ITERS_PER_SPLIT), ) @@ -3169,7 +3175,7 @@ def validate_q_token_kv_block_sparse_group_size( num_qo_heads: int, num_kv_heads: int, *, - group_size: int, + group_size: int = 1, ) -> int: """Validate a QToken-KvBlock-Sparse-Attention Q1--Q8 group. @@ -3200,7 +3206,8 @@ def validate_q_token_kv_block_sparse_group_size( Number of key/value heads. group_size : int Caller-selected number of query rows per QToken-KvBlock-Sparse-Attention route. Supported values - are Q1--Q8, subject to the TileQ128 head-capacity bound. + are Q1--Q8, subject to the TileQ128 head-capacity bound. Defaults to + one (no grouping). Returns ------- @@ -3222,16 +3229,17 @@ def make_q_token_kv_block_sparse_qo_indptr( query_start_loc_cpu: torch.Tensor, num_query_tokens: int, *, - group_size: int, + group_size: int = 1, device: Optional[Union[int, str, torch.device]] = None, ) -> torch.Tensor: """Build request-safe QToken-KvBlock-Sparse-Attention route offsets. ``query_start_loc_cpu`` contains cumulative offsets for real requests. Every request is chunked independently, so its final route may contain - fewer than ``group_size`` rows. Any inert CUDA-graph padding suffix is - partitioned separately. The returned Int32 ``qo_indptr`` can be passed to - the packed-Q wrapper together with planned ``seq_len_q=group_size``. + fewer than ``group_size`` rows. The default is one (no query grouping). + Any inert CUDA-graph padding suffix is partitioned separately. Pass the + returned Int32 ``qo_indptr`` to a packed-Q wrapper planned with the same + ``group_size``. Parameters ---------- @@ -3243,7 +3251,7 @@ def make_q_token_kv_block_sparse_qo_indptr( Total flattened query-row count, including any inert CUDA-graph padding suffix. group_size : int - Maximum number of rows in each request-safe QToken-KvBlock-Sparse-Attention route. + Maximum rows per request-safe query group, default one (no grouping). device : int, str, torch.device, or None Destination device for the returned offsets. By default, PyTorch's default tensor device is used. @@ -3350,6 +3358,7 @@ def _prepare_prims_ts_batch_decode_plan( page_size: Optional[int], kv_scale_factors: Optional[PagedKVScaleFactors] = None, q_token_kv_block_sparse_page_memberships: Optional[torch.Tensor] = None, + query_major_memberships: bool = False, use_q_token_kv_block_sparse_route: bool = False, use_pdl: bool = False, split_kv: bool = True, @@ -3357,6 +3366,7 @@ def _prepare_prims_ts_batch_decode_plan( direct_q1_max_model_len: Optional[int] = None, direct_q1_sparse_block_size: int = 4, share_pattern_across_kv_heads: bool = True, + allow_head_aware_locators: bool = False, ) -> tuple[PrimsTSBatchDecodePlan, torch.Tensor]: """Validate and freeze one dense-block-table PrimTS launch contract.""" @@ -3432,7 +3442,6 @@ def _prepare_prims_ts_batch_decode_plan( k_cache.dtype, v_cache.dtype, output_dtype, - allow_fp8_bf16_output=use_q_token_kv_block_sparse_route, ) k_sf_cache, v_sf_cache = _normalize_paged_kv_scale_factors( kv_scale_factors, @@ -3466,6 +3475,12 @@ def _prepare_prims_ts_batch_decode_plan( share_pattern_across_kv_heads, ) spec = _resolve_decode_launch_spec(*policy_args) + if query_major_memberships: + if not spec.config.supports_query_major_memberships: + raise ValueError( + "query-major memberships require one G2..G8 Keeps/page-4 CTA" + ) + spec = replace(spec, config=replace(spec.config, query_major_memberships=True)) if spec.config.uses_q_token_kv_block_sparse_page_membership: if q_token_kv_block_sparse_page_memberships is None: raise ValueError( @@ -3525,6 +3540,16 @@ def _prepare_prims_ts_batch_decode_plan( kv_block_size=direct_q1_sparse_block_size, page_capacity=max_num_pages, ) + flat_native_kv_tma = ( + use_q_token_kv_block_sparse_route + and spec.config.uses_staged_one_inst_tmem_p + and (spec.config.use_bf16_qkv or spec.config.use_fp8_qkv) + and spec.config.has_storage_subpages + and (num_kv_heads == 1 or (allow_head_aware_locators and not direct_q1_inputs)) + and k_cache.is_contiguous() + and normalized_cache.v_cache.is_contiguous() + and int(k_cache.shape[0]) * storage_page_size * num_kv_heads < (1 << 31) + ) compile_spec = _make_decode_compile_spec( spec, device_index=device_index, @@ -3542,6 +3567,7 @@ def _prepare_prims_ts_batch_decode_plan( kv_prefix_mode="dynamic", kv_lengths_mode="dynamic", direct_q1_spec=direct_q1_spec, + flat_native_kv_tma=flat_native_kv_tma, ) compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) workspace = _bind_decode_workspace(workspace_buffer, layout) @@ -3572,6 +3598,7 @@ def _prepare_prims_ts_batch_decode_plan( _compiled_main=compiled_main, _compiled_reducer=compiled_reducer, _direct_q1_inputs=direct_q1_inputs, + _flat_locator_heads=num_kv_heads if flat_native_kv_tma else 1, ) return plan, out diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py index 9b23b5dd7d38..0c5a915d0bcf 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py @@ -137,6 +137,8 @@ from collections.abc import Callable from typing import Any, Tuple +import importlib + import cutlass import cutlass.experimental.cuda as cuda import cutlass.cute as cute @@ -2821,6 +2823,9 @@ def __init__( ) validate_head_paired_head_ratio(head_paired=head_paired, h_r=h_r) if use_paged_kv: + # Import before task bodies enter staged GPU control flow. The + # helper module imports FmhaConfig, so it cannot load in resources. + importlib.import_module(f"{__package__}.helpers_paged") if num_tokens_per_page not in _SUPPORTED_CONTEXT_PAGE_SIZES: raise ValueError( "paged context requires num_tokens_per_page in " diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py index e6cdddf8d5cb..af69ee7d51bc 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py index da1e558acc79..78648bd18861 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py @@ -28,9 +28,12 @@ import functools from collections.abc import Callable +from dataclasses import dataclass import cutlass import cutlass.cute as cute +from ..._compile_cache import build_and_load_attention_kernel + import torch from cuda.bindings import driver as cuda_drv @@ -54,6 +57,29 @@ _ERROR_INVALID_PHYSICAL_PAGE_ID = 5 +@dataclass(frozen=True) +class _BlockSparseInspectCompileSpec: + """Static BSR geometry shared by equivalent inspection specializations.""" + + batch_size: int + num_pattern_heads: int + num_q_block_rows: int + num_kv_blocks: int + + +@dataclass(frozen=True) +class _PagedBlockSparseInspectCompileSpec: + """Static BSR geometry, paging divisors, and request-validation bounds.""" + + batch_size: int + num_pattern_heads: int + num_q_block_rows: int + minimum_seq_len_kv: int + max_seq_len_kv: int + kv_block_size: int + page_size: int + + @cute.jit def _validate_bsr_row_lane( block_indices: cute.Tensor, @@ -464,12 +490,19 @@ def compile_block_sparse_inspection( ) -> Callable[..., None]: """Compile one geometry specialization while keeping ``indices[nnz]`` dynamic. - Tensor ranks, dtypes, compact strides, and all attention geometry are part - of the specialization. Only the flat ``block_indices`` extent is symbolic; - tensor contents may of course vary between calls to the cached function. + The artifact key uses rounded Q/KV block counts: exact lengths and block + sizes that produce the same BSR tensor shape and index bounds share code. + ``block_indices[nnz]`` and tensor contents remain runtime inputs. """ num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + num_kv_blocks = (seq_len_kv + kv_block_size - 1) // kv_block_size + compile_spec = _BlockSparseInspectCompileSpec( + batch_size=batch_size, + num_pattern_heads=num_pattern_heads, + num_q_block_rows=num_q_block_rows, + num_kv_blocks=num_kv_blocks, + ) logical_nnz = cute.sym_int() stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) kernel = _InspectBlockSparseBsr( @@ -480,7 +513,8 @@ def compile_block_sparse_inspection( q_block_size=q_block_size, kv_block_size=kv_block_size, ) - with torch.cuda.device(device_index): + + def compile_kernel(): return cute.compile( kernel, _fake_compact( @@ -495,6 +529,16 @@ def compile_block_sparse_inspection( options=_COMPILE_OPTIONS, ) + with torch.cuda.device(device_index): + return build_and_load_attention_kernel( + "block_sparse", + "inspect", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(), + ) + @functools.cache def compile_paged_block_sparse_metadata_inspection( @@ -509,9 +553,23 @@ def compile_paged_block_sparse_metadata_inspection( kv_block_size: int, page_size: int, ) -> Callable[..., None]: - """Compile one paged metadata entry that launches request then live-BSR.""" + """Compile request and live-BSR validation keyed by its static geometry. + + Exact Q lengths sharing a Q-block count reuse one artifact. Page/block + divisors and inclusive length limits stay in the key because they become + constants in the generated validation code; table widths remain symbolic. + """ num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + compile_spec = _PagedBlockSparseInspectCompileSpec( + batch_size=batch_size, + num_pattern_heads=num_pattern_heads, + num_q_block_rows=num_q_block_rows, + minimum_seq_len_kv=minimum_seq_len_kv, + max_seq_len_kv=max_seq_len_kv, + kv_block_size=kv_block_size, + page_size=page_size, + ) logical_nnz = cute.sym_int() runtime_page_columns = cute.sym_int() runtime_page_row_stride = cute.sym_int64(divisibility=1) @@ -536,7 +594,7 @@ def compile_paged_block_sparse_metadata_inspection( inspect_bsr=inspect_bsr, ) - with torch.cuda.device(device_index): + def compile_kernel(): return cute.compile( inspect_paged_block_sparse_metadata, _fake_compact( @@ -559,6 +617,16 @@ def compile_paged_block_sparse_metadata_inspection( options=_COMPILE_OPTIONS, ) + with torch.cuda.device(device_index): + return build_and_load_attention_kernel( + "block_sparse", + "inspect_paged", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(), + ) + __all__ = [ "compile_block_sparse_inspection", diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py index af54b06004ce..42ba3d0e971d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py @@ -61,6 +61,7 @@ MMA_K_STEP_BYTES, PACKED_REGISTER_BYTES, PARALLEL_REDUCTION_BYTES_PER_SLICE, + PARALLEL_REDUCTION_SLOT_LANES, PARALLEL_REDUCTION_THREADS_PER_CTA, PARTIAL_O_ELEMENT_BYTES, PARTIAL_STATS_VALUES_PER_ROW, @@ -124,6 +125,7 @@ } _GROUPED_KEEPS_STATIC_ONLY_PROFILES = { (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float16, 128, 0, 2, 2), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, BFloat16, 128, 0, 2, 2), (BFloat16, BFloat16, BFloat16, BFloat16, 64, 0, 2, 2), (Float16, Float16, Float16, Float16, 256, 128, 1, 1), } @@ -152,7 +154,8 @@ KV_TILE_256_CORRECTION_TASK_REGISTERS = 152 # Sparse route unions reuse the dense Keeps staging: D64/D128 use two -# unstaged K/V instances; D256 uses one instance with D128 head bands. +# unstaged K/V instances; D256 uses one instance with D128 bands or a full +# D256 FP8 stage assembled from 128-byte swizzled chunks. _Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES = { ( dtype, @@ -172,6 +175,10 @@ (Float8E4M3FN, BFloat16), ) } +_Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES.update( + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, output, 256, 256, 1, 1) + for output in (Float16, BFloat16) +) _KV_TILE_256_PHYSICAL_DEFAULTS: Mapping[str, ConfigValue] = { "tmem_s_cols": 128, @@ -221,6 +228,7 @@ def _dtype_bytes(dtype: type) -> int: (BFloat16, BFloat16, BFloat16), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN), (Float8E4M3FN, Float8E4M3FN, Float16), + (Float8E4M3FN, Float8E4M3FN, BFloat16), } @@ -230,10 +238,29 @@ def _dtype_bytes(dtype: type) -> int: # families rather than shape exceptions: batch size, KV length, tile choice, # and legal GMEM split fanout remain unrestricted by this declaration. _GROUPED_KEEPS_PAGED_FP8_PROFILES = { + ( + Float8E4M3FN, + Float8E4M3FN, + Float8E4M3FN, + output_dtype, + head_dim, + 128 if head_dim == 256 else 0, + 1 if head_dim == 256 else 2, + 1 if head_dim == 256 else 2, + ) + for head_dim in (64, 128, 256) + for output_dtype in (Float16, BFloat16, Float8E4M3FN) +} + +# Automatic selection retains the established cost-model recipes. A working +# output epilogue alone does not qualify another MMA layout for that model; +# explicit kernel configurations can use the wider capability matrix above. +_AUTO_GROUPED_KEEPS_PAGED_FP8_PROFILES = { (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 64, 0, 2, 2), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 128, 0, 2, 2), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 256, 128, 1, 1), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float16, 256, 128, 1, 1), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, BFloat16, 256, 128, 1, 1), } @@ -607,6 +634,9 @@ class FmhaDecodeConfig: # sparse routes consume query membership from a separate packed-word table; # Q1 uses the same scattered page route without membership masking. use_q_token_kv_block_sparse_route: bool = False + # Private prepared-plan layout: eight query words per 32 page-4 IDs. + # Legacy public membership tensors keep the default page-byte encoding. + query_major_memberships: bool = False # None preserves each sparse family's default in raw/static configs too. share_pattern_across_kv_heads: bool | None = None # Allow the attention grid to acquire a programmatic launch dependency. @@ -834,6 +864,19 @@ def uses_q_token_kv_block_sparse_page_route(self) -> bool: and self.grouped_q_fits_tile ) + @property + def supports_query_major_memberships(self) -> bool: + """Whether this CTA can read eight-slot/page-4 query masks for G2..G8.""" + return ( + self.uses_q_token_kv_block_sparse_page_membership + and self.use_keeps_mma_ab + and 2 <= self.max_seq_len_q <= 8 + and self.has_single_q_cta + and self.num_tokens_per_page == 4 + and self.tile_size_kv == 128 + and self.tile_size_q in (64, 128) + ) + @property def uses_held_encoded_locator_window(self) -> bool: """Whether encoded paging keeps the complete CTA route in SMEM.""" @@ -872,14 +915,23 @@ def uses_task_register_reallocation(self) -> bool: @property def softmax_task_num_registers(self) -> int | None: + """Softmax's preferred grant, capped by the CTA's initial register pool. + + Dense and sparse Keeps kernels share this rule. Full 128-score rows + can use 232 registers only when the pool permits; active producer and + correction grants are unchanged. KV256 and transformed KV keep their + separate budgets. + """ if not self.uses_task_register_reallocation: return None if self.use_transform_kv: preferred = 168 if self.headdim == 64 and self.num_insts_kv != 1 else 184 else: - preferred = ( - KV_TILE_256_SOFTMAX_TASK_REGISTERS if self.tile_size_kv == 256 else 184 - ) + # Full Q128 rows keep 128 scores live through masking and P + # conversion. Give them headroom without taking registers from TMA. + preferred = 232 if self.softmax_score_fragment_regs == 128 else 184 + if self.tile_size_kv == 256: + preferred = KV_TILE_256_SOFTMAX_TASK_REGISTERS # setmaxnreg redistributes the CTA's initial register pool, not the # whole SM register file. Complete 8-register/thread quanta leave # unused SM registers when a CTA has more than 16 warps. @@ -903,6 +955,14 @@ def softmax_task_num_registers(self) -> int | None: transform_regs = self.transform_kv_task_num_registers assert transform_regs is not None available -= self.transform_kv_num_warps * (transform_regs - producer_regs) + # Only whole padding warpgroups may donate independently. Padding + # beside active tasks must execute the same setmaxnreg as its peers. + for wg_idx in range(MAX_WARP_GROUPS): + padding_warps = getattr(self, f"wg{wg_idx}_padding_num_warps") + if padding_warps == 4: + padding_regs = self.padding_task_num_registers(wg_idx * 4, 4) + assert padding_regs is not None + available += padding_warps * (producer_regs - padding_regs) return min(preferred, available // (softmax_warps * 8) * 8) @property @@ -931,6 +991,22 @@ def transform_kv_task_num_registers(self) -> int | None: return 88 return 56 + def padding_task_num_registers(self, warp_idx: int, num_warps: int) -> int | None: + """Release unused registers only from complete padding warpgroups.""" + if not self.uses_task_register_reallocation: + return None + if ( + not self.use_transform_kv + and self.tile_size_kv != 256 + and self.softmax_score_fragment_regs == 128 + and warp_idx % 4 == 0 + and num_warps == 4 + ): + # setmaxnreg's minimum legal grant; no attention work runs here. + # Smaller score fragments and KV transforms keep their budgets. + return 24 + return self.mma_load_task_num_registers + # ------------------------------------------------------------------ # SMEM allocation alignment # ------------------------------------------------------------------ @@ -952,6 +1028,13 @@ def q_tokens_per_cta(self) -> int: self.groups_tokens_heads_q, ) + @property + def q_tma_tokens_per_cta(self) -> int: + """Q tokens copied by a box, excluding spare sparse-group slots.""" + if self.use_variable_seqlens_q and self.uses_q_token_kv_block_sparse_page_route: + return min(self.max_seq_len_q, self.q_tokens_per_cta) + return self.q_tokens_per_cta + @property def num_q_ctas(self) -> int: """Return Q-CTA groups per ``(batch, KV head)`` launch tile. @@ -979,15 +1062,13 @@ def has_single_q_cta(self) -> bool: @property def q_tma_rows_per_cta(self) -> int: """Q rows represented by the configured tensor-map box.""" - return _q_tma_rows_per_cta( - self.tile_size_q, - self.heads_q_per_kv, - self.groups_tokens_heads_q, - ) + if self.groups_tokens_heads_q: + return self.q_tma_tokens_per_cta * self.heads_q_per_kv + return self.tile_size_q @property def q_manual_padding_rows(self) -> int: - """Grouped structural Q rows completed without a TMA instruction.""" + """Grouped Q rows whose transaction bytes are completed without TMA.""" if self.groups_tokens_heads_q: return self.tile_size_q - self.q_tma_rows_per_cta return 0 @@ -1076,6 +1157,23 @@ def num_head_dim_stages_kv(self) -> int: """Number of K/V head-dim stages needed to cover headdim.""" return (self.headdim + self.head_dim_kv_stage - 1) // self.head_dim_kv_stage + @property + def uses_2d_flat_kv_tma(self) -> bool: + """Whether warp-owned issuing can omit the flat map's singleton axes.""" + return ( + self.use_flat_native_kv_tma + and (self.tile_size_kv // self.num_tokens_per_page) % self.load_num_warps + == 0 + and ( + self.k_dtype == self.v_dtype == BFloat16 + or ( + self.use_fp8_qkv + and self.use_persistent_scheduler + and self.head_dim_kv_stage // 128 > 1 + ) + ) + ) + @property def tmem_o_cols_per_head_dim_stage(self) -> int: """TMEM O columns owned by one head-dim stage.""" @@ -1294,13 +1392,7 @@ def supports_reduction_dtypes(self) -> bool: and self.out_dtype in (Float16, BFloat16) ) or ( self.q_dtype == Float8E4M3FN - and ( - self.out_dtype in (Float16, Float8E4M3FN) - or ( - self.uses_q_token_kv_block_sparse_page_route - and self.out_dtype == BFloat16 - ) - ) + and self.out_dtype in (Float16, BFloat16, Float8E4M3FN) ) @property @@ -1354,6 +1446,11 @@ def use_fp8_kv(self) -> bool: # page-four table: # locator = physical_page * subpages_per_storage_page + subpage. storage_tokens_per_page: int = 0 + # Prepared compact caches can view physical pages and heads as tensor rows. + # The paired metadata plan selects the corresponding locator head span. + use_flat_native_kv_tma: bool = False + # Prepared locators count the complete physical head span between pages. + flat_native_kv_num_heads: int = 1 # Maximum number of pages per (batch, head_kv) — sizes the page index # table stride. max_num_pages_per_seq_kv: int = 1 @@ -1461,7 +1558,7 @@ def parallel_reduction_cluster_size(self) -> int: return 1 return { 8: 1, - 16: 2, + 16: 1, 32: 4, 64: 8, 128: 16, @@ -1477,12 +1574,23 @@ def parallel_reduction_splits_per_cta(self) -> int: // self.parallel_reduction_cluster_size ) + @property + def parallel_reduction_slot_lanes(self) -> int: + """Share a single-CTA fragment's split slots across adjacent lanes.""" + if ( + self.use_compact_parallel_reduction + or self.parallel_reduction_cluster_size != 1 + or self.parallel_reduction_splits_per_cta % PARALLEL_REDUCTION_SLOT_LANES + ): + return 1 + return PARALLEL_REDUCTION_SLOT_LANES + @property def parallel_reduction_threads_per_cta(self) -> int: """Return the thread count selected by the reducer schedule.""" if self.use_compact_parallel_reduction: return REDUCTION_THREADS_PER_CTA - return PARALLEL_REDUCTION_THREADS_PER_CTA + return PARALLEL_REDUCTION_THREADS_PER_CTA * self.parallel_reduction_slot_lanes @property def parallel_reduction_bytes_per_slice(self) -> int: @@ -2197,6 +2305,21 @@ def _kv256_dtypes_qualified(self) -> bool: return self._uses_contiguous_grouped_keeps_recipe return (self.q_dtype, self.k_dtype, self.out_dtype) in _KV256_DTYPE_RECIPES + @property + def supports_auto_grouped_keeps(self) -> bool: + """Whether grouped Keeps is also qualified for automatic selection.""" + if ( + self.use_paged_kv + and self.tile_size_kv == 128 + and self.q_dtype == self.k_dtype == self.v_dtype == Float8E4M3FN + ): + return ( + self._grouped_keeps_profile_key + in _AUTO_GROUPED_KEEPS_PAGED_FP8_PROFILES + and self.supports_grouped_keeps + ) + return self.supports_grouped_keeps + @property def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: """Whether inactive grouped rows are excluded solely at output stores. @@ -2210,6 +2333,7 @@ def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: profile = self._grouped_keeps_profile_key return ( self.use_keeps_mma_ab + and self.use_paged_kv and self.groups_tokens_heads_q and self.max_seq_len_q == 1 and not self.use_variable_seqlens_q @@ -2236,6 +2360,7 @@ def uses_guarded_grouped_keeps_output_rows(self) -> bool: profile = self._grouped_keeps_profile_key return ( self.use_keeps_mma_ab + and self.use_paged_kv and self.groups_tokens_heads_q and not self.use_variable_seqlens_q and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES @@ -2368,7 +2493,14 @@ def uses_staged_one_inst_tmem_p(self) -> bool: return ( self.use_keeps_mma_ab and self.headdim == 256 - and self.head_dim_per_stage_kv == 128 + and ( + self.head_dim_per_stage_kv == 128 + or ( + self.head_dim_per_stage_kv == 256 + and self.use_fp8_qkv + and self.uses_scattered_page_route + ) + ) and self.num_insts_kv == 1 and self.o_stages == 1 ) @@ -2873,7 +3005,7 @@ def supports_grouped_keeps(self) -> bool: if self.use_sage_attention: return False - if profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES: + if self.use_paged_kv and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES: fixed_q1 = ( self.max_seq_len_q == 1 and 1 <= self.heads_q_per_kv <= self.tile_size_q ) @@ -3826,28 +3958,19 @@ def _select_auto_launch_mode( def get_max_active_clusters_for_cluster_size(cluster_size: int) -> int: - """Query cluster occupancy after establishing CUDA's primary context. + """Query device occupancy, preserving the CPU-only policy fallback.""" + from ..._hardware import get_max_active_clusters - ``HardwareInfo`` compiles and retains a tiny occupancy-query module on its - first call. If that first call precedes CUDA context creation, retrying can - reuse an invalid module handle. Establish the context before constructing - ``HardwareInfo`` so a fresh Python process is reliable as well. - """ - cuda_context_ready = False try: import torch if torch.cuda.is_available(): - torch.cuda.init() - torch.empty(0, device="cuda") - cuda_context_ready = True - except (ImportError, RuntimeError): + return get_max_active_clusters(cluster_size) + except ImportError: pass try: return utils.HardwareInfo().get_max_active_clusters(cluster_size) except RuntimeError: - if cuda_context_ready: - raise return max(FALLBACK_SM_COUNT_B200 // max(cluster_size, 1), 1) @@ -4035,6 +4158,8 @@ def _try_apply_default_wide_keeps_config( probe, explicit_fields | {"tile_size_q"}, ) + if probe.groups_tokens_heads_q and not probe.supports_auto_grouped_keeps: + continue _validate_profile_support( cfg=probe, seq_len_q=seq_len_q, @@ -4164,6 +4289,8 @@ def _resolve_grouped_q_launch_candidates( probe, explicit_fields | {"tile_size_q"}, ) + if probe.use_keeps_mma_ab and not probe.supports_auto_grouped_keeps: + return () cost_num_insts_kv = probe.num_insts_kv _validate_profile_support( cfg=probe, @@ -4744,9 +4871,9 @@ def _validate_profile_support( heads_q_per_kv=heads_q_per_kv, groups_tokens_heads_q=True, ) - # Preserve the legacy fixed FP8 Q8/D128 profile. The D64/D256 extensions - # are deliberately limited to grouped page-32 decode; broader dtypes, - # head ratios, and variable-Q layouts remain unqualified. + # Fixed FP8 Q8 producers and reducers support both 16-bit output types. + # D64/D256 retain their grouped page-32 geometry; head ratios and + # variable-Q layouts remain constrained by the same resource profile. qualified_fp8_q8_separate_reduction_supported = ( not use_keeps_mma_ab and cfg.use_split_kv @@ -4756,14 +4883,14 @@ def _validate_profile_support( and cfg.k_dtype == Float8E4M3FN and cfg.v_dtype == Float8E4M3FN and ( - (headdim == 128 and cfg.out_dtype == Float8E4M3FN) + headdim == 128 or ( use_groups_tokens_heads_q and cfg.use_paged_kv and cfg.num_tokens_per_page == 32 and ( - (headdim == 64 and cfg.out_dtype == Float8E4M3FN) - or (headdim == 256 and cfg.out_dtype == Float16) + headdim == 64 + or (headdim == 256 and cfg.out_dtype in (Float16, BFloat16)) ) ) ) @@ -4781,7 +4908,7 @@ def _validate_profile_support( padded_splits = cfg.parallel_reduction_padded_splits default_cluster_size = { 8: 1, - 16: 2, + 16: 1, 32: 4, 64: 8, 128: 16, @@ -4794,7 +4921,7 @@ def _validate_profile_support( clustered_topology_supported = ( not cfg.use_compact_parallel_reduction and default_cluster_size == cluster_size - and splits_per_cta in (2, 4, 8) + and splits_per_cta in (2, 4, 8, 16) and cluster_size * splits_per_cta == padded_splits ) if not ( @@ -4913,22 +5040,27 @@ def _validate_profile_support( effective_head_dim_stage = cfg.head_dim_per_stage_kv effective_num_insts_kv = cfg.num_insts_kv effective_o_stages = cfg.o_stages + supported_head_stage = effective_head_dim_stage == 128 or ( + effective_head_dim_stage == 256 + and cfg.use_fp8_qkv + and cfg.uses_q_token_kv_block_sparse_page_route + ) if effective_num_insts_kv == 1 and ( - headdim != 256 or effective_head_dim_stage != 128 or effective_o_stages != 1 + headdim != 256 or not supported_head_stage or effective_o_stages != 1 ): raise ValueError( "one-instance KeepsMmaAb is enabled only for the staged " - "headDim=256 profile with head_dim_per_stage_kv=128 and " + "headDim=256 profile with D128 bands (or D256 for sparse FP8) and " "o_stages=1" ) if headdim == 256 and ( - effective_head_dim_stage != 128 + not supported_head_stage or effective_num_insts_kv != 1 or effective_o_stages != 1 ): raise ValueError( "fmha_decode keepsMmaAb headDim=256 requires " - "head_dim_per_stage_kv=128, num_insts_kv=1, and o_stages=1" + "D128 bands (or D256 for sparse FP8), num_insts_kv=1, and o_stages=1" ) if headdim != 256 and effective_head_dim_stage != 0: raise ValueError( @@ -4939,20 +5071,6 @@ def _validate_profile_support( raise ValueError( "fmha_decode keepsMmaAb requires numHeadsQPerKv == tile_size_q" ) - fp8_q_token_kv_block_sparse_bf16_output = ( - cfg.out_dtype == BFloat16 - and use_groups_tokens_heads_q - and is_q_token_kv_block_sparse_grouped_keeps - ) - if ( - cfg.q_dtype == Float8E4M3FN - and cfg.out_dtype not in (Float16, Float8E4M3FN) - and not fp8_q_token_kv_block_sparse_bf16_output - and not cfg.use_sage_attention - ): - raise ValueError( - "fmha_decode keepsMmaAb fp8 qkv path supports fp16 or fp8 output" - ) use_split_kv = split_kv_mode != "disabled" or cfg.use_split_kv if use_split_kv: if cfg.q_dtype not in (Float16, BFloat16, Float8E4M3FN): @@ -4987,7 +5105,7 @@ def _validate_profile_support( and cfg.uses_q_token_kv_block_sparse_page_route ) ) - and effective_head_dim_stage == 128 + and supported_head_stage and effective_num_insts_kv == 1 and effective_o_stages == 1 ) @@ -5002,7 +5120,7 @@ def _validate_profile_support( and cfg.q_dtype == Float8E4M3FN and cfg.k_dtype == Float8E4M3FN and cfg.v_dtype == Float8E4M3FN - and cfg.out_dtype == (Float16 if headdim == 256 else Float8E4M3FN) + and (headdim == 64 or cfg.out_dtype in (Float16, BFloat16)) and cfg.splits_kv == cfg.max_splits_kv and 2 <= cfg.max_splits_kv <= 128 ) @@ -5049,7 +5167,7 @@ def _validate_profile_support( "{64,128,256}, TileSizeQ in {8,16,32}, " "either ungrouped single-token or complete-token grouped Q, " "at least two split CTAs, " - "and fp16/bf16 or fp8 qkv with fp16/fp8 output" + "and fp16/bf16 or fp8 qkv with fp16/bf16/fp8 output" ) cluster_reason = cluster_smem_reduction_unsupported_reason( max_splits_kv=cfg.max_splits_kv, @@ -5078,9 +5196,9 @@ def _validate_profile_support( raise ValueError( "separate reduction SwapsMmaAb profiles require fixed SQ=1, " "fixed grouped Q, or packed variable Q; tile_size_q in {16,32}, " - "QToken-KvBlock-Sparse-Attention TileQ8, or a fixed FP8 Q8/HqPerKv8 profile (legacy D128 with FP8 output, " - "plus grouped paged-KV/page-32 D64 with FP8 output or D256 with " - "FP16 output); valid reduction dtypes and static split-KV are required" + "QToken-KvBlock-Sparse-Attention TileQ8, or a fixed FP8 Q8/HqPerKv8 profile " + "(D128, plus grouped paged-KV/page-32 D64 or D256 with " + "FP16/BF16 output); valid reduction dtypes and static split-KV are required" ) if tile_size_q and tile_size_q not in (8, 16, 32): raise ValueError("fmha_decode SwapsMmaAb supports tile_size_q in {8,16,32}") diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py index 909de7947f81..f120a13f02cb 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py @@ -186,14 +186,21 @@ # Clustered standalone reducer shape. One 128-thread CTA covers a contiguous # 2 KiB partial-O slice with one 16-byte vector per thread. Each cluster rank -# owns a compile-time 2, 4, or 8 split slots. Loads are batched in groups of at -# most four; padded split slots remain neutral and never form GMEM pointers. +# owns a compile-time 2, 4, or 8 split slots. Padded and unpublished slots +# remain neutral and never form GMEM pointers. PARALLEL_REDUCTION_THREADS_PER_CTA = 128 PARALLEL_REDUCTION_BYTES_PER_SLICE = ( PARALLEL_REDUCTION_THREADS_PER_CTA * REDUCTION_BYTES_PER_THREAD ) -PARALLEL_REDUCTION_LOAD_BATCH = 4 +# Batch independent published-slot loads before the serial fold. +PARALLEL_REDUCTION_LOAD_BATCH = 16 PARALLEL_REDUCTION_FINAL_REDUCERS = 4 +# Single-CTA reducer profiles spread one output fragment's split slots +# over this many adjacent lanes and merge them with warp shuffles. The fold is +# a serial chain of dependent FP32 instructions at very low occupancy, so more +# lanes with fewer slots shorten that per-lane dependency chain. Must divide +# the slots owned by each CTA. +PARALLEL_REDUCTION_SLOT_LANES = 4 # Each reduction thread produces an 8-element O vector backed by four packed # 16-bit registers. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py index 2292eb6544bb..7a2235d8bc28 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py @@ -134,10 +134,9 @@ ) _PERSISTENT_SCHEDULE_TOKEN_STAGES = 2 -# Response slots of the CLC work-queue fetch pipeline. Every role consumes -# the queue one tile behind the scheduler warp, so a second slot adds no -# lookahead, only a loop-carried stage index that spills to local memory. -_WORK_QUEUE_STAGES = 1 +# Double-buffer CLC responses independently of the schedule-token throttle. +# A single response slot triggers register spills in grouped sparse work loops. +_WORK_QUEUE_STAGES = 2 def _block_sparse_bshd_tma_strides( @@ -404,7 +403,7 @@ def _build_decode_gen_schedule( if cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1 and not cfg.uses_tmem_p: raise ValueError( "one-instance KeepsMmaAb is enabled only for the staged headDim=256 " - "profile with head_dim_per_stage_kv=128 and o_stages=1" + "profile with D128 bands (or D256 for sparse FP8) and o_stages=1" ) if use_native_paged_kv and not cfg.use_paged_kv: raise ValueError("native paged-KV ABI requires cfg.use_paged_kv=True") @@ -534,7 +533,18 @@ def _build_decode_gen_schedule( cfg.k_dtype_bytes != cfg.v_dtype_bytes and cfg.tile_size_kv != 256 ) one_inst_tmem_stages = 2 if use_one_inst_qkv else 1 - one_inst_kv_stages = cfg.num_head_dim_stages_kv if use_one_inst_qkv else 1 + # Keep complete head-dimension rounds in each balanced K/V ring. FP8's + # smaller stage tiles can prefetch farther without increasing the budget. + one_inst_kv_stages = ( + max( + cfg.num_head_dim_stages_kv, + cfg.kv_stages + // (2 * cfg.num_head_dim_stages_kv) + * cfg.num_head_dim_stages_kv, + ) + if use_one_inst_qkv + else 1 + ) use_distributed_split_kv_stages = not use_one_inst_qkv if cfg.tile_size_q == 128 and use_distributed_split_kv_stages: # Q128's four instruction-local K0/K1/V0/V1 rings need equal depth. @@ -836,15 +846,14 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: advance_on_wait=True, ) - # Two-instance Keeps keeps stats outside S and orders each same-instance PV - # before the next QK, so it needs no stats-done credit. - # The staged one-instance path needs an overwrite-credit gate across its - # double-buffered S/P overlay: correction returns the stage credit before - # MMA can reissue QK into those columns. + # The single MMA issuer orders PV before QK reuses the same S/P slot. + # PTX's same-warp mma(A-read) -> mma(D-write) pipeline protects that alias. + # Separate SMEM stats have their own softmax-local handoff, just as in + # two-instance Keeps; correction need not additionally gate S reuse. stats_done0_cfg = None stats_done1_cfg = None resource_dependency_graph: dict[MemoryResource, list[MemoryResource]] - if use_one_inst_qkv: + if use_one_inst_qkv and (not cfg.keeps_stats_via_smem or cfg.mma_num_warps != 1): stats_done0_cfg = PipelineConfig.create_async_async_pipeline_cfg( num_stages=one_inst_tmem_stages, producer_group=mma_grp, @@ -3375,12 +3384,6 @@ def fmha_decode_launch( if cutlass.const_expr(use_native_paged_kv): total_pages = num_physical_kv_pages storage_tokens_per_page = Int32(cfg.effective_storage_tokens_per_page) - kv_shape = ( - kv_tma_d, - storage_tokens_per_page, - h_k, - total_pages, - ) k_tma_page_stride = k_page_stride k_tma_head_stride = k_head_stride k_tma_token_stride = k_token_stride @@ -3396,24 +3399,60 @@ def fmha_decode_launch( v_tma_page_stride = v_page_stride * Int64(2) v_tma_head_stride = v_head_stride * Int64(2) v_tma_token_stride = v_token_stride * Int64(2) - k_layout = cute.make_layout( - kv_shape, - stride=( - 1, - k_tma_token_stride, - k_tma_head_stride, - k_tma_page_stride, - ), - ) - v_layout = cute.make_layout( - kv_shape, - stride=( - 1, - v_tma_token_stride, - v_tma_head_stride, - v_tma_page_stride, - ), - ) + if cutlass.const_expr(cfg.use_flat_native_kv_tma): + # Reinterpret compact HND rows, without copying/repacking KV + # or requiring contiguous sparse selections. Planning checked + # signed-32-bit row bounds; other bindings use the general map. + # For S tokens/page, F tokens/fragment and H heads, metadata + # emits A = page*H*(S/F) + subpage. Attention addresses row + # (A + head*(S/F))*F, avoiding per-copy page/subpage divmod. + # Bytes and TMA count are unchanged; rank-2 issuing below is a + # separate optimization of this descriptor-only flat view. + flat_tokens = ( + num_physical_kv_pages + * Int64(storage_tokens_per_page) + * Int64(cfg.flat_native_kv_num_heads) + ) + kv_shape = (kv_tma_d, flat_tokens, Int32(1), Int32(1)) + k_tma_head_stride = k_tma_token_stride * flat_tokens + k_tma_page_stride = k_tma_head_stride + v_tma_head_stride = v_tma_token_stride * flat_tokens + v_tma_page_stride = v_tma_head_stride + else: + kv_shape = ( + kv_tma_d, + storage_tokens_per_page, + h_k, + total_pages, + ) + if cutlass.const_expr(cfg.uses_2d_flat_kv_tma): + # The prepared flat view has no head/page coordinates. Keep + # the same physical rows and SW128 tile without repacking KV. + k_layout = cute.make_layout( + (kv_shape[0], kv_shape[1]), stride=(1, k_tma_token_stride) + ) + v_layout = cute.make_layout( + (kv_shape[0], kv_shape[1]), stride=(1, v_tma_token_stride) + ) + else: + k_layout = cute.make_layout( + kv_shape, + stride=( + 1, + k_tma_token_stride, + k_tma_head_stride, + k_tma_page_stride, + ), + ) + v_layout = cute.make_layout( + kv_shape, + stride=( + 1, + v_tma_token_stride, + v_tma_head_stride, + v_tma_page_stride, + ), + ) k_tma = cute.make_tensor(k_iter, k_layout) v_tma = cute.make_tensor(v_iter, v_layout) else: @@ -3496,8 +3535,8 @@ def fmha_decode_launch( q_box_dims: tuple[object, ...] if cutlass.const_expr(cfg.use_variable_seqlens_q): # Packed Q is physically [sum_q_tokens, num_heads_q, head_dim]. Flatten - # the two head modes into Hq so the ragged token axis still fits in a - # rank-5 tensor map after the helper inserts its two synthetic modes. + # the head modes into Hq. Sparse Q uses this rank-3 map directly; + # other packed paths retain per-sequence ragged zero fill. q_tma = cute.make_tensor( q_iter, cute.make_layout( @@ -3509,7 +3548,7 @@ def fmha_decode_launch( q_box_dims = ( tma_box0_q, cfg.heads_q_per_kv, - cfg.q_tokens_per_cta, + cfg.q_tma_tokens_per_cta, ) q_groups = Int32( (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta @@ -3520,13 +3559,24 @@ def fmha_decode_launch( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) q_groups = Int32(cfg.max_seq_len_q) * head_ctas_per_token - tma_desc_q = create_tensor_map_ragged_from_tensor( - q_tma, - box_dims=q_box_dims, - ragged_dim=2, - stride_order=(0, 1, 2), - swizzle=tma_swizzle_q, - ) + if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_route): + # A partial group may read the next sequence into inactive Q rows. + # QK/softmax/PV rows are independent; row masks and guarded stores + # discard them. Physical allocation bounds still get TMA zero fill. + tma_desc_q = create_tensor_map_tiled_from_view( + q_tma, + box_dims=q_box_dims, + stride_order=(0, 1, 2), + swizzle=tma_swizzle_q, + ) + else: + tma_desc_q = create_tensor_map_ragged_from_tensor( + q_tma, + box_dims=q_box_dims, + ragged_dim=2, + stride_order=(0, 1, 2), + swizzle=tma_swizzle_q, + ) else: q_tma = cute.make_tensor( q_iter, @@ -3557,19 +3607,27 @@ def fmha_decode_launch( stride_order=(0, 1, 2, 3, 4), swizzle=tma_swizzle_q, ) + if cutlass.const_expr(cfg.uses_2d_flat_kv_tma): + k_box_dims = (tma_box0_k, tma_kv_tokens) + v_box_dims = (tma_box0_v, tma_kv_tokens) + kv_stride_order = (0, 1) + else: + k_box_dims = (tma_box0_k, tma_kv_tokens, 1, 1) + v_box_dims = (tma_box0_v, tma_kv_tokens, 1, 1) + kv_stride_order = (0, 1, 2, 3) if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): tma_desc_k = create_tensor_map_tiled_from_view( k_tma, - box_dims=(tma_box0_k, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=k_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_k, dtype=cutlass.Float4E2M1FNx2, tma_format=cuda.TensorMapDataFormat.B4X16_P64, ) tma_desc_v = create_tensor_map_tiled_from_view( v_tma, - box_dims=(tma_box0_v, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=v_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_v, dtype=cutlass.Float4E2M1FNx2, tma_format=cuda.TensorMapDataFormat.B4X16_P64, @@ -3577,14 +3635,14 @@ def fmha_decode_launch( else: tma_desc_k = create_tensor_map_tiled_from_view( k_tma, - box_dims=(tma_box0_k, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=k_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_k, ) tma_desc_v = create_tensor_map_tiled_from_view( v_tma, - box_dims=(tma_box0_v, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=v_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_v, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py index fe49a7c67592..3e2031504ef7 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py @@ -94,6 +94,26 @@ from .smem_block_sparse_metadata import SmemBlockSparseKvMetadataResource +@cute.jit +def _q_token_membership_keep_word16( + word0: Uint32, + word1: Uint32, + word2: Uint32, + word3: Uint32, + query_idx: Int32, +) -> Uint32: + """Extract one query's membership bit from sixteen packed page bytes.""" + keep_word = Uint32(0) + shift = query_idx & Int32(7) + words = (word0, word1, word2, word3) + for word_idx in cutlass.range_constexpr(4): + bits = (words[word_idx] >> shift) & Uint32(0x01010101) + # The disjoint partial products gather four byte bits without carries. + nibble = (bits * Uint32(0x01020408)) >> Uint32(24) + keep_word |= nibble << Uint32(word_idx * 4) + return cutlass.select_(Uint32(query_idx) < Uint32(8), keep_word, Uint32(0)) + + def _paged_sparse_kv_tma_transaction_geometry( *, tile_size_kv: int, @@ -158,8 +178,20 @@ def _kv_dtype_bytes_for_kind( def _decode_native_page_locator( cfg: Constexpr[FmhaDecodeConfig], page_locator: Int32, + kv_head: Int32 = 0, ) -> tuple[Int32, Int32]: - """Return ``(token_offset, physical_page)`` for one native table entry.""" + """Return token/page TMA coordinates for one prepared locator.""" + if cutlass.const_expr(cfg.use_flat_native_kv_tma): + # The prepared locator is the head-zero origin in the compact cache. + # Keep the entire fragment OOB when adding a head offset to a sentinel. + if cutlass.const_expr(cfg.flat_native_kv_num_heads > 1): + head_fragments = ( + cfg.effective_storage_tokens_per_page // cfg.num_tokens_per_page + ) + page_locator = (page_locator + kv_head * Int32(head_fragments)) | ( + page_locator >> Int32(31) + ) + return page_locator * Int32(cfg.num_tokens_per_page), Int32(0) token_offset = Int32(0) physical_page = page_locator if cutlass.const_expr(cfg.has_storage_subpages): @@ -224,14 +256,84 @@ def _issue_sparse_page_copies( head_dim_stage: Constexpr[int], head_dim_stage_offset: Constexpr[int], ): - """Partition (D64 chunk, physical fragment) copies across loader lanes. + """Issue physical page fragments through the selected loader schedule. The cache writer zero-fills page padding: full-fragment tail loads rely on this because score masking alone does not sanitize NaN/Inf in V. """ - chunk_hd = min(head_dim_stage, 64) + chunk_hd = min(head_dim_stage, 128 if cfg.use_fp8_qkv else 64) chunks = head_dim_stage // chunk_hd fragments = cfg.tile_size_kv // cfg.num_tokens_per_page + if cutlass.const_expr( + cfg.uses_2d_flat_kv_tma and cfg.k_dtype == cfg.v_dtype == cutlass.BFloat16 + ): + # Prefetch warp-owned locators before TMA issue and reuse each across + # BF16's head chunks. Static indexing keeps the small array in registers. + fragments_per_warp = fragments // cfg.load_num_warps + warp = _load_task_warp_rank(cfg) + if prims.elect_sync(): + locators = cutlass.Array( + Int32, (fragments_per_warp,), space=cutlass.AddressSpace.rmem + ) + for i in cutlass.range_constexpr(fragments_per_warp): + fragment = warp * Int32(fragments_per_warp) + Int32(i) + locators[i] = page_offsets.page_id(tile_idx, local_tile_idx, fragment) + for i in cutlass.range_constexpr(fragments_per_warp): + fragment = warp * Int32(fragments_per_warp) + Int32(i) + token_offset, _ = _decode_native_page_locator(cfg, locators[i], kv_head) + for chunk in cutlass.range_constexpr(chunks): + smem_offset = Int32( + chunk * chunk_hd * cfg.tile_size_kv + ) + fragment * Int32(chunk_hd * cfg.num_tokens_per_page) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + tma_desc, + ( + Int32(head_dim_stage_offset + chunk * chunk_hd), + token_offset, + ), + barrier, + ) + return + if cutlass.const_expr(cfg.uses_2d_flat_kv_tma): + # A warp owns complete fragments across head planes. Its staged token + # coordinate is uniform: expose that before election so both TMA copies + # reuse uniform coordinates instead of serializing lane operands. + # General addressing retains its lane-parallel decode. One-shot + # kernels keep parallel shared loads to minimize startup latency. + uniform_fragments_per_warp = fragments // cfg.load_num_warps + uniform_warp = _load_task_warp_rank(cfg) + uniform_first_fragment = uniform_warp * Int32(uniform_fragments_per_warp) + uniform_page_ptr = page_offsets.page_id_pointer( + tile_idx, local_tile_idx, uniform_first_fragment + ) + uniform_destination_ptr = stage_base.data_ptr( + uniform_first_fragment * Int32(chunk_hd * cfg.num_tokens_per_page) + ) + # Reuse a compact set of uniform address operands across fragments; + # advance staged locators and SMEM destinations independently of the + # scattered global coordinates. Head-plane copies stay unrolled. + for _fragment_idx in cutlass.range(uniform_fragments_per_warp, unroll=1): + uniform_token_offset = cute.arch.make_warp_uniform( + Int32(uniform_page_ptr[0]) + ) + if prims.elect_sync(): + for head_chunk in cutlass.range_constexpr(chunks): + prims.cp_async_bulk_tensor_shared_cta_global( + uniform_destination_ptr + + Int32(head_chunk * chunk_hd * cfg.tile_size_kv), + tma_desc, + ( + Int32(head_dim_stage_offset + head_chunk * chunk_hd), + uniform_token_offset, + ), + barrier, + ) + uniform_page_ptr = uniform_page_ptr + Int32(1) + uniform_destination_ptr = uniform_destination_ptr + Int32( + chunk_hd * cfg.num_tokens_per_page + ) + return copies = fragments * chunks copies_per_warp = (copies + cfg.load_num_warps - 1) // cfg.load_num_warps lane = cute.arch.thread_idx()[0] & Int32(31) @@ -243,7 +345,9 @@ def _issue_sparse_page_copies( chunk = copy // Int32(fragments) fragment = copy % Int32(fragments) locator = page_offsets.page_id(tile_idx, local_tile_idx, fragment) - token_offset, physical_page = _decode_native_page_locator(cfg, locator) + token_offset, physical_page = _decode_native_page_locator( + cfg, locator, kv_head + ) smem_offset = chunk * Int32(chunk_hd * cfg.tile_size_kv) + fragment * Int32( chunk_hd * cfg.num_tokens_per_page ) @@ -253,7 +357,9 @@ def _issue_sparse_page_copies( ( Int32(head_dim_stage_offset) + chunk * Int32(chunk_hd), token_offset, - kv_head, + Int32(0) + if cutlass.const_expr(cfg.use_flat_native_kv_tma) + else kv_head, physical_page, ), barrier, @@ -338,7 +444,7 @@ def _create_initial_task_locals( leading_byte_offset = q_leading_bytes stride_byte_offset = Int32(1024) if cutlass.const_expr(self.cfg.q_dtype_bytes == 1): - q_head_dim_stage = self.cfg.head_dim_kv_stage + q_head_dim_stage = min(self.cfg.head_dim_kv_stage, 128) q_tile_bytes = Int32( self.cfg.tile_size_q * q_head_dim_stage * self.cfg.q_dtype_bytes ) @@ -389,7 +495,7 @@ def _tma_coords( logical_h_k_idx: Int32, logical_b_idx: Int32, head_dim_offset: Int32, - ) -> tuple[Int32, Int32, Int32, Int32, Int32]: + ) -> tuple[Int32, ...]: """Map a logical Q CTA to token-major tensor-map coordinates.""" cfg = self.cfg if cutlass.const_expr(cfg.use_variable_seqlens_q): @@ -407,6 +513,12 @@ def _tma_coords( cfg.heads_q_per_kv ) + head_cta_idx * Int32(cfg.tile_size_q) tokens_per_box = 1 + if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_route): + return ( + head_dim_offset, + global_head_idx, + self.q_token_offset + token_idx, + ) packed_coords = transform_ragged_coords( ( head_dim_offset, @@ -459,7 +571,7 @@ def _tma_coords( @cute.jit def _complete_grouped_q_padding(self, stage_info: StageInfo) -> None: - """Account structural grouped rows that have no corresponding TMA.""" + """Credit the Q tile rows not covered by the configured TMA box.""" cfg = self.cfg if cutlass.const_expr( cfg.groups_tokens_heads_q and cfg.q_manual_padding_rows > 0 @@ -486,7 +598,7 @@ def tma_load(self, stage_info: StageInfo) -> None: stage_base = self._smem_base_q.subview(stage_info.stage_idx * stage_elems) if cutlass.const_expr(cfg.q_dtype_bytes == 1): if _is_load_task_warp_leader(cfg): - if cutlass.const_expr(cfg.num_head_dim_stages_kv == 1): + if cutlass.const_expr(cfg.headdim <= 128): # FP8 with one head-dim stage is one tensor copy into the # complete staged Q tile. prims.cp_async_bulk_tensor_shared_cta_global( @@ -503,9 +615,9 @@ def tma_load(self, stage_info: StageInfo) -> None: else: # H256 FP8 stages Q by head-dim slices so each QK head-dim # stage sees a contiguous SMEM tile. - q_chunk_dim = cfg.head_dim_kv_stage + q_chunk_dim = 128 for q_chunk_idx in cutlass.range_constexpr( - cfg.num_head_dim_stages_kv + cfg.headdim // q_chunk_dim ): head_dim_offset = q_chunk_idx * q_chunk_dim smem_offset = head_dim_offset * cfg.tile_size_q @@ -734,9 +846,6 @@ def _create_initial_task_locals( # FP16, BF16, and FP8. self._desc_base = prims.Tcgen05SmemDesc(0) return {"kv_desc": cutlass.Int64(0), "v_desc": cutlass.Int64(0)} - kv_tile_bytes = Int32( - self.cfg.tile_size_kv * self.cfg.head_dim_kv_stage * inst_dtype_bytes - ) leading_byte_offset = Int32( self.cfg.tile_size_kv * min(self.cfg.head_dim_kv_stage, 64) @@ -744,13 +853,18 @@ def _create_initial_task_locals( ) stride_byte_offset = Int32(1024) if cutlass.const_expr(inst_dtype_bytes == 1): - leading_byte_offset = kv_tile_bytes + leading_byte_offset = Int32( + self.cfg.tile_size_kv * min(self.cfg.head_dim_kv_stage, 128) + ) stride_byte_offset = Int32( _major_k_stride_bytes(inst_dtype_bytes, self.cfg.head_dim_kv_stage) ) if cutlass.const_expr( self.kv_kind == KV_KIND_V - and (inst_dtype_bytes == 1 or self.cfg.headdim == 64) + and ( + (inst_dtype_bytes == 1 and self.cfg.head_dim_kv_stage <= 128) + or self.cfg.headdim == 64 + ) ): leading_byte_offset = Int32(0) inst_swizzle = prims.Tcgen05SmemSwizzle.SWIZZLE_128B @@ -1152,6 +1266,22 @@ def _producer_load( tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) inst_dtype_bytes = _kv_dtype_bytes_for_kind(self.cfg, self.kv_kind) if cutlass.const_expr(inst_dtype_bytes == 1): + if cutlass.const_expr( + cfg.use_fp8_qkv and cfg.uses_scattered_page_route + ): + _issue_sparse_page_copies( + cfg, + self.page_offsets_kv, + self._stage_base(stage_info), + tma_desc, + stage_info.barrier, + tile_idx, + local_tile_idx, + logical_h_k_idx, + head_dim_stage, + head_dim_stage_offset, + ) + return if prims.elect_sync() and _load_task_warp_rank(cfg) < Int32( min(page_fragments, cfg.load_num_warps) ): @@ -1162,7 +1292,6 @@ def _producer_load( # elected lane replay the full loop over-completes the TMA # transaction barrier when load_num_warps > 1. stage_base = self._stage_base(stage_info) - page_ids = self.page_offsets_kv.page_ids(tile_idx, local_tile_idx) sf_desc = ( self.tma_desc_v_sf if cutlass.const_expr(self.kv_kind == KV_KIND_V) @@ -1182,8 +1311,14 @@ def _producer_load( transactions_per_warp ): page_frag = first_page_frag + Int32(warp_transaction) + # A runtime warp rank cannot index a register vector + # without materializing it in local memory. Read only + # this issuer's locator from the retained SMEM window. + locator = self.page_offsets_kv.page_id( + tile_idx, local_tile_idx, page_frag + ) token_offset, physical_page = _decode_native_page_locator( - cfg, Int32(page_ids[page_frag]) + cfg, locator ) smem_page_offset = Int32( page_frag @@ -1658,10 +1793,10 @@ def page_ids(self, tile_idx: Int32, local_tile_idx: Int32) -> cutlass.Array: return self._smem_page_offsets.load(offset, vector_size=1, alignment=4) @cute.jit - def page_id( + def page_id_pointer( self, tile_idx: Int32, local_tile_idx: Int32, page_frag: Int32 - ) -> Int32: - """Load a tile's locator; offset stages are independent of K/V data stages.""" + ) -> cutlass.Pointer: + """Locate a staged ID independently of the K/V data-buffer stage.""" cfg = self.cfg pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page if cutlass.const_expr(self.holds_encoded_locator_window): @@ -1675,7 +1810,14 @@ def page_id( else: group_page_idx = (tile_idx * Int32(pages_per_tile)) & Int32(31) offset = self.consumer_work_stage * Int32(32) + group_page_idx - return Int32(self._smem_page_offsets[offset + page_frag]) + return self._smem_page_offsets.data_ptr(offset + page_frag) + + @cute.jit + def page_id( + self, tile_idx: Int32, local_tile_idx: Int32, page_frag: Int32 + ) -> Int32: + """Load one locator from the active page-offset stage.""" + return Int32(self.page_id_pointer(tile_idx, local_tile_idx, page_frag)[0]) @cute.jit def _q_token_kv_block_sparse_membership_word( @@ -1761,37 +1903,101 @@ def q_token_kv_block_sparse_page_membership( ) @cute.jit - def q_token_kv_block_sparse_page_memberships4( + def query_membership_keep_word( + self, stage_info: StageInfo, local_tile_idx: Int32, q_token_idx: Int32 + ) -> Uint32: + """Read a private, tail-padded query bit-plane without transposing it.""" + cfg = self.cfg + assert cfg.query_major_memberships and cfg.num_tokens_per_page == 4 + assert cfg.tile_size_kv == 128 and 2 <= cfg.max_seq_len_q <= 8 + query = q_token_idx & Int32(7) + word = Uint32(0) + if cutlass.const_expr(self.cache_memberships_in_smem): + self._create_initial_task_locals(stage_info.context) + word = Uint32( + self._smem_q_token_kv_block_sparse_memberships[ + local_tile_idx * Int32(8) + query + ] + ) + else: + task_cache = _decode_gen_task_cache(stage_info) + tile = local_tile_idx + Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + if tile * Int32(32) <= Int32(task_cache[_TASK_CACHE_KV_PAGE_IDX_UB]): + head, batch = _logical_head_batch(stage_info, self.h_k_idx, self.b_idx) + row = cutlass.Int64(batch) + if cutlass.const_expr( + cfg.use_persistent_scheduler and not cfg.shares_sparse_pattern + ): + row = row * cutlass.Int64(self.num_heads_kv) + cutlass.Int64(head) + word = Uint32( + self.q_token_kv_block_sparse_page_memberships[ + row + * cutlass.Int64( + self.q_token_kv_block_sparse_page_membership_stride + ) + + cutlass.Int64(tile * Int32(8) + query) + ] + ) + return cutlass.select_(Uint32(q_token_idx) < Uint32(8), word, Uint32(0)) + + @cute.jit + def q_token_kv_block_sparse_page_membership_word4( self, stage_info: StageInfo, local_tile_idx: Int32, page_frag: Int32, - ) -> cutlass.Array: - """Load four aligned grouped-Q membership bytes from one packed word. - - The ``tmem_s`` Keeps Softmax path calls this helper four times to - assemble each lane-local 16-page membership word. - """ + ) -> Uint32: + """Load the packed word for four aligned grouped-Q page memberships.""" assert self.cfg.uses_q_token_kv_block_sparse_page_membership pages_per_tile = Int32(self.cfg.tile_size_kv // self.cfg.num_tokens_per_page) membership_idx = local_tile_idx * pages_per_tile + page_frag - memberships = cutlass.Array( - Uint32, - 4, - space=cutlass.AddressSpace.rmem, - ) membership_word, _ = self._q_token_kv_block_sparse_membership_word( stage_info, membership_idx ) - for elem_idx in cutlass.range_constexpr(4): - memberships[elem_idx] = ( - membership_word - >> Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS * elem_idx) - ) & Uint32( - Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_MASK, + return membership_word + + @cute.jit + def q_token_kv_block_sparse_page_keep_word16( + self, + stage_info: StageInfo, + local_tile_idx: Int32, + page_frag: Int32, + query_idx: Int32, + ) -> Uint32: + """Extract sixteen memberships, widening only padded SMEM reads.""" + assert self.cfg.uses_q_token_kv_block_sparse_page_membership + assert not self.cfg.query_major_memberships + if cutlass.const_expr(self.cache_memberships_in_smem): + self._create_initial_task_locals(stage_info.context) + pages_per_tile = Int32( + self.cfg.tile_size_kv // self.cfg.num_tokens_per_page + ) + word_idx = (local_tile_idx * pages_per_tile + page_frag) // Int32(4) + # Callers select aligned sixteen-page windows; the producer owns + # and initializes the full padded SMEM tile, including its tail. + words = ( + self._smem_q_token_kv_block_sparse_memberships.subview(word_idx) + .data_ptr() + .load(count=4, alignment=16) + ) + word0, word1, word2, word3 = words[0], words[1], words[2], words[3] + else: + # Public GMEM rows need not have sixteen readable tail bytes. + # Preserve the existing live-prefix checks for each packed word. + word0 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + ) + word1 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + Int32(4) + ) + word2 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + Int32(8) + ) + word3 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + Int32(12) ) - return memberships + return _q_token_membership_keep_word16(word0, word1, word2, word3, query_idx) @consumer_work(returns=cached_page_ids) @cute.jit @@ -1840,14 +2046,19 @@ def _prefetch_page_offsets_async( lane, warp_rank, route_state, + kv_head, ): """Stage int32 IDs with cp.async and resolve direct inputs in place. Copy and mapping use the same lane ownership, so the per-thread wait suffices before in-place mapping. The existing TS full barrier then publishes all lanes' final locators to the TMA load task. - Materialized locators need no translation. Direct inputs contain only - selected block IDs; their causal tail locators are synthesized here. + Direct inputs contain only selected block IDs; their causal tail + locators are synthesized here. + The uniform FP8 issuer consumes flat token coordinates, decoded here + per staged entry rather than in the issuer. Held windows share these + across K/V; the streamed K/V staging cadence is unchanged. Other + issuers consume IDs. """ entries = self.page_ids_per_stage stage_base = stage_info.stage_idx * Int32(entries) @@ -1909,6 +2120,14 @@ def _prefetch_page_offsets_async( elif slot > last: # cp.async zero-fill is not the invalid-locator sentinel. self._smem_page_offsets[stage_base + offset] = Int32(-1) + if cutlass.const_expr( + self.cfg.use_fp8_qkv and self.cfg.uses_2d_flat_kv_tma + ): + locator = Int32(self._smem_page_offsets[stage_base + offset]) + token_offset, _ = _decode_native_page_locator( + self.cfg, locator, kv_head + ) + self._smem_page_offsets[stage_base + offset] = token_offset @cute.jit def _producer_load_page_offsets( @@ -2034,6 +2253,12 @@ def _producer_load_page_offsets( + membership_word_idx * Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD) ) + if cutlass.const_expr(cfg.query_major_memberships): + # Every word in a query tile refers to the same 32 pages. + # The private producer initializes all eight tail words. + first_logical_page_idx = ( + raw_tile_base + membership_word_idx // Int32(8) + ) * Int32(32) if first_logical_page_idx <= page_idx_ub: membership_word = Uint32( self.q_token_kv_block_sparse_page_memberships[ @@ -2087,16 +2312,17 @@ def _producer_load_page_offsets( # The producer owns only the live byte prefix. Mask the # final word here so padding bytes remain unobservable # even when the caller reuses an uninitialized workspace. - live_memberships = ( - page_idx_ub - first_logical_page_idx + Int32(1) - ) - if live_memberships < Int32( - Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD - ): - live_bits = live_memberships * Int32( - Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + if cutlass.const_expr(not cfg.query_major_memberships): + live_memberships = ( + page_idx_ub - first_logical_page_idx + Int32(1) ) - membership_word &= (Uint32(1) << live_bits) - Uint32(1) + if live_memberships < Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ): + live_bits = live_memberships * Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + ) + membership_word &= (Uint32(1) << live_bits) - Uint32(1) self._smem_q_token_kv_block_sparse_memberships[ membership_word_idx ] = membership_word @@ -2117,6 +2343,7 @@ def _producer_load_page_offsets( lane_idx, page_warp_rank, direct_route, + logical_h_idx, ) elif cutlass.const_expr(self.holds_encoded_locator_window): # Stripe whole KV tiles: with two producers, warps own even/odd diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py index e75003484df3..d08b770708f0 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py @@ -505,7 +505,9 @@ def _vp_mma( ) else: v_operand = v_operand + Int32( - (cfg.head_dim_kv_stage * 2) if cfg.use_fp8_pv else 128 + (min(cfg.head_dim_kv_stage, 128) * 2) + if cfg.use_fp8_pv + else 128 ) if cutlass.const_expr(not cfg.uses_tmem_p): if cutlass.const_expr( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py index 715a2ec91bcc..73982aa0a912 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py @@ -901,6 +901,12 @@ def _qk_mma( elif cutlass.const_expr(cfg.qk_dtype_bytes == 1): k_desc_offset = ki * Int32(2) q_desc_offset = ki * Int32(2) + if cutlass.const_expr(cfg.head_dim_kv_stage > 128): + # A full FP8 head stage consists of two SW128 planes, + # not a row-major D256 tile. Cross planes at K128. + chunk_idx = (ki * Int32(_mma_k_step_qk(cfg))) // Int32(128) + k_desc_offset += chunk_idx * Int32(8 * cfg.tile_size_kv - 8) + q_desc_offset += chunk_idx * Int32(8 * cfg.tile_size_q - 8) iter_k_desc = k_desc + k_desc_offset else: chunk_idx = (ki * Int32(_mma_k_step_qk(cfg))) // Int32(64) @@ -1072,24 +1078,42 @@ def _q_token_kv_block_sparse_keeps_kv128_membership_word( lane_idx, cfg.softmax_score_fragment_regs, ) + if cutlass.const_expr(cfg.query_major_memberships): + return self.page_offsets_ref.query_membership_keep_word( + stage_info, local_tile_idx, q_token_idx + ) >> Uint32(col_base // Int32(4)) keep_word = Uint32(0) page_span = min(cfg.num_tokens_per_page, cfg.num_s_regs_per_thread) pages_per_lane = cfg.num_s_regs_per_thread // page_span + if cutlass.const_expr(pages_per_lane >= 16): + for vector_idx in cutlass.range_constexpr(pages_per_lane // 16): + bits = self.page_offsets_ref.q_token_kv_block_sparse_page_keep_word16( + stage_info, + local_tile_idx, + col_base // Int32(cfg.num_tokens_per_page) + Int32(vector_idx * 16), + q_token_idx, + ) + keep_word |= bits << Uint32(vector_idx * 16) + return keep_word for page_vector_idx in cutlass.range_constexpr(pages_per_lane // 4): memberships = ( - self.page_offsets_ref.q_token_kv_block_sparse_page_memberships4( + self.page_offsets_ref.q_token_kv_block_sparse_page_membership_word4( stage_info, local_tile_idx, col_base // Int32(cfg.num_tokens_per_page) + Int32(page_vector_idx * 4), ) ) - for vector_elem_idx in cutlass.range_constexpr(4): - local_page_idx = page_vector_idx * 4 + vector_elem_idx - page_is_member = Uint32( - (memberships[vector_elem_idx] & membership_bit) != Uint32(0) - ) - keep_word = keep_word | (page_is_member << Uint32(local_page_idx)) + # Select query q's bit from each byte (positions 0/8/16/24), then + # pack them into bits 0..3 with a carry-free uint32 multiply, not + # a first-set-bit search. Zero q >= 8 so the bounded shift cannot + # alias another query's bit. + byte_bits = (memberships >> (q_token_idx & Int32(7))) & Uint32(0x01010101) + packed_bits = (byte_bits * Uint32(0x01020408)) >> Uint32(24) + packed_bits = cutlass.select_( + Uint32(q_token_idx) < Uint32(8), packed_bits, Uint32(0) + ) + keep_word = keep_word | (packed_bits << Uint32(page_vector_idx * 4)) if cutlass.const_expr(pages_per_lane < 4): for local_page_idx in cutlass.range_constexpr(pages_per_lane): membership = ( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py index 18f2c272cf5c..a4a87200ca5d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,6 +21,8 @@ import cutlass import cutlass.cute as cute from cutlass import Int32, Int64, Uint8 +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import BaseDSL from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.enums import WorkAttr from cutlass.experimental.task_scheduling.memory import ( @@ -34,6 +37,7 @@ producer_work, ) +from flashinfer.cute_dsl.availability import is_rubin_cute_dsl_available from ..fmha_decode_config import FmhaDecodeConfig from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, @@ -46,14 +50,27 @@ from .smem_resources import SmemKvResource +def _has_packed_e4m3_e2m1_mul() -> bool: + """Whether the compile target has the packed E4M3 x E2M1 multiply. + + PTX 9.4 (CUDA 13.4) adds it for SM100 and SM103; ptxas rejects it for + SM107, which takes the f16x2 path instead. + """ + if not cutlass.target_version(min_version="13.4"): + return False + if not is_rubin_cute_dsl_available(): + return True + arch = BaseDSL._get_dsl().get_arch_enum() + return not (Arch.sm_107 <= arch <= Arch.sm_107f) + + @cute.jit def _mul_e2m1x4_e4m3x4(packed_fp4: Int32, packed_sf: Int32) -> Int32: """Apply four E4M3 scale factors to four unpacked E2M1 values. - PTX 9.4 (CUDA 13.4) supports the byte-padded E2M1 operand directly. Each FP4 value occupies the low nibble of one byte in the 32-bit input. """ - if cutlass.const_expr(cutlass.target_version(min_version="13.4")): + if cutlass.const_expr(_has_packed_e4m3_e2m1_mul()): return cute.arch.inline_ptx( "mul.e4m3x4.e2m1p4x4.e4m3x4.satfinite {$w0}, {$r0}, {$r1};", write_only_types=[Int32], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py index 7cea30d06cb4..7d2393db28f0 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py @@ -2620,7 +2620,7 @@ def create_mma_task_one_inst_qkv( work_queue: WorkQueue | None, cfg: FmhaDecodeConfig, *, - tmem_stats_done: MemoryResource, + tmem_stats_done: MemoryResource | None, domain: int | cutlass.Int32, warp_idx: int | None = None, num_warps: int | None = None, @@ -2636,14 +2636,15 @@ def mma_schedule_body( tmem_s: MemoryResource, smem_p: MemoryResource, tmem_o: MemoryResource, - tmem_stats_done: MemoryResource, + tmem_stats_done: MemoryResource | None, q_desc: Any, ) -> None: """Schedule single-instance QK and PV waves across HEAD/LOOP/TAIL.""" def qk_mma(q_desc, qk_mma_label: str, section: FmhaStage) -> None: """Issue one scheduled single-instance QK wave.""" - tmem_stats_done.acquire() + if tmem_stats_done is not None: + tmem_stats_done.acquire() tmem_s.acquire() _issue_score_seed(tmem_s, section, cfg) for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): @@ -2662,7 +2663,8 @@ def qk_mma(q_desc, qk_mma_label: str, section: FmhaStage) -> None: ) smem_k.release() tmem_s.commit() - tmem_stats_done.commit() + if tmem_stats_done is not None: + tmem_stats_done.commit() def pv_mma(vp_mma_label: str, section: FmhaStage) -> None: """Issue one scheduled single-instance PV wave.""" @@ -2710,7 +2712,7 @@ def mma_schedule_prelude( smem_v.init_descriptor_state() smem_p.init_descriptor_state() - @schedule + @_schedule_with_optional_resources def mma_schedule( smem_q: MemoryResource, smem_k: MemoryResource, @@ -2718,8 +2720,8 @@ def mma_schedule( tmem_s: MemoryResource, smem_p: MemoryResource, tmem_o: MemoryResource, - tmem_stats_done: MemoryResource, - work_queue: WorkQueue | None = None, + tmem_stats_done: MemoryResource | None, + work_queue: WorkQueue | None, ) -> None: """Wrap one-inst MMA work in packed persistent skip handling.""" _decode_work_tile_schedule_with_invariant_bridge( @@ -2740,34 +2742,16 @@ def mma_schedule( ), ) - schedule_result = ( - mma_schedule( - smem_q, - smem_k, - smem_v, - tmem_s, - smem_p, - tmem_o, - tmem_stats_done, - ) - if work_queue is None - else mma_schedule( - smem_q, - smem_k, - smem_v, - tmem_s, - smem_p, - tmem_o, - tmem_stats_done, - work_queue, - ) + schedule_result = mma_schedule( + smem_q, smem_k, smem_v, tmem_s, smem_p, tmem_o, tmem_stats_done, work_queue ) src = [smem_q, smem_k, smem_v, smem_p] if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[tmem_s, tmem_o, tmem_stats_done], + dst_resources=[tmem_s, tmem_o] + + ([tmem_stats_done] if tmem_stats_done is not None else []), cfg=cfg, warp_idx=cfg.mma_warp_idx if warp_idx is None else warp_idx, num_warps=cfg.mma_num_warps if num_warps is None else num_warps, @@ -4149,7 +4133,7 @@ def padding_body() -> None: warp_idx=warp_idx, num_warps=num_warps, schedule=captured_schedule, - num_registers=cfg.mma_load_task_num_registers, + num_registers=cfg.padding_task_num_registers(warp_idx, num_warps), name="PaddingTask", **kw, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py index bfd4d2229592..f28769a10c8d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py @@ -23,9 +23,11 @@ * S2-S4 use one 512-thread CTA per 8 KiB output slice and fold the exact split count in registers. -* S5 and larger use 128-thread CTAs over 2 KiB slices. Each cluster rank folds - its split slots locally, publishes ``(LSE, O)`` through SMEM, and rank zero - merges the distributed states. G16 uses a 4x4 two-level merge. +* S5-S16 use one 512-thread CTA per 2 KiB slice, with four adjacent lanes + folding interleaved split slots and merging through warp shuffles. +* S17 and larger use 128-thread CTAs over 2 KiB slices. Each cluster rank + publishes its local ``(LSE, O)`` through SMEM; rank zero merges the states. + Sixteen-CTA clusters use a 4x4 two-level merge. A 512-thread exact-split kernel remains as the serial reference schedule. Both schedules merge the optional attention-sink denominator, pack the requested @@ -459,9 +461,12 @@ def decode_gen_parallel_separate_reduction_kernel( """Reduce split partials with a compact or clustered constexpr schedule. S2-S4 use one 512-thread CTA per 8 KiB output slice and reduce the exact - split count directly. Larger schedules use 128-thread CTAs over 2 KiB - slices. Each cluster rank owns 2, 4, or 8 split slots; G1 stores directly, - G16 uses a two-level 4x4 merge, and G2/G4/G8 finalize through rank zero. + split count directly. Larger schedules cover 2 KiB slices with 128 output + fragments per CTA. A single CTA (up to 16 slots) spreads each fragment's slots + over ``PARALLEL_REDUCTION_SLOT_LANES`` adjacent lanes, merges them with + warp shuffles, and stores directly; clustered profiles keep one lane per + fragment, each cluster rank owns 2, 4, or 8 split slots, G16 uses a + two-level 4x4 merge, and G2/G4/G8 finalize through rank zero. """ if cutlass.const_expr(cfg.use_compact_parallel_reduction): @@ -494,9 +499,15 @@ def decode_gen_parallel_separate_reduction_kernel( bytes_per_output_row = Int32(cfg.headdim * PARTIAL_O_ELEMENT_BYTES) slice_idx = block_idx_x // Int32(cfg.parallel_reduction_cluster_size) + # ``slot_lanes`` adjacent lanes share one 16-byte output fragment and each + # own an interleaved subset of its split slots (lane ``l`` folds slots + # ``l, l + slot_lanes, ...`` so a short runtime prefix stays balanced). + slot_lanes = cfg.parallel_reduction_slot_lanes + fragment_idx = thread_idx // Int32(slot_lanes) + lane_slot = thread_idx % Int32(slot_lanes) reduce_base_offset = slice_idx * Int32( PARALLEL_REDUCTION_BYTES_PER_SLICE - ) + thread_idx * Int32(REDUCTION_BYTES_PER_THREAD) + ) + fragment_idx * Int32(REDUCTION_BYTES_PER_THREAD) reduce_row_idx = reduce_base_offset // bytes_per_output_row reduce_col_idx = (reduce_base_offset % bytes_per_output_row) // Int32( PARTIAL_O_ELEMENT_BYTES @@ -531,10 +542,11 @@ def decode_gen_parallel_separate_reduction_kernel( output_vals[elem_idx] = Float32(0.0) global_lse = Float32(-Float32.inf) - # Batch up to four independent GMEM loads before folding them. The final - # batch width is compile-time. Padded split slots stay neutral and never - # form GMEM pointers, including non-power-of-two split counts. + # Issue every owned slot's independent GMEM loads before folding them. The + # final batch width is compile-time. Padded split slots stay neutral and + # never form GMEM pointers, including non-power-of-two split counts. local_splits = cfg.parallel_reduction_splits_per_cta + lane_splits = local_splits // slot_lanes partial_lse = cutlass.Array( Float32, PARALLEL_REDUCTION_LOAD_BATCH, space=cutlass.AddressSpace.rmem ) @@ -575,12 +587,18 @@ def decode_gen_parallel_separate_reduction_kernel( ) for split_base_i in cutlass.range_constexpr( - 0, local_splits, PARALLEL_REDUCTION_LOAD_BATCH + 0, lane_splits, PARALLEL_REDUCTION_LOAD_BATCH ): - batch_width = min(PARALLEL_REDUCTION_LOAD_BATCH, local_splits - split_base_i) + batch_width = min(PARALLEL_REDUCTION_LOAD_BATCH, lane_splits - split_base_i) split_base = Int32(split_base_i) + # Prefetch only published slots; inactive slots and padded Q rows + # are in bounds but may contain uninitialized caller-owned scratch. for jj in cutlass.range_constexpr(batch_width): - split_idx = cluster_rank * Int32(local_splits) + split_base + Int32(jj) + split_idx = ( + cluster_rank * Int32(local_splits) + + (split_base + Int32(jj)) * Int32(slot_lanes) + + lane_slot + ) valid_split_idx = cutlass.const_expr( cfg.parallel_reduction_padded_splits == cfg.max_splits_kv ) or split_idx < Int32(cfg.max_splits_kv) @@ -622,7 +640,11 @@ def decode_gen_parallel_separate_reduction_kernel( partial_regs[regs_base + reg_idx] = loaded_partial_regs[reg_idx] for jj in cutlass.range_constexpr(batch_width): - split_idx = cluster_rank * Int32(local_splits) + split_base + Int32(jj) + split_idx = ( + cluster_rank * Int32(local_splits) + + (split_base + Int32(jj)) * Int32(slot_lanes) + + lane_slot + ) valid_split_idx = cutlass.const_expr( cfg.parallel_reduction_padded_splits == cfg.max_splits_kv ) or split_idx < Int32(cfg.max_splits_kv) @@ -649,11 +671,50 @@ def decode_gen_parallel_separate_reduction_kernel( ) global_lse = new_lse - # G1 stores its register accumulator directly and compiles out SMEM - # publication, mapa, and cluster barriers. The branch is CTA-uniform and - # resolved at compile time. + # G1 merges the lanes sharing one fragment with a warp butterfly, then the + # first lane stores the register accumulator directly and compiles out + # SMEM publication, mapa, and cluster barriers. The branch is CTA-uniform + # and resolved at compile time. if cutlass.const_expr(cfg.parallel_reduction_cluster_size == 1): - if valid_reduce_row: + # Shuffles execute warp-uniformly outside every row/slot predicate; + # lanes of invalid fragments exchange neutral or unused state. Names + # stay distinct from the clustered finalize below so the DSL sees one + # stable type per variable across the compile-time branches. + if cutlass.const_expr(slot_lanes > 1): + for lane_level in cutlass.range_constexpr((slot_lanes - 1).bit_length()): + lane_xor = 1 << lane_level + lane_peer_lse = Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=global_lse, + offset=lane_xor, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + lane_peer_vals = cutlass.Array( + Float32, OUTPUT_VALUES_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + lane_peer_vals[elem_idx] = Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=output_vals[elem_idx], + offset=lane_xor, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + lane_merged_lse, lane_own_weight, lane_peer_weight = merge_log2_lse( + global_lse, lane_peer_lse + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + output_vals[elem_idx] = ( + output_vals[elem_idx] * lane_own_weight + + lane_peer_vals[elem_idx] * lane_peer_weight + ) + global_lse = lane_merged_lse + if valid_reduce_row and lane_slot == Int32(0): _store_parallel_reduction_output( output_vals, global_lse, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py index aa68b9efa0f1..1e47d5c5b237 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py @@ -34,8 +34,10 @@ WARPGROUP_WARPS = 4 WARPGROUP_THREADS = WARPGROUP_WARPS * WARP_LANES -# Kernel-level TMEM lifecycle synchronization. Barrier ID 15 is reserved for -# alloc/dealloc phases in the 1CTA path; the 2CTA path uses config-owned IDs. +# Kernel-level 1CTA TMEM synchronization: 14 publishes allocation to compute +# warps; 15 handles full-CTA initialization/deallocation. The 2CTA path owns +# its barrier IDs in its configuration. +TMEM_READY_BARRIER_ID = 14 TMEM_LIFECYCLE_BARRIER_ID = 15 TMEM_DEALLOC_MBAR_THREADS = WARP_LANES @@ -69,6 +71,10 @@ TCGEN05_16X32BX2_FP8_P_STRIDE = 16 TCGEN05_SECOND_PANEL_ADDR_OFFSET = 16 << 16 +# Softmax uses base-two exponentials internally; public sparse LSE uses ln. +LN_2 = 0.6931471805599453 +LOG2_E = 1.4426950408889634 + # Softmax scratch stores the max state first and the sum state in a second # fixed-size panel. The offset is in Uint32 scratch words. SOFTMAX_SCRATCH_SUM_WORD_OFFSET = 384 @@ -87,6 +93,16 @@ PAGE_OFFSET_BYTES = 4 CP_ASYNC_CACHE_CA = "ca" +# gather4 always copies four selected KV rows. This instruction width is +# independent of the number of issuing warps or the KV tile size. +TMA_GATHER_ROWS = 4 + +# Sparse 2CTA offset caches derive their sizes from cfg token geometry: +# K owns 128/2=64 rows per CTA; V needs all 128 rows. Lane-distributed +# caches therefore hold 64/32=2 K and 128/32=4 V words per lane. A uniform +# cache instead holds the issuing warp's share (64/W K and 128/W V words +# per lane). Pending V stays lane-distributed until it becomes current. + # Dense MLA kernels specialize paged-KV addressing for this explicit ABI set. # Each listed size exactly partitions a 128-token KV tile; other page sizes are # outside the current kernel ABI. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py new file mode 100644 index 000000000000..d36f2ad068f2 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py @@ -0,0 +1,452 @@ +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Blackwell four-row TMA gather with explicit five-coordinate PTX operands. + +The CUTLASS DSL 4.7 experimental wrapper validates gather coordinates as a pair; +it cannot express four independent row coordinates through that interface. +Keep the native instruction isolated here instead of patching the dependency. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +from .constants import TMA_GATHER_ROWS +from .sparse_indices import SparseKVIndices, SparseSourceTile + + +@dsl_user_op +def _address(pointer, dtype, *, loc=None, ip=None): + if hasattr(pointer, "data_ptr"): + pointer = pointer.data_ptr(loc=loc, ip=ip) + if hasattr(pointer, "toint"): + pointer = pointer.toint(loc=loc, ip=ip) + return dtype(pointer) + + +@cute.jit +def select_gather_map(primary, extra, use_extra): + """Select numeric addresses: Pointer-object rebinding is not an SSA select.""" + primary_addr = _address(primary, Int64) + extra_addr = _address(extra, Int64) + return extra_addr if use_extra else primary_addr + + +@cute.jit +def load_sparse_rows(routes, request, offset, *, fragments: cutlass.Constexpr[int]): + """Coalesced row window, distributed as one coordinate per lane/fragment.""" + lane = cute.arch.thread_idx()[0] & Int32(31) + rows = cutlass.Array(Int32, fragments, space=cutlass.AddressSpace.rmem) + for fragment in cutlass.range_constexpr(fragments): + index = offset + lane + Int32(fragment * 32) + value = Int32(0x7FFFFFFF) + if cutlass.const_expr(isinstance(routes, SparseSourceTile)): + value = routes.load_row(index) + elif index < routes.shape[0]: + value = Int32(routes[index, request]) + rows[fragment] = value + return rows + + +@cute.jit +def load_sparse_quad(routes, request, offset): + """Keep one contiguous row quadruple in its issuing lane's registers.""" + if cutlass.const_expr(isinstance(routes, SparseKVIndices)): + source = routes.select_source_tile(offset, request) + return source.load_quad(offset) + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + rows[j] = Int32(0x7FFFFFFF) + if offset + Int32(3) < routes.shape[0]: + if cutlass.const_expr(routes.stride[0] == 1): + values = cutlass.Pointer( + routes.iterator + offset + request * routes.stride[1], dtype=Int32 + ).load(count=4, alignment=4) + for j in cutlass.range_constexpr(4): + rows[j] = values[j] + else: + for j in cutlass.range_constexpr(4): + rows[j] = Int32(routes[offset + j, request]) + else: + for j in cutlass.range_constexpr(4): + if offset + Int32(j) < routes.shape[0]: + rows[j] = Int32(routes[offset + j, request]) + return rows + + +@cute.jit +def decode_sparse_quad(primary, extra, raw): + """Select a source and map masked row sentinels to TMA zero-fill rows.""" + rows = cutlass.Array(Int32, TMA_GATHER_ROWS, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(TMA_GATHER_ROWS): + # Native sparse pools have fewer than INT32_MAX storage rows. The + # reserved sentinel is already an out-of-bounds TMA coordinate; + # retaining it gives zero fill without a compare/select per row. + rows[j] = raw[j] & Int32(0x7FFFFFFF) + # Internally generated invalid coordinates retain the source bit: the + # first coordinate selects the shared descriptor even when it is invalid. + return select_gather_map(primary, extra, raw[0] < 0), rows + + +@cute.jit +def gather4_quad( + dst, primary, extra, column, raw, barrier, *, cta_group: cutlass.Constexpr[int] = 1 +): + """Issue a quad from the owning lane, without election or warp shuffles.""" + descriptor, rows = decode_sparse_quad(primary, extra, raw) + gather4( + dst, + descriptor, + column, + rows[0], + rows[1], + rows[2], + rows[3], + barrier, + cta_group=cta_group, + ) + + +@cute.jit +def gather4_cached( + dst, + primary, + extra, + column, + cached_rows, + token_offset, + barrier, + *, + cache_base: cutlass.Constexpr[int] = 0, + cta_group: cutlass.Constexpr[int] = 1, +): + """Assemble a row quadruple using only warp-local register shuffles.""" + raw0 = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_rows[cache_base + token_offset // 32], + offset=token_offset % 32, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + token = token_offset + Int32(j) + raw = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_rows[cache_base + token // 32], + offset=token % 32, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + row = raw & Int32(0x7FFFFFFF) + rows[j] = row if row != Int32(0x7FFFFFFF) else Int32(-1) + # The same source-tag invariant applies to cached invalid anchors. + descriptor = select_gather_map(primary, extra, raw0 < 0) + if prims.elect_sync(): + gather4( + dst, + descriptor, + column, + rows[0], + rows[1], + rows[2], + rows[3], + barrier, + cta_group=cta_group, + ) + + +@cute.jit +def invalid_sparse_token( + token, length, routes, request, *, sparse: cutlass.Constexpr[bool] +): + """Mask prefix tails and tagged invalid routes before softmax. + + Bit 31 selects the source; bits 0..30 hold its TMA row coordinate. + All-one low bits denote an invalid row for either source. + """ + if cutlass.const_expr( + sparse and isinstance(routes, (SparseSourceTile, CachedSparseMask)) + ): + return (token >= length) | routes.is_invalid(token, request) + invalid = token >= length + if cutlass.const_expr(sparse): + if token < length: + invalid = (Int32(routes[token, request]) & Int32(0x7FFFFFFF)) == Int32( + 0x7FFFFFFF + ) + return invalid + + +@cute.jit +def gather4( + dst, + tensor_map, + column, + row0, + row1, + row2, + row3, + barrier, + *, + cta_group: cutlass.Constexpr[int] = 1, +): + """Issue one four-row copy from the calling thread. + + Multiple lanes may issue disjoint copies. The caller accounts for all + copies in the shared completion barrier, including masked/OOB rows. + + The descriptor is 2D with a one-row box and no interleave. Destination + alignment/swizzle and coordinate bounds follow its tensor-map contract. + Group 2 delivers completion to the even CTA of the issuing CTA pair. + """ + dst_addr = _address(dst, Int32) + bar_addr = _address(barrier, Int32) + if cutlass.const_expr(cta_group == 2): + bar_addr = bar_addr & Int32(-16777217) # Clear peer-CTA address bit 24. + instruction = ( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4." + "mbarrier::complete_tx::bytes.cta_group::2 " + ) + else: + assert cta_group == 1 + instruction = ( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4." + "mbarrier::complete_tx::bytes.cta_group::1 " + ) + cute.arch.inline_ptx( + instruction + "[{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], [{$r7}];", + read_only_args=[ + dst_addr, + _address(tensor_map, Int64), + Int32(column), + Int32(row0), + Int32(row1), + Int32(row2), + Int32(row3), + bar_addr, + ], + ) + + +class CachedSparseMask: + """One/two per-warp bitsets reused across all score registers and heads.""" + + def __init__(self, values, swaps): + self.values = values + self.swaps = swaps + self.low, self.high = values + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return CachedSparseMask( + cutlass.new_from_mlir_values(self.values, values), self.swaps + ) + + @cute.jit + def is_invalid(self, token, request): + del request + word = self.low if self.swaps or (token & Int32(32)) == Int32(0) else self.high + return ((word >> (token & Int32(31))) & Int32(1)) != Int32(0) + + +@cute.jit +def cache_sparse_mask( + routes, + tile_offset, + length, + request, + *, + swaps: cutlass.Constexpr[bool], + column_half=None, + tile_size_kv: cutlass.Constexpr[int] = 128, +): + lane = cute.arch.thread_idx()[0] % 32 + warp = cute.arch.thread_idx()[0] % 128 // 32 + if cutlass.const_expr(isinstance(routes, SparseSourceTile) and routes.valid_prefix): + low, high = Int32(0), Int32(0) + if not routes.tile_is_full(tile_offset, length): + if cutlass.const_expr(swaps): + low = routes.prefix_mask32(tile_offset + warp * Int32(32), length) + high = low + else: + if cutlass.const_expr(column_half is None): + column_half = lane >> Int32(4) + start = tile_offset + if cutlass.const_expr(tile_size_kv == 128): + start += column_half * Int32(64) + low = routes.prefix_mask32(start, length) + high = routes.prefix_mask32(start + Int32(32), length) + return CachedSparseMask((low, high), swaps) + if cutlass.const_expr(swaps): + token = tile_offset + warp * Int32(32) + lane + invalid = invalid_sparse_token(token, length, routes, request, sparse=True) + bits = Int32(cute.arch.vote_ballot_sync(invalid)) + return CachedSparseMask((bits, bits), True) + else: + masks = cutlass.Array( + Int32, tile_size_kv // 32, space=cutlass.AddressSpace.rmem + ) + for part in cutlass.range_constexpr(tile_size_kv // 32): + token = tile_offset + Int32(part * 32) + lane + invalid = invalid_sparse_token(token, length, routes, request, sparse=True) + masks[part] = Int32(cute.arch.vote_ballot_sync(invalid)) + if cutlass.const_expr(column_half is None): + column_half = lane >> Int32(4) + if cutlass.const_expr(tile_size_kv == 64): + low, high = masks[0], masks[1] + else: + low = masks[0] if column_half == Int32(0) else masks[2] + high = masks[1] if column_half == Int32(0) else masks[3] + return CachedSparseMask((low, high), False) + + +@dsl_user_op +def _gather4_slices( + dst, + tensor_map, + column, + rows, + barrier, + *, + num_slices, + dst_stride, + col_stride, + loc=None, + ip=None, +): + # All callers broadcast the quad before entering this collective. Keep + # feature slices in one asm block so election/uniformization is shared. + instruction = "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::2 " + asm = [ + "{ .reg .b32 slice_dst, slice_col; .reg .pred issue; elect.sync _|issue, 0xffffffff;", + "mov.b32 slice_dst, {$r0};", + "mov.b32 slice_col, {$r2};", + ] + for i in range(num_slices): + asm.append( + "@issue " + + instruction + + "[slice_dst], [{$r1}, {slice_col, {$r3}, {$r4}, {$r5}, {$r6}}], [{$r7}];" + ) + if i + 1 < num_slices: + asm.append(f"add.s32 slice_dst, slice_dst, {dst_stride};") + asm.append(f"add.s32 slice_col, slice_col, {col_stride};") + asm.append("}") + cute.arch.inline_ptx( + "\n".join(asm), + read_only_args=[ + _address(dst, Int32), + _address(tensor_map, Int64), + Int32(column), + rows[0], + rows[1], + rows[2], + rows[3], + _address(barrier, Int32) & Int32(-16777217), + ], + loc=loc, + ip=ip, + ) + + +@cute.jit +def gather4_uniform_quad_slices( + dst, + primary, + extra, + column, + raw, + barrier, + *, + num_slices: cutlass.Constexpr[int], + dst_stride: cutlass.Constexpr[int], + col_stride: cutlass.Constexpr[int], +): + """Issue 2CTA feature slices from a warp-uniform cached row quadruple. + + Every lane must participate with identical arguments. One elected lane + issues all copies; dst_stride is bytes and col_stride is source elements. + The caller accounts for all copied bytes in the completion barrier. + """ + descriptor, rows = decode_sparse_quad(primary, extra, raw) + _gather4_slices( + dst, + descriptor, + column, + rows, + barrier, + num_slices=num_slices, + dst_stride=dst_stride, + col_stride=col_stride, + ) + + +@dsl_user_op +def gather4_cached_warp( + dst, + tensor_map, + column, + cached_rows, + barrier, + *, + cache_start, + num_quads, + num_slices, + quad_stride, + slice_stride, + col_stride, + loc=None, + ip=None, +): + """Issue one warp's cached quads with one election for every feature slice. + + All lanes participate. Rows are already mapped, cleaned and warp-uniform; + the caller retains them across the head-dimension stages of one K tile. + Strides for destinations are bytes, and the column stride is elements. + """ + assert num_quads <= 8 + args = [ + _address(dst, Int32), + _address(tensor_map, Int64), + Int32(column), + _address(barrier, Int32), + ] + args += [cached_rows[cache_start + i] for i in range(num_quads * 4)] + instruction = "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::1 " + asm = [ + "{ .reg .b32 copy_dst, copy_col; .reg .pred issue; elect.sync _|issue, 0xffffffff;" + ] + for part in range(num_slices): + asm.append(f"add.s32 copy_col, {{$r2}}, {part * col_stride};") + for quad in range(num_quads): + asm.append( + f"add.s32 copy_dst, {{$r0}}, {part * slice_stride + quad * quad_stride};" + ) + rows = ", ".join("{$r" + str(4 + quad * 4 + j) + "}" for j in range(4)) + asm.append( + "@issue " + + instruction + + "[copy_dst], [{$r1}, {copy_col, " + + rows + + "}], [{$r3}];" + ) + asm.append("}") + cute.arch.inline_ptx("\n".join(asm), read_only_args=args, loc=loc, ip=ip) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py index fe45148142c3..a9d28a54ecef 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py @@ -54,7 +54,7 @@ def num_softmax_scale_groups(cfg: MlaConfig) -> int: """Return the number of independent softmax scale groups per thread.""" if cfg.kernel_variant == "keeps_mma_ab": return 1 - return max(cfg.tile_size_q // 4, 4) + return max(cfg.tile_size_q // 4, 2) def num_s_regs_per_thread(cfg: MlaConfig) -> int: @@ -73,6 +73,8 @@ def softmax_scratch_words(cfg: MlaConfig) -> int: def num_packed_p_regs(cfg: MlaConfig) -> int: """Return the number of packed P registers per softmax thread.""" + if cfg.kernel_variant == "keeps_mma_ab": + return cfg.tile_size_kv // (8 if cfg.is_fp8_qkv() else 4) q_repeats = max(cfg.tile_size_q // 8, 1) return (2 if cfg.is_fp8_qkv() else 4) * q_repeats diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py index e07a06325518..b2d9933ec306 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py @@ -21,9 +21,7 @@ from cutlass import Float32, Int32, Uint32 from cutlass.experimental import primitives as prims -# Softmax converts exp2 inputs with log2(e) and initializes masked scores to -# negative Float32 max. -LOG2_E = 1.4426950408889634074 +# Masked scores start at negative Float32 max. NEG_FLT_MAX = -3.4028235e38 # E4M3FN finite maximum used when clamping FP8 output conversion. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py index 86f1ca7cee6b..021b9ebbd79f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py @@ -223,6 +223,12 @@ def fp8_log2_quant_scale(): return Float32(8.8073549) +@cute.jit +def fp8_quant_scale(): + """Return the E4M3 probability scale used by sink normalization.""" + return Float32(448.0) + + @cute.jit def fp8_quant_scale_rcp(): """Return the reciprocal of the shared E4M3 probability scale.""" @@ -334,14 +340,19 @@ def store_transposed_smem8b_x4( @cute.jit -def tcgen05_store_p_16x32bx2_x16(tmem_addr, regs_p, start_idx: cutlass.Constexpr[int]): +def tcgen05_store_p_16x32bx2_x16( + tmem_addr, + regs_p, + start_idx: cutlass.Constexpr[int], + stride: cutlass.Constexpr[int] = TCGEN05_16X32BX2_BF16_P_STRIDE, +): """Store 16 packed BF16 P registers to a split TMEM-P tile.""" # The stride immediate is the BF16 lane-to-column spacing for the split # 16x32bx2 P layout; keep it named because the FP8 path uses a different # immediate with the same instruction shape. inline_ptx( "tcgen05.st.sync.aligned.16x32bx2.x16.b32 " - f"[{{$r0}}], {TCGEN05_16X32BX2_BF16_P_STRIDE}, " + f"[{{$r0}}], {stride}, " "{ {$r1}, {$r2}, {$r3}, {$r4}, {$r5}, {$r6}, {$r7}, {$r8}, " "{$r9}, {$r10}, {$r11}, {$r12}, {$r13}, {$r14}, {$r15}, {$r16} };", read_only_args=[ diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/output_quant.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/output_quant.py new file mode 100644 index 000000000000..e620e90dfb2f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/output_quant.py @@ -0,0 +1,251 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Final sparse MLA output: conjugate RoPE and block E4M3 quantization.""" + +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T +from cutlass.experimental import primitives as prims + +from .math import mul_packed_f32x2, fma_packed_f32x2 +from .ops import pack_float4_to_fp8_e4m3 + + +@cutlass.dsl_user_op +def _max_abs3(x, y, z, *, loc=None, ip=None): + """SM100 max-of-magnitudes with native floating-point operand constraints.""" + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(v).ir_value(loc=loc, ip=ip) for v in (x, y, z)], + "max.abs.ftz.f32 $0, $1, $2, $3;", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dataclass(frozen=True) +class MlaOutputQuant: + """Static geometry; runtime args are (flat O, positions, cos/sin, scales). + + O is E4M3 stored group-major, physically [G, T, (H/G) * 512] (T is derived + from its extent); positions are INT32; cos/sin is FP32; scales are a + token-major logical view of column-major FP32 or packed INT32 storage. + Rows are token * H + head and columns are channels within the D512 head. + """ + + block_size: int + ue8m0: bool + num_heads: int + num_groups: int + + @cute.jit + def rotate(self, args, values, row, column): + """Rotate a mutable final-output fragment in place.""" + _, positions, cos_sin, _ = args + token = Int64(row) // self.num_heads + count = cutlass.const_expr(values.shape[0]) + fragment = values + + # The cache contains forward sin; conjugation changes its sign here. + if column + count - 1 >= 448: + position = positions[token] + # Load independent coefficients together before consuming them. + pairs = cutlass.const_expr(min(32, count // 2)) + for i in cutlass.range_constexpr(0, count, pairs * 2): + d = column + i + if d >= 448: + # Callers provide aligned contiguous channel vectors. + base = ( + cos_sin.iterator.raw_ptr() + + Int64(position) * 64 + + (d - 448) // 2 + ) + alignment = cutlass.const_expr(min(16, pairs * 4)) + cos = base.load(count=pairs, alignment=alignment) + sin = (base + 32).load(count=pairs, alignment=alignment) + for j in cutlass.range_constexpr(pairs): + x, y = fragment[i + 2 * j], fragment[i + 2 * j + 1] + sine = mul_packed_f32x2((y, x), (sin[j], sin[j]), ftz=True) + rotated = fma_packed_f32x2( + (x, y), (cos[j], cos[j]), (sine[0], -sine[1]), ftz=True + ) + fragment[i + 2 * j] = rotated[0] + fragment[i + 2 * j + 1] = rotated[1] + + @cute.jit + def store( + self, + args, + values, + row, + column, + lane_stride: cutlass.Constexpr[int] = 1, + rotary: cutlass.Constexpr[bool] = True, + shared_output=None, + ): + """Finalize contiguous FP32 fragments after all attention reductions. + + Blocks fit in one thread or one warp; ``lane_stride=16`` joins the + two lanes owning the halves of a keep-MMA head. Nonrotary D stages + remove the rotation at compile time, limiting epilogue code size. + ``shared_output`` is the optional 2CTA FP8 output stage: two swizzled + 64-row by 128-byte boxes, consumed by its asynchronous TMA-store task. + """ + if cutlass.const_expr(rotary): + # Some epilogues pass an immutable register vector. Two-CTA + # epilogues rotate their mutable accumulator before this store. + count = cutlass.const_expr(values.shape[0]) + fragment = cutlass.Array(Float32, count, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(count): + fragment[i] = values[i] + self.rotate(args, fragment, row, column) + else: + fragment = values + out, _, _, scales = args + token = Int64(row) // self.num_heads + head = Int64(row) % self.num_heads + # Physical [G, T, (H/G) * 512] output; the 512-channel head rows of one + # (group, token) pair are contiguous. Planning bounds T * H below 2**31, + # so the row index fits Int32. + heads_per_group = cutlass.const_expr(self.num_heads // self.num_groups) + total_tokens = Int32(Int64(cute.size(out)) // (self.num_heads * 512)) + head32 = Int32(head) + if cutlass.const_expr(heads_per_group & (heads_per_group - 1) == 0): + shift = cutlass.const_expr(heads_per_group.bit_length() - 1) + group = head32 >> shift + head_in_group = head32 & Int32(heads_per_group - 1) + else: + group = head32 // heads_per_group + head_in_group = head32 % heads_per_group + out_row = Int64( + (group * total_tokens + Int32(token)) * heads_per_group + head_in_group + ) + lane = cute.arch.thread_idx()[0] % 32 + count = cutlass.const_expr(fragment.shape[0]) + per_block = cutlass.const_expr(min(count, self.block_size)) + lanes = cutlass.const_expr(self.block_size // per_block) + # RoPE branches can leave partners at different program counters. + # Name every lane in the quant block instead of sampling activemask. + group_base = lane % lane_stride + lane // (lanes * lane_stride) * ( + lanes * lane_stride + ) + group_bits = cutlass.const_expr( + sum(1 << (i * lane_stride) for i in range(lanes)) + ) + group_mask = Uint32(group_bits) << Uint32(group_base) + for start in cutlass.range_constexpr(0, count, per_block): + # Two independent max3 chains avoid a long scalar amax dependency. + m0, m1 = Float32(1e-4), Float32(1e-4) + for i in cutlass.range_constexpr(0, per_block, 4): + m0 = _max_abs3(m0, fragment[start + i], fragment[start + i + 1]) + m1 = _max_abs3(m1, fragment[start + i + 2], fragment[start + i + 3]) + amax = _max_abs3(m0, m1, Float32(1e-4)) + for step in cutlass.range_constexpr(lanes.bit_length() - 1): + peer = prims.shfl_sync( + thread_mask=group_mask, + val=amax, + offset=(1 << step) * lane_stride, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + amax = cute.math.max(amax, peer) + scale = amax * Float32(1.0 / 448.0) + if cutlass.const_expr(self.ue8m0): + exponent = (scale.bitcast(Uint32) + Uint32(0x007FFFFF)) >> 23 + scale = (exponent << 23).bitcast(Float32) + if cutlass.const_expr(self.ue8m0): + inverse_bits = (Uint32(254) - exponent) << 23 + # UE8M0 scales are powers of two; keep reciprocal(inf) == 0. + inv_scale = ( + Float32(0) if exponent == 255 else inverse_bits.bitcast(Float32) + ) + else: + inv_scale = cute.math.rcp(scale, approx=True) + block = (column + start) // self.block_size + # The validated [G, columns, pad4(T)] backing is contiguous in + # head/block order, so group decomposition is unnecessary. + if lane // lane_stride % lanes == 0: + if cutlass.const_expr(self.ue8m0): + word = head * (512 // self.block_size // 4) + block // 4 + byte = block % 4 + offset = token + word * Int64(scales.stride[2]) + byte_ptr = cutlass.inttoptr( + (scales.iterator.raw_ptr() + offset).toint(Int64) + byte, + mem_space=1, + dtype=cutlass.Uint8, + ) + byte_ptr.store(exponent.to(cutlass.Uint8)) + else: + offset = token + (head * (512 // self.block_size) + block) * Int64( + scales.stride[2] + ) + (scales.iterator.raw_ptr() + offset).store(scale) + + words = cutlass.const_expr(min(8, per_block // 4)) + for i in cutlass.range_constexpr(0, per_block, 4 * words): + v = start + i + packed = cutlass.Array(Int32, words, space=cutlass.AddressSpace.rmem) + for w in cutlass.range_constexpr(words): + j = v + 4 * w + lo = mul_packed_f32x2( + (fragment[j], fragment[j + 1]), (inv_scale, inv_scale), ftz=True + ) + hi = mul_packed_f32x2( + (fragment[j + 2], fragment[j + 3]), + (inv_scale, inv_scale), + ftz=True, + ) + packed[w] = pack_float4_to_fp8_e4m3(lo[0], lo[1], hi[0], hi[1]) + if cutlass.const_expr(shared_output is not None): + for part in cutlass.range_constexpr(0, words, 4): + col = Int32(column + v + part * 4) & Int32(255) + smem_offset = ( + (col // 128) * 8192 + (Int32(head) & 63) * 128 + (col & 127) + ) + # TMA 128B swizzle XORs byte bits [4:7] with [7:10]. + # The 1024B-aligned base leaves these swizzle bits unchanged. + swizzled = smem_offset ^ ((smem_offset >> 3) & 0x70) + destination = cutlass.inttoptr( + shared_output + swizzled, mem_space=3, dtype=Int32 + ) + destination.store(packed.load(part, 4), alignment=16) + else: + ptr = cutlass.inttoptr( + (out.iterator.raw_ptr() + out_row * 512 + column + v).toint( + Int64 + ), + mem_space=1, + dtype=Int32, + ) + if cutlass.const_expr(words == 8): + # Native allocations support a 256-bit store. Retain the + # documented 16-byte alignment for caller-owned buffers. + if (ptr.toint(Int64) & Int64(31)) == 0: + ptr.store(packed.load(0, 8), alignment=32) + else: + ptr.store(packed.load(0, 4), alignment=16) + (ptr + 4).store(packed.load(4, 4), alignment=16) + else: + ptr.store(packed.load(0, words), alignment=4 * words) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py index 133c3e15fe48..6e3626b3467c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py @@ -173,6 +173,7 @@ def staged_qk_mma( head_dim_stages=None, consumer_label=None, include_acquire=True, + release_k=True, ): """Consume staged K descriptors and produce one QK MMA score tile.""" if consumer_label is None: @@ -192,7 +193,8 @@ def staged_qk_mma( smem_kv.wait() kv_desc = getattr(smem_kv, consumer_label)(k_subtile_idx=k_subtile_idx) tmem_s.qk_mma(kv_desc=kv_desc, k_subtile_idx=k_subtile_idx) - smem_kv.release() + if release_k: + smem_kv.release() tmem_s.commit() @@ -206,6 +208,7 @@ def staged_pv_mma( consumer_label=None, producer_label=None, is_tail: bool = False, + reuse_k=False, ): """Consume staged V descriptors and produce one PV MMA output tile.""" smem_p.wait() @@ -232,7 +235,8 @@ def staged_pv_mma( producer_label is not None and producer_label.find("tail") >= 0 ) for v_subtile_idx in range(head_dim_stages): - smem_kv.wait() + if not reuse_k: + smem_kv.wait() if producer_label.endswith("_0"): v_desc_0 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) getattr(tmem_o, producer_label)( @@ -316,6 +320,7 @@ def staged_pv_mma_tmem_p( head_dim_stages, consumer_label, producer_label="pv_mma_loop_tmem_p", + reuse_k=False, ): """Consume TMEM P and staged V descriptors to produce one PV output tile.""" tmem_p.wait() @@ -323,7 +328,8 @@ def staged_pv_mma_tmem_p( tmem_o.acquire() is_tail = producer_label == "pv_mma_tail_tmem_p" for v_subtile_idx in range(head_dim_stages): - smem_kv.wait() + if not reuse_k: + smem_kv.wait() v_desc_0 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) getattr(tmem_o, producer_label)( p_stage_idx=p_stage_idx, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py new file mode 100644 index 000000000000..d944aa034bd1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py @@ -0,0 +1,308 @@ +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Select a source index buffer from a tile's logical KV offset. + + row = indices.bind_query_row(query_row) + tile = row.select_source_tile(tile_start) + quad = tile.load_quad(kv_offset) + +The primary source's logical span is rounded to 128 so BK64/BK128 tiles never +cross pools, even when physical rows are random. Caller arrays/live lengths +need no such padding; bounded loads synthesize invalid tails. A 32-entry +boundary would suffice for the grouped TMA descriptors alone, but not for +this whole-tile source selection. These compiler-side objects carry existing +pointers/scalars through CuTe calls; they allocate no merged metadata buffer. +Only SparseKVLengths retains tensor indexing for the shared scheduler. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 + + +class SparseKVIndices: + """Caller-owned index lists; optionally cache one query row's live lengths.""" + + def __init__(self, values, capacity, assume_valid_prefix=False, has_extra=True): + self.values = values + ( + self.primary_indices, + self.extra_indices, + self.primary_lengths, + self.extra_lengths, + self.query_row, + self.primary_length, + self.extra_length, + ) = values + self.capacity = capacity + self.assume_valid_prefix = assume_valid_prefix + self.has_extra = has_extra + self.is_bound = self.query_row is not None + if self.is_bound: + self.primary_span = ( + (self.primary_length + Int32(127)) // Int32(128) * Int32(128) + ) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseKVIndices( + cutlass.new_from_mlir_values(self.values, values), + self.capacity, + self.assume_valid_prefix, + self.has_extra, + ) + + @cute.jit + def bind_query_row(self, query_row, primary_length=None, extra_length=None): + if cutlass.const_expr(self.is_bound): + return self + row = Int64(query_row) + # Once bound, carry scalar lengths instead of the two length tensors. + return SparseKVIndices( + ( + self.primary_indices, + self.extra_indices, + None, + None, + row, + Int32(self.primary_lengths[row]) + if cutlass.const_expr(primary_length is None) + else primary_length, + ( + Int32(self.extra_lengths[row]) + if cutlass.const_expr(extra_length is None) + else extra_length + ) + if cutlass.const_expr(self.has_extra) + else Int32(0), + ), + self.capacity, + self.assume_valid_prefix, + self.has_extra, + ) + + @cute.jit + def select_source_tile( + self, tile_start, query_row=None, *, for_mask: cutlass.Constexpr[bool] = False + ): + row = self.bind_query_row(query_row) + if cutlass.const_expr(not self.has_extra): + # A single-pool plan needs no second length, pointer selection or tag. + if cutlass.const_expr(for_mask and self.assume_valid_prefix): + return SparseSourceTile( + (None, Int32(0), row.primary_length, None), valid_prefix=True + ) + return SparseSourceTile( + ( + row.primary_indices[row.query_row, None], + Int32(0), + row.primary_length, + Int32(0), + ), + valid_prefix=self.assume_valid_prefix, + ) + use_extra = tile_start >= row.primary_span + origin = row.primary_span if use_extra else Int32(0) + length = row.extra_length if use_extra else row.primary_length + if cutlass.const_expr(for_mask and self.assume_valid_prefix): + # Analytic masks need no index pointer or index loads. + return SparseSourceTile((None, origin, length, None), valid_prefix=True) + primary_ptr = row.primary_indices.iterator.toint() + row.query_row * Int64( + row.primary_indices.stride[0] + ) * Int64(4) + extra_ptr = row.extra_indices.iterator.toint() + row.query_row * Int64( + row.extra_indices.stride[0] + ) * Int64(4) + # Source tensors can have different layouts; select numeric addresses. + address = extra_ptr if use_extra else primary_ptr + indices = cute.make_tensor( + cute.make_ptr(Int32, address, assumed_align=4), + cute.make_layout((self.capacity,), stride=(1,)), + ) + # The first coordinate chooses a shared TMA descriptor, even when + # invalid. Preserve this source tag on every synthesized tail entry. + tag = Int32(-2147483648) if use_extra else Int32(0) + return SparseSourceTile( + (indices, origin, length, tag), valid_prefix=self.assume_valid_prefix + ) + + +class SparseSourceTile: + """Selected source and bounds; analytic prefix masks need no index pointer.""" + + def __init__(self, values, valid_prefix=False): + self.values = values + self.indices, self.origin, self.length, self.tag = values + self.valid_prefix = valid_prefix + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseSourceTile( + cutlass.new_from_mlir_values(self.values, values), self.valid_prefix + ) + + @cute.jit + def _bounded_index(self, token): + position = token - self.origin + # A zero-capacity source binds one dummy index. Even masked scalar + # loads stay inside its backing; validity is determined separately. + safe = cute.math.min( + cute.math.max(position, Int32(0)), + cute.math.max(self.length - Int32(1), Int32(0)), + ) + index = Int32(self.indices[safe]) + valid = (position >= Int32(0)) & (position < self.length) + if cutlass.const_expr(not self.valid_prefix): + valid = valid & (index >= Int32(0)) + return index, valid + + @cute.jit + def load_row(self, token): + index, valid = self._bounded_index(token) + return (index if valid else Int32(0x7FFFFFFF)) | self.tag + + @cute.jit + def load_full_quad(self, token): + """Load an in-bounds quad; callers establish its source-local bounds.""" + values = cutlass.Pointer( + self.indices.iterator + token - self.origin, dtype=Int32 + ).load(count=4, alignment=4) + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + index = Int32(values[j]) + if cutlass.const_expr(self.valid_prefix): + rows[j] = index | self.tag + else: + rows[j] = (index if index >= Int32(0) else Int32(0x7FFFFFFF)) | self.tag + return rows + + @cute.jit + def load_quad(self, token): + """One vector load for an active quad; bounded scalar loads for tails.""" + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + rows[j] = Int32(0x7FFFFFFF) | self.tag + position = token - self.origin + if position >= Int32(0) and position + Int32(3) < self.length: + values = self.load_full_quad(token) + for j in cutlass.range_constexpr(4): + rows[j] = values[j] + elif position + Int32(3) >= Int32(0) and position < self.length: + for j in cutlass.range_constexpr(4): + rows[j] = self.load_row(token + Int32(j)) + return rows + + @cute.jit + def is_invalid(self, token, request): + del request + if cutlass.const_expr(self.valid_prefix): + return (token < self.origin) | (token >= self.origin + self.length) + _, valid = self._bounded_index(token) + return ~valid + + @cute.jit + def tile_is_full(self, start, live_length): + return start + Int32(128) <= cute.math.min( + self.origin + self.length, live_length + ) + + @cute.jit + def prefix_mask32(self, start, live_length): + # Source/tile boundaries are aligned; start is within this source. + end = cute.math.min(self.origin + self.length, live_length) + valid = cute.math.min(cute.math.max(end - start, Int32(0)), Int32(32)) + # Avoid an undefined shift by 32, including on an unselected SSA arm. + bits = Int32(-1) << (valid & Int32(31)) + return bits if valid < Int32(32) else Int32(0) + + +class SparseKVLengths: + """Adapter for the scheduler's lengths[query_row] interface only.""" + + def __init__(self, values, has_extra=True): + self.values = values + self.has_extra = has_extra + self.primary_lengths, self.extra_lengths, self.bound_length, rows = values + self.shape = ( + (rows,) if self.bound_length is not None else self.primary_lengths.shape + ) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseKVLengths( + cutlass.new_from_mlir_values(self.values, values), self.has_extra + ) + + @cute.jit + def __getitem__(self, query_row): + if cutlass.const_expr(self.bound_length is not None): + return self.bound_length + return _sparse_execution_length( + Int32(self.primary_lengths[query_row]), + Int32(self.extra_lengths[query_row]) + if cutlass.const_expr(self.has_extra) + else Int32(0), + ) + + +@cute.jit +def _sparse_execution_length(primary_length, extra_length): + primary_tiles = (primary_length + Int32(127)) // Int32(128) + extra_tiles = (extra_length + Int32(127)) // Int32(128) + # Empty rows still execute a masked tile to publish O=0 and LSE=-inf. + return cute.math.max((primary_tiles + extra_tiles) * Int32(128), Int32(1)) + + +@cute.jit +def bind_sparse_kv_indices( + primary_indices, + extra_indices, + primary_lengths, + extra_lengths, + capacity: cutlass.Constexpr[int], + *, + query_row=None, + assume_valid_prefix: cutlass.Constexpr[bool] = False, + has_extra: cutlass.Constexpr[bool] = True, +): + """Bind index readers and scheduler lengths without copying metadata.""" + indices = SparseKVIndices( + ( + primary_indices, + extra_indices, + primary_lengths, + extra_lengths, + None, + None, + None, + ), + capacity, + assume_valid_prefix, + has_extra, + ) + if cutlass.const_expr(query_row is not None): + indices = indices.bind_query_row(query_row) + length = _sparse_execution_length(indices.primary_length, indices.extra_length) + lengths = SparseKVLengths( + (None, None, length, primary_lengths.shape[0]), has_extra + ) + else: + lengths = SparseKVLengths( + (primary_lengths, extra_lengths, None, None), has_extra + ) + return indices, lengths diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py index a5e7e810d770..e35517b09e93 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py @@ -45,7 +45,7 @@ def local_kv_tile_idx( return Int32(inst_id) if cutlass.const_expr(section == MlaStage.Loop): base = stage_info.loop_offset * Int32(cfg.num_insts_kv) - if cutlass.const_expr(is_v): + if cutlass.const_expr(is_v or (cfg.serial_swap_reuse and inst_id == 1)): return base + Int32(inst_id) return base + Int32(cfg.num_insts_kv + inst_id) if cutlass.const_expr(cfg.total_kv_tiles <= cfg.num_insts_kv): @@ -109,6 +109,8 @@ def runtime_split_pruning_is_profitable(configured_splits_kv: int) -> bool: def _runtime_configured_local_kv_tiles(cfg: MlaConfig, seq_len_kv): """Return the instruction-aligned local span for configured split capacity.""" + if cutlass.const_expr(cfg.fixed_sparse_kv_tiles > 0): + return Int32(cfg.fixed_sparse_kv_tiles) total_kv_tiles = runtime_total_kv_tiles(cfg, seq_len_kv) num_insts_kv = Int32(cfg.num_insts_kv) tiles_per_group = Int32(cfg.num_ctas_per_seq_kv) * num_insts_kv @@ -134,6 +136,16 @@ def runtime_num_ctas_kv(cfg: MlaConfig, seq_len_kv): @cute.jit def runtime_local_kv_tiles(cfg: MlaConfig, seq_len_kv): """Return the padded local KV tile count assigned to each KV CTA group.""" + if cutlass.const_expr(cfg.fixed_sparse_kv_tiles > 0): + return Int32(cfg.fixed_sparse_kv_tiles) + if cutlass.const_expr( + cfg.sparse_direct and cfg.seq_len_kv <= 256 and cfg.num_ctas_per_seq_kv == 1 + ): + # The short direct schedule executes its planned padded tiles. Every + # consumer must use that same domain: using the smaller live length + # here makes tail PV overwrite O with its final masked tile instead + # of accumulating into the earlier nonempty tile. + return Int32(cfg.local_kv_tiles(cfg.total_kv_tiles)) if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): return runtime_total_kv_tiles(cfg, seq_len_kv) return _runtime_configured_local_kv_tiles(cfg, seq_len_kv) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py index c0e29277c516..3949e8ca5eb8 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py @@ -161,7 +161,7 @@ def is_throughput_2cta_mla_supported_shape( dtype in ("bf16", "e4m3") and out_dtype in ("bf16", "e4m3") and latent_dim == 512 - and rope_dim == 64 + and rope_dim in (0, 64) and page_size in SUPPORTED_MLA_PAGE_SIZES and 1 <= num_heads <= 128 and seq_len_q >= 1 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py new file mode 100644 index 000000000000..8f6100825f04 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py @@ -0,0 +1,496 @@ +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0. +"""Sparse MLA policy: choose a schedule, size buffers, then fill the SM grid. + +``rows`` is planned B*Q, ``query_length`` is planned Q, and ``capacity`` counts +selected slots (each source rounded to 128, at least 256 in total), not raw KV +context. Live lengths and indices remain runtime inputs. A query/head grid has +rows*ceil(heads/tile_Q) work tiles; a 2CTA tile consumes two CTAs. + +Selection tries qualified schedules before the small-grid defaults, then +sizes buffers and fills the remaining grid with KV/V partitions. Row-to-SM +ratios estimate available work, not achieved occupancy. Head/capacity/Q and +wave cutoffs are measured GB300 crossovers, not hardware or model limits; +resource and layout eligibility is checked separately by the kernel config. +""" + +from dataclasses import dataclass, replace +from math import ceil + + +@dataclass(frozen=True) +class SparseMlaProfile: + """Kernel schedule; zero tile/split/V-partition values are resolved below.""" + + family: str + tile_size_q: int = 0 + split_kv: int = 1 + scheduler: str = "nonpersistent" + reduction: str = "gmem_separate" + head_dim_ctas: int = 1 + gather_issue_warps: int = 4 + balanced_registers: bool = False + reuse_kv: bool = False + reuse_kv_stages: int = 0 + paired_correction: bool = False + single_kv_stream: bool = False + defer_max_update: bool = False + uniform_offset_cache: bool = False + kv_pipeline_stages: int = 0 + kv_tile_size: int = 128 + page_pipeline_stages: int = 0 + + +def _base_fp8(rows, heads, capacity, sm_count): + """Small-grid tile and split selection for FP8.""" + # M64 shares KV across more heads but can leave a small grid underfilled. + # M16 exposes more tasks at the cost of repeating KV work for head tiles. + throughput_h32 = heads == 32 and rows * 2 >= sm_count and (capacity >= 512) + long_list = capacity >= 8192 + m16_work = rows * ((heads + 15) // 16) + m16_one_step_work = m16_work * ((capacity + 255) // 256) + use_keep = heads >= 64 and ( + capacity > 1024 and m16_one_step_work > sm_count or m16_work >= sm_count // 2 + ) + use_keep = use_keep or (long_list and heads >= 16 and (rows >= 8)) + use_keep = use_keep or (capacity <= 256 and heads >= 16) + use_keep = use_keep or throughput_h32 + base_clusters = rows * ((heads + 127) // 128) + keep_work = rows * ((heads + 63) // 64) + kv_tiles = (capacity + 127) // 128 + split_limit = min(kv_tiles, max(1, sm_count // 4 // keep_work)) + kv_steps = (kv_tiles + split_limit - 1) // split_limit + keep_splits = (kv_tiles + kv_steps - 1) // kv_steps + prefer_wide_q = ( + heads >= 64 + and capacity >= 512 + and (kv_steps <= 2) + and (keep_work * keep_splits * 4 >= sm_count // 2) + ) + cluster_splits = 1 << ((capacity + 255) // 256 - 1).bit_length() + m8_work = rows * ((heads + 7) // 8) + one_tile_splits = (capacity + 127) // 128 + if heads == 32 and capacity >= 512 and rows >= sm_count: + # Enough query rows amortize padded M64 work. Retaining two FP8 K + # tiles saves V reloads; uniform coordinates amortize repeated D slices. + return SparseMlaProfile( + family="keep", + reuse_kv=True, + reuse_kv_stages=10, + uniform_offset_cache=True, + scheduler="clc", + ) + if ( + heads >= 64 + and 512 <= capacity <= 1024 + and (sm_count // 2 <= keep_work * one_tile_splits <= sm_count) + ): + # One K tile per split fills about half to one device wave. Further + # splitting cannot shorten the tile and only increases merge traffic. + return SparseMlaProfile( + family="keep", split_kv=one_tile_splits, reuse_kv=True, reuse_kv_stages=10 + ) + if heads == 32 and 32 <= rows <= sm_count // 2 and (capacity >= 512): + # This intermediate grid favors M64 reuse, with a second split only + # for longer lists; additional splits lose to their merge overhead. + return SparseMlaProfile( + family="keep", + split_kv=1 if capacity <= 1024 else 2, + reuse_kv=True, + reuse_kv_stages=10, + ) + if heads <= 16 and rows * 4 >= sm_count and (capacity >= 512): + # Narrow Q tiles avoid padded heads; twelve stages retain two FP8 + # tiles plus four prefetched D slices (four slices per D512 tile). + return SparseMlaProfile( + family="swap", + tile_size_q=8 if heads <= 8 else 16, + split_kv=2 if capacity > 1024 and rows * 2 <= sm_count else 1, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=12, + ) + if capacity <= 256 and heads >= 16 and (m8_work * 4 <= sm_count // 2): + # Short lists offer little KV parallelism. V partitions expose work + # despite repeating QK/softmax, while M8 limits padded arithmetic. + return SparseMlaProfile( + family="swap", tile_size_q=8, head_dim_ctas=4, balanced_registers=True + ) + # Partition V when the wide tile alone cannot fill a resident wave. + if prefer_wide_q: + return SparseMlaProfile(family="keep", split_kv=keep_splits, head_dim_ctas=4) + elif ( + 1 < cluster_splits <= 4 + and heads % 16 == 0 + and (m16_work * cluster_splits <= sm_count // 2) + ): + # A small two/four-split cluster avoids a standalone GMEM finisher. + # Its coordination/residency cost loses for larger grids or split sets. + return SparseMlaProfile( + family="swap", + split_kv=cluster_splits, + reduction="cluster", + head_dim_ctas=0, + balanced_registers=True, + ) + if ( + long_list + and heads >= 128 + and (rows >= 8) + and (base_clusters * 2 <= sm_count // 2) + ): + # Wide heads amortize the 2CTA layout; long lists provide enough KV + # work to split when the query grid alone is small. + return SparseMlaProfile(family="2cta", split_kv=0, gather_issue_warps=1) + elif heads >= 128 and capacity >= 512 and (base_clusters * 2 >= sm_count // 2): + return SparseMlaProfile(family="2cta", gather_issue_warps=1) + return SparseMlaProfile( + family="keep" if use_keep else "swap", + split_kv=1 if throughput_h32 or (use_keep and capacity <= 1024) else 0, + head_dim_ctas=1 if throughput_h32 else 0, + balanced_registers=not use_keep, + reuse_kv=capacity >= 512 and use_keep, + reuse_kv_stages=10 if capacity >= 512 and use_keep else 0, + ) + + +def _base_bf16(rows, heads, capacity, sm_count): + """Small-grid tile and split selection for BF16.""" + # BF16 BK128 cannot retain two complete D512 tiles in SMEM. A narrow + # grid favors more small-Q tasks; a large grid amortizes M64's padding + # and benefits from retaining one tile instead of reloading V. + throughput_h32 = heads == 32 and rows * 2 >= sm_count and capacity >= 512 + throughput_bf16 = ( + heads >= 32 and (rows >= sm_count or throughput_h32) and (capacity >= 512) + ) + short_list = capacity <= 256 and heads <= 64 + tile = (8 if heads <= 16 else 16) if short_list else 16 if heads <= 32 else 64 + if throughput_bf16: + tile = 64 + profile = SparseMlaProfile( + family="swap" if tile < 64 else "keep", + tile_size_q=tile, + split_kv=1 if short_list or throughput_bf16 else 0, + head_dim_ctas=1 if throughput_bf16 else 0, + balanced_registers=tile < 64, + reuse_kv=throughput_bf16, + reuse_kv_stages=4 if throughput_bf16 else 0, + uniform_offset_cache=throughput_bf16 + and rows >= sm_count + and ( + capacity >= 2048 + or (capacity >= 1024 and rows >= 2 * sm_count) + or rows >= 8 * sm_count + ), + ) + cluster_splits = 1 << ((capacity + 255) // 256 - 1).bit_length() + m16_work = rows * ((heads + 15) // 16) + if ( + not short_list + and tile == 16 + and (1 < cluster_splits <= 4) + and (heads % 16 == 0) + and (m16_work * cluster_splits <= sm_count // 2) + ): + # Same small-grid tradeoff as FP8: cluster fusion saves a launch, but + # only the qualified two/four-split M16 layouts justify its overhead. + profile = replace(profile, split_kv=cluster_splits, reduction="cluster") + return profile + + +def _choose_fp8(rows, heads, query_length, capacity, sm_count): + """Prefer resident split grids or sustained-throughput schedules when eligible.""" + # These short-Q regions favor one KV stream with retained tiles and a + # split grid near one device wave. The Q<=8 cutoff is empirical; the + # kernel also supports prefill. Powers of two are policy, not a KV ABI. + if ( + query_length <= 8 + and capacity >= 512 + and ( + heads <= 16 + and 32 <= rows <= sm_count + and (not (heads <= 8 and capacity >= 2048 and (rows * 2 >= sm_count))) + or (heads == 32 and 16 <= rows <= sm_count // 8 and (capacity >= 2048)) + ) + ): + base_ctas = rows * ((heads + 15) // 16) + wave_splits = max(1, sm_count // base_ctas) + wave_splits = 1 << (wave_splits.bit_length() - 1) + return SparseMlaProfile( + family="swap", + split_kv=wave_splits, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=10, + single_kv_stream=True, + ) + if ( + heads == 128 + and query_length <= 8 + and (capacity >= 512) + and ((8 if capacity <= 6 * 128 else 16) <= rows <= sm_count // 4) + ): + # Two CTAs cover H128 without repeating KV loads across M64 tiles. + # Up to six KV tiles, avoiding the separate split-KV merge wins from + # eight query rows. Longer lists keep the measured sixteen-row + # crossover; these are GB300 latency crossovers, not kernel limits. + # Four splits help only the small, long-list grid measured here. + small_long_grid = capacity >= 2048 and rows <= sm_count // 8 + return SparseMlaProfile( + family="2cta", split_kv=4 if small_long_grid else 1, gather_issue_warps=1 + ) + if heads == 64 and capacity >= 512 and (rows * 2 >= sm_count): + # Sustained M64 work amortizes eight issuers, expanded coordinate + # registers and CLC setup. Paired correction reduces stats traffic. + return SparseMlaProfile( + family="keep", + gather_issue_warps=8, + reuse_kv=True, + reuse_kv_stages=10, + uniform_offset_cache=True, + scheduler="clc", + paired_correction=True, + ) + if heads == 128 and rows * 2 >= sm_count: + # Planned capacity is at least 256. Once the two-CTA grid fills a + # device wave, share KV across H128 even for short selected lists. + # BF16's CLC choice below is a separately qualified schedule. + return SparseMlaProfile(family="2cta", scheduler="static", gather_issue_warps=1) + if ( + heads <= 16 + and capacity >= 512 + and ( + capacity <= 1024 + and rows >= (sm_count if heads > 8 else 2 * sm_count) + or rows >= 8 * sm_count + ) + ): + # At high row counts, saved KV traffic outweighs M64 head padding. + # Smaller grids retain M8/M16 to preserve independent work. + return SparseMlaProfile( + family="keep", + reuse_kv=True, + reuse_kv_stages=10, + uniform_offset_cache=True, + scheduler="clc", + ) + return _base_fp8(rows, heads, capacity, sm_count) + + +def _choose_bf16(rows, heads, query_length, capacity, sm_count): + """Prefer resident split grids or sustained-throughput schedules when eligible.""" + # Small-head workloads can retain one BK128 tile with five/six 32-KiB + # stages: four hold D512, the rest prefetch. PV must release that tile + # before next QK; this favors a single stream in the measured regions. + if ( + heads <= 16 + and capacity >= 512 + and (rows * 2 >= sm_count) + and ( + query_length <= 8 + or ( + heads <= 8 + and capacity <= 1024 + and (query_length >= 128) + and (rows >= 8 * sm_count) + ) + ) + ): + return SparseMlaProfile( + family="swap", + tile_size_q=8 if heads <= 8 else 16, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=6 + if heads > 8 or (query_length <= 8 and rows >= 4 * sm_count) + else 5, + single_kv_stream=True, + defer_max_update=True, + scheduler="clc", + ) + if ( + heads in (8, 16, 32, 64) + and capacity >= 512 + and (query_length >= 128) + and (rows >= 8 * sm_count) + ): + # Long-Q grids have ample tasks. Static persistence amortizes CTA + # setup here; wider heads favor M64 reuse, narrow heads avoid padding. + keep = heads >= 32 + return SparseMlaProfile( + family="keep" if keep else "swap", + tile_size_q=64 if keep else 8 if heads <= 8 else 16, + balanced_registers=not keep, + reuse_kv=True, + reuse_kv_stages=5 if keep else 6, + uniform_offset_cache=keep, + scheduler="static", + ) + if ( + heads in (32, 64) + and rows * 2 >= sm_count + and (query_length <= 8) + and (capacity >= 512) + ): + # Decode with enough rows favors M64 reuse. Buffer sizing below + # supplies its qualified gather/cache budget and scheduler crossover. + return SparseMlaProfile( + family="keep", + reuse_kv=True, + reuse_kv_stages=5, + uniform_offset_cache=capacity >= 2048 + and rows >= sm_count + or (capacity >= 1024 and rows >= 4 * sm_count) + or rows >= 8 * sm_count, + ) + if heads <= 16 and rows * 2 >= sm_count and (capacity >= 512): + return SparseMlaProfile( + family="swap", + tile_size_q=8 if heads <= 8 else 16, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=6, + ) + if heads == 128 and 2 <= sm_count // (2 * rows) <= 4 and (capacity >= 512): + # Two M64 head tiles with BK64 reuse and a few KV splits fill this + # small H128 grid. Larger grids favor the wide 2CTA path below. + wave_splits = max(1, sm_count // 2 // rows) + return SparseMlaProfile( + family="keep", + split_kv=wave_splits, + gather_issue_warps=8, + uniform_offset_cache=True, + reuse_kv=True, + reuse_kv_stages=10, + kv_tile_size=64, + page_pipeline_stages=2, + defer_max_update=True, + ) + if ( + capacity >= 512 + and rows * 4 >= sm_count + and (heads == 128 or (heads in (32, 64) and rows * 2 <= sm_count)) + ): + # BF16 2CTA amortizes wide-head work once enough clusters are present. + # More K/V stages overlap gathers; coordinate caching costs registers, + # so it is enabled only for the qualified H128 configuration. + return SparseMlaProfile( + family="2cta", + kv_pipeline_stages=8, + uniform_offset_cache=heads == 128, + scheduler="clc" if heads == 128 and rows >= sm_count else "nonpersistent", + ) + return _base_bf16(rows, heads, capacity, sm_count) + + +def _size_buffers(profile, *, rows, heads, query_length, capacity, dtype, sm_count): + # These are coupled measured budgets, not independent universal wins: + # issuer/cache registers compete with softmax/correction, and extra KV + # stages compete with Q, P, metadata and scheduler SMEM. + if ( + dtype == "fp8" + and heads == 128 + and (capacity >= 2048) + and (query_length <= 8) + and (rows >= 4 * sm_count) + ): + profile = replace(profile, balanced_registers=True) + if ( + dtype == "bf16" + and profile.family == "keep" + and (profile.reuse_kv_stages == 5) + and (heads in (32, 64)) + ): + # Both prefill and decode wide-BF16 reuse schedules share this budget. + profile = replace( + profile, + gather_issue_warps=8, + uniform_offset_cache=True, + defer_max_update=True, + ) + if query_length <= 8 and rows >= 2 * sm_count: + profile = replace(profile, scheduler="clc") + if ( + heads == 32 + and query_length <= 8 + and (rows >= 4 * sm_count) + and (capacity >= 1024) + ): + # BK64 halves each BF16 stage to 16 KiB: ten stages retain two + # four-slice K tiles plus two slices ahead. This restores QK/PV + # overlap, at the cost of twice as many token-tile iterations. + profile = replace( + profile, kv_tile_size=64, reuse_kv_stages=10, page_pipeline_stages=2 + ) + if ( + dtype == "bf16" + and profile.family == "keep" + and (heads == 64) + and (rows >= sm_count) + and (capacity >= 512) + and (profile.split_kv == 1) + ): + # The qualified H64 BK64 schedule serves both prefill and decode. + # Two metadata stages let preparation advance while issuers reuse + # current coordinates. The bounded maximum update is BF16-only. + profile = replace( + profile, + kv_tile_size=64, + reuse_kv=True, + reuse_kv_stages=10, + page_pipeline_stages=2, + gather_issue_warps=8, + uniform_offset_cache=True, + scheduler="clc", + defer_max_update=True, + ) + return profile + + +def select_sparse_mla_profile(*, rows, heads, query_length, capacity, dtype, sm_count): + """Resolve one complete schedule from the validated public plan bounds.""" + choose = _choose_fp8 if dtype == "fp8" else _choose_bf16 + profile = choose(rows, heads, query_length, capacity, sm_count) + profile = _size_buffers( + profile, + rows=rows, + heads=heads, + query_length=query_length, + capacity=capacity, + dtype=dtype, + sm_count=sm_count, + ) + tile = profile.tile_size_q or {"swap": 16, "keep": 64, "2cta": 128}[profile.family] + work = rows * ceil(heads / tile) + splits = profile.split_kv + if splits == 0: + # Fill a resident wave using instruction-aligned KV partitions. + # More splits shorten each CTA's work but increase partial-O/LSE merge + # traffic. Two-CTA work targets half as many cluster tasks as SMs. + step = 128 if tile >= 64 else 256 + steps = max(1, ceil(capacity / step)) + target = sm_count // 2 if profile.family == "2cta" else sm_count + partitions = min(steps, max(1, target // work), 128) + splits = ceil(steps / ceil(steps / partitions)) + v_ctas = profile.head_dim_ctas + if v_ctas == 0: + # V partitioning exposes more tasks but repeats QK/softmax across them. + # Prefer it only while the query/KV grid still lacks enough work. + remaining = work * splits + v_ctas = ( + 1 + if profile.family == "2cta" or remaining * 2 > sm_count + else 4 + if remaining * 4 <= sm_count + else 2 + ) + # Retaining K through PV requires a single owner of the V dimension. + reuse = profile.reuse_kv and v_ctas == 1 + return replace( + profile, + tile_size_q=tile, + split_kv=splits, + head_dim_ctas=v_ctas, + reuse_kv=reuse, + reuse_kv_stages=profile.reuse_kv_stages if reuse else 0, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py index af89becc56eb..fe2461ad74c4 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py @@ -25,9 +25,6 @@ from ..helpers.mask import MaskType, normalize_mask_type -# Softmax converts natural-scale scores to exp2 with log2(e). -LOG2_E = 1.4426950408889634074 - # The separate reducer follows the public MLA output contract: partial O is # stored as BF16 while LSE and the final accumulation remain FP32. One # 512-thread CTA owns eight D512 rows, with each thread moving one 16-byte @@ -49,6 +46,10 @@ # Each transpose-TMA issue stages one 64-element slice of the latent dimension. V_TMA_LATENT_ELEMENTS = 64 +# Two N slices of the final M128xD512 FP8 output, split across two CTAs. +OUTPUT_TMA_STAGES = 2 +OUTPUT_TMA_STAGE_BYTES = 64 * 256 + def ceil_div(a: int, b: int) -> int: """Return ``ceil(a / b)`` for positive integer divisors.""" @@ -159,8 +160,12 @@ class MlaDecodeConfig: iterations_qk_stages: int = 5 iterations_pv_stages: int = 4 + # Final quantized output can use a staged TMA store on padding warp 10. + use_tma_output: bool = False + # Pipeline stage counts for the captured schedule resources. The combined # K/V stage count keeps enough delayed-V stages live for K-before-V overlap. + sparse_offset_stages: int = 4 load_q_stage: int = 1 load_k_stage: int = 3 load_v_stage: int = 2 @@ -179,6 +184,9 @@ class MlaDecodeConfig: correction_warp_ids: Tuple[int, ...] = (4, 5, 6, 7) mma_warp_id: int = 8 load_tma_warp_id: int = 9 + load_k_warp_id: int = 9 + load_v_warp_id: int = 10 + load_num_warps: int = 1 pv_mma_warp_id: int = 11 empty_warp_ids: Tuple[int, ...] = (11,) second_compute_warp_ids: Tuple[int, ...] = () @@ -258,6 +266,7 @@ class MlaDecodeConfig: # Use block_split_kvs[batch] as a per-batch cap before runtime K contracts # the useful split prefix. Grid and workspace geometry retain the maximum. is_var_split_kv: bool = False + cache_uniform_sparse_pages: bool = False # Causal is bottom-right aligned for speculative decode. Dense still masks # the ordinary per-batch KV tail at ``cache_seqs[batch]``. mask_type: str = MaskType.CAUSAL.value @@ -334,6 +343,11 @@ def make_mla_decode_config( is_var_seq: bool = False, is_var_split_kv: bool = False, mask_type: MaskType | str = MaskType.CAUSAL, + sparse_gather_warps: int = 1, + sparse_kv_stages: int = 0, + cache_uniform_sparse_pages: bool = False, + balance_sparse_registers: bool = False, + fuse_output_quant: bool = False, ) -> MlaDecodeConfig: """Create and populate a MlaDecodeConfig from problem parameters.""" cfg = MlaDecodeConfig() @@ -373,7 +387,9 @@ def _require_divisible( if rope_dim < 0: raise ValueError(f"rope_dim must be non-negative, got {rope_dim}") _require_positive("page_size", page_size) - if page_size not in SUPPORTED_MLA_PAGE_SIZES: + if page_size not in SUPPORTED_MLA_PAGE_SIZES and not ( + page_size == 1 and rope_dim == 0 + ): raise ValueError( f"page_size must be one of {SUPPORTED_MLA_PAGE_SIZES}, got {page_size}" ) @@ -403,6 +419,81 @@ def _require_divisible( cfg.other_reg_num = 32 cfg.mma_o_stage = 2 cfg.epilogue_sync_bar_id = 4 + if page_size == 1: + # Independent QK/PV streams retain two whole K/V tiles each. A + # third V tile costs 32 KiB per CTA without improving overlap in + # the qualified sparse schedule. Keep four metadata slots so + # index prefetch can run ahead of both streams. + cfg.load_k_stage = 2 + cfg.load_v_stage = 2 + cfg.use_fp8_dual_softmax_schedule = False + cfg.second_compute_warp_ids = () + cfg.num_softmax_groups = 1 + cfg.load_k_warp_id = 12 + cfg.load_v_warp_id = 16 + cfg.load_num_warps = 4 + cfg.threads_per_cta = cfg.threads_per_warp * 20 + cfg.other_reg_num = 48 + + if sparse_gather_warps != 1: + if sparse_gather_warps not in (4, 8) or qkv_dtype != "bf16" or page_size != 1: + raise ValueError("multi-warp combined loading requires BF16 sparse 2CTA") + cfg.load_num_warps = sparse_gather_warps + cfg.load_tma_warp_id = 12 + cfg.load_k_warp_id = 12 + cfg.threads_per_cta = 384 + 32 * sparse_gather_warps + cfg.softmax_reg_num = 128 if sparse_gather_warps == 8 else 192 + cfg.correction_reg_num = 160 if sparse_gather_warps == 8 else 192 + cfg.other_reg_num = 64 + + if sparse_kv_stages: + if qkv_dtype != "bf16" or page_size != 1 or sparse_gather_warps not in (4, 8): + raise ValueError("pipeline tuning requires multi-warp BF16 sparse 2CTA") + if sparse_kv_stages < 4: + raise ValueError("combined KV pipeline requires at least four stages") + cfg.load_kv_stage = sparse_kv_stages + + if cache_uniform_sparse_pages: + if qkv_dtype != "bf16" or page_size != 1 or sparse_gather_warps not in (4, 8): + raise ValueError("uniform K cache requires multi-warp BF16 sparse loading") + cfg.cache_uniform_sparse_pages = True + cfg.softmax_reg_num = 128 + cfg.correction_reg_num = 160 if sparse_gather_warps == 8 else 192 + cfg.other_reg_num = 64 if sparse_gather_warps == 8 else 96 + + if balance_sparse_registers: + if page_size != 1 or qkv_dtype != "e4m3": + raise ValueError("balanced registers require sparse FP8") + # Sparse FP8 keeps 64 registers for the gather issuers. Split the + # remaining budget equally: 128 softmax registers spilled live loop + # state on every KV tile, while correction fits in 144 registers. + # 4*144 + 4*144 + 12*64 = 1920 warp-registers, the same CTA total + # as the previous 128/160/64 allocation (61440 registers). + cfg.softmax_reg_num = 144 + cfg.correction_reg_num = 144 + cfg.other_reg_num = 64 + + if fuse_output_quant and cfg.is_fp8_qkv() and cfg.threads_per_cta == 640: + # Quantization keeps an output block and rotary coefficients live. + # The gather tasks fit in 48 registers; use their surplus for the + # correction task while preserving the CTA's 61,440-register budget. + cfg.softmax_reg_num = 160 + cfg.correction_reg_num = 176 + cfg.other_reg_num = 48 + + # SETMAXNREG redistributes the CTA's initial allocation, rounded down + # to eight registers/thread; the remainder of the SM file is not credit. + initial_regs = 65536 // (cfg.threads_per_cta * 8) * 8 + smx_warps = cfg.num_softmax_groups * cfg.num_compute_warps + cor_warps = cfg.num_compute_warps + other_warps = cfg.threads_per_cta // 32 - smx_warps - cor_warps + claimed = ( + smx_warps * cfg.softmax_reg_num + + cor_warps * cfg.correction_reg_num + + other_warps * cfg.other_reg_num + ) * 32 + if claimed > initial_regs * cfg.threads_per_cta: + raise ValueError("2CTA task registers exceed the CTA's initial allocation") # Derived MMA tilers. FP8 latent QK uses K=128 while the separate RoPE MMA # keeps K=64. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py index 5ea23ffb4238..ffc5c7fb609c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py @@ -32,7 +32,7 @@ import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.experimental.cuda as cuda -from cutlass import Int32, Int64 +from cutlass import Float32, Int32, Int64 from cutlass.cute.testing import assert_ as runtime_assert from cutlass.cute.nvgpu import OperandMajorMode, tcgen05 from cutlass.experimental import primitives as prims @@ -50,14 +50,15 @@ ) from .config import ( - LOG2_E, + OUTPUT_TMA_STAGES, + OUTPUT_TMA_STAGE_BYTES, REDUCTION_ROWS_PER_CTA, REDUCTION_THREADS_PER_ROW, V_TMA_LATENT_ELEMENTS, MlaDecodeConfig, make_mla_decode_config, ) -from ..helpers.constants import MAX_MLA_SPLITS_KV, TMEM_DEALLOC_MBAR_THREADS +from ..helpers.constants import LOG2_E, MAX_MLA_SPLITS_KV, TMEM_DEALLOC_MBAR_THREADS from ..helpers.mask import MaskType, mask_visible_k_length, normalize_mask_type from ..helpers.query import ( FlatQueryTileLayout, @@ -86,6 +87,7 @@ from ..helpers.math import qkv_dtype from ..parallel_reduction_topology import ( ParallelReductionTopology, + make_balanced_parallel_reduction_topology, make_q128_wave_limited_parallel_reduction_topology, should_use_q128_g1_parallel_reducer, validate_parallel_reduction_workspace, @@ -107,6 +109,7 @@ MlaInterleavedTask, MlaTask, create_load_k_task, + create_sparse_offsets_task, create_load_v_task, create_load_tma_task, create_mma_task, @@ -115,6 +118,7 @@ create_softmax_task, create_correction_task, create_padding_task, + create_output_tma_task, create_scheduler_task, ) @@ -132,6 +136,7 @@ def build_mla_decode_task_manager( smem_kc_arr=None, smem_vc_arr=None, smem_p_arr=None, + sparse_offsets_smem=None, # TMA descriptors (None for validate-only) tma_desc_q_latent=None, tma_desc_q_rope=None, @@ -145,6 +150,12 @@ def build_mla_decode_task_manager( tidx=None, # GMEM output tensors output=None, + atten_sinks=None, + lse_in_natural_log=False, + output_quant=None, + output_quant_args=None, + smem_o=None, + tma_o_desc=None, acc_output=None, lse=None, acc_lse=None, @@ -200,7 +211,7 @@ def build_mla_decode_task_manager( ) # Cooperative groups - tma_producer_group = pipeline.CooperativeGroup(Agent.Thread) # elect_one TMA + tma_producer_group = pipeline.CooperativeGroup(Agent.Thread) umma_hw_group = pipeline.CooperativeGroup(Agent.Thread) # UMMA hardware # Softmax warps: 4 warps * 32 threads * 2 CTAs for cluster-scoped consumer @@ -227,10 +238,15 @@ def build_mla_decode_task_manager( num_stages=cfg.load_q_stage, num_bytes=cfg.tma_copy_q_bytes, producer_group=tma_producer_group, + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_copy_q_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_copy_q_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), ) if cfg.use_fp8_split_mma_schedule: @@ -241,19 +257,29 @@ def build_mla_decode_task_manager( num_stages=cfg.load_k_stage, num_bytes=cfg.tma_copy_k_tile_bytes, producer_group=tma_producer_group, + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_copy_k_tile_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_copy_k_tile_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), ) smem_v_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.load_v_stage, num_bytes=cfg.tma_copy_v_tile_bytes, producer_group=tma_producer_group, + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_copy_v_tile_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_copy_v_tile_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), ) smem_kv_pipeline_cfg = None else: @@ -265,7 +291,12 @@ def build_mla_decode_task_manager( consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_kc_subtile_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_kc_subtile_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, ) smem_k_pipeline_cfg = None smem_v_pipeline_cfg = None @@ -329,10 +360,31 @@ def build_mla_decode_task_manager( # Create resource instances # ────────────────────────────────────────────────────────────── + sparse_offsets_enabled = cfg.page_size == 1 and cfg.use_fp8_split_mma_schedule + shared_mask_enabled = ( + sparse_offsets_enabled + and not cfg.use_fp8_dual_softmax_schedule + and not getattr(page_offsets, "assume_valid_prefix", False) + ) + # Prefix masks use live lengths only: softmax no longer consumes the + # offset ring. Remove its arrivals along with its task dependency; K/V + # loaders remain consumers with the same addressing and synchronization. + offsets_cfg = None + if sparse_offsets_enabled: + offsets_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=cfg.sparse_offset_stages, + producer_group=pipeline.CooperativeGroup(Agent.Thread, 32), + consumer_group=pipeline.CooperativeGroup( + Agent.Thread, + (2 * cfg.load_num_warps + (4 if shared_mask_enabled else 0)) * 32, + ), + cta_layout_vmnk=cluster_shape_vmnk, + ) page_offset_window = PageOffsetWindowResource( page_offsets=page_offsets, + smem_sparse_offsets=sparse_offsets_smem, cfg=cfg, - pipeline_config=None, + pipeline_config=offsets_cfg, name="page_offset_window", ) @@ -351,6 +403,7 @@ def build_mla_decode_task_manager( if cfg.use_fp8_split_mma_schedule: smem_k = SmemKResource( + smem_sparse_offsets=sparse_offsets_smem, smem_k=smem_kc_arr, page_offsets=page_offsets, tma_desc_c_latent=tma_desc_c_latent, @@ -361,9 +414,11 @@ def build_mla_decode_task_manager( name="smem_k", ) smem_v = SmemVResource( + smem_sparse_offsets=sparse_offsets_smem, smem_v=smem_vc_arr, page_offsets=page_offsets, tma_desc_c_transpose=tma_desc_c_transpose, + tma_desc_extra_v=tma_desc_q_rope, logical_seq_len_q=logical_seq_len_q, cfg=cfg, pipeline_config=smem_v_pipeline_cfg, @@ -376,6 +431,7 @@ def build_mla_decode_task_manager( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_c_transpose=tma_desc_c_transpose, + tma_desc_extra_v=tma_desc_q_rope, cfg=cfg, pipeline_config=smem_kv_pipeline_cfg, name="smem_kv", @@ -384,6 +440,7 @@ def build_mla_decode_task_manager( smem_v = None tmem_s = TmemSResource( + page_offsets=page_offsets, smem_q_latent=smem_q_latent_arr, smem_q_rope=smem_q_rope_arr, smem_p=smem_p_arr, @@ -420,6 +477,14 @@ def build_mla_decode_task_manager( name="tmem_o", ) + output_pipeline = None + if cfg.use_tma_output: + output_pipeline = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=OUTPUT_TMA_STAGES, + producer_group=compute_group_local, + consumer_group=pipeline.CooperativeGroup(Agent.Thread, 32), + cta_layout_vmnk=cluster_shape_vmnk, + ) gmem_o = GmemOResource( cfg=cfg, output=output, @@ -428,6 +493,12 @@ def build_mla_decode_task_manager( partial_lse=acc_lse, tmem_o_ref=tmem_o, tmem_corr_ref=tmem_corr, + atten_sinks=atten_sinks, + lse_in_natural_log=lse_in_natural_log, + output_quant=output_quant, + output_quant_args=output_quant_args, + smem_o=smem_o, + tma_o_desc=tma_o_desc, output_scale=None, # set at runtime softmax_scale_log2=None, # set at runtime smem_exchange=None, # set at runtime @@ -436,6 +507,7 @@ def build_mla_decode_task_manager( logical_num_heads_q=logical_num_heads_q, logical_seq_len_q=logical_seq_len_q, name="gmem_o", + pipeline_config=output_pipeline, ) # ────────────────────────────────────────────────────────────── @@ -451,12 +523,14 @@ def build_mla_decode_task_manager( load_k_task = create_load_k_task( smem_q, smem_k, + page_offsets=page_offset_window if sparse_offsets_enabled else None, work_queue=work_queue, domain=domain, num_registers=wg2_reg_count, ) load_v_task = create_load_v_task( smem_v, + page_offsets=page_offset_window if sparse_offsets_enabled else None, work_queue=work_queue, domain=domain, num_registers=wg2_reg_count, @@ -521,6 +595,7 @@ def build_mla_decode_task_manager( tmem_s, tmem_corr, smem_p, + page_offsets=page_offset_window if shared_mask_enabled else None, work_queue=work_queue, task_class=( MlaInterleavedTask @@ -580,13 +655,41 @@ def build_mla_decode_task_manager( ] task_list.extend([mma_task, softmax_task, correction_task]) + if sparse_offsets_enabled: + task_list.append( + create_sparse_offsets_task( + page_offset_window, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + ) + if cfg.use_tma_output: + task_list.append( + create_output_tma_task( + gmem_o, + work_queue=work_queue, + num_registers=wg2_reg_count, + ) + ) + elif cfg.use_fp8_split_mma_schedule and cfg.load_num_warps > 1: + task_list.append( + create_padding_task( + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + warp_idx=10 if sparse_offsets_enabled else 9, + num_warps=1 if sparse_offsets_enabled else 2, + ) + ) if not cfg.use_fp8_split_mma_schedule and use_clc_dynamic: padding_task = create_padding_task( work_queue=work_queue, task_class=non_interleaved_task_class, domain=task_domain, num_registers=wg2_reg_count, - warp_idx=10, + warp_idx=9 if cfg.load_num_warps > 1 else 10, + num_warps=2 if cfg.load_num_warps > 1 else 1, ) task_list.append(padding_task) scheduler_task = create_scheduler_task( @@ -604,8 +707,8 @@ def build_mla_decode_task_manager( work_queue=work_queue, domain=domain, num_registers=wg2_reg_count, - warp_idx=10, - num_warps=2, + warp_idx=9 if cfg.load_num_warps > 1 else 10, + num_warps=3 if cfg.load_num_warps > 1 else 2, ) task_list.append(padding_task) @@ -664,6 +767,12 @@ def build_mla_decode_task_manager( else: resource_dependency_graph[work_queue] = [] + if sparse_offsets_enabled: + # K/V TMA reads register snapshots; metadata has no asynchronous SMEM read edge. + resource_dependency_graph[page_offset_window] = ( + [work_queue] if work_queue is not None else [] + ) + # ────────────────────────────────────────────────────────────── # Create TaskManager # ────────────────────────────────────────────────────────────── @@ -804,6 +913,19 @@ def __init__( raise ValueError("batch_size must be positive") self.batch_size = batch_size self.mask_type = normalize_mask_type(mask_type) + self.fuse_sparse_reduction = False + self.fuse_sparse_epilogue = False + self.output_quant = None + # Sparse attention binds separate source lists; dense attention uses + # the ordinary page table. Both reuse the staged KV loaders. + self.direct_sparse = False + self.sparse_gather_warps = 1 + self.sparse_kv_stages = 0 + self.sparse_uniform_pages = False + self.sparse_balanced_registers = False + self.assume_valid_prefix = False + self.direct_static_scales = False + self.direct_sparse_capacities = (0, 0) self.query_tile_layout = FlatQueryTileLayout.for_tile( num_heads, seq_len_q, mma_qk_tiler_mn[0] ) @@ -827,6 +949,13 @@ def _effective_reduction_shape(self) -> tuple[int, int]: return self.mma_qk_tiler_mn[0], self.num_q_tiles + def _reduction_topology(self): + if self.fuse_sparse_reduction: + return make_balanced_parallel_reduction_topology( + self.reduction_split_capacity, cluster_size=1 + ) + return self.parallel_reduction_topology + def compile_signature(self) -> tuple[object, ...]: """Return the complete batch-independent JIT identity.""" @@ -854,7 +983,18 @@ def compile_signature(self) -> tuple[object, ...]: self.tail_q_rows, self._parallel_reduction_shape_is_eligible, self.use_parallel_reduction, - self.parallel_reduction_topology, + self._reduction_topology(), + self.fuse_sparse_reduction, + self.fuse_sparse_epilogue, + self.output_quant, + self.direct_sparse, + self.assume_valid_prefix, + self.direct_static_scales, + self.direct_sparse_capacities, + self.sparse_gather_warps, + self.sparse_kv_stages, + self.sparse_uniform_pages, + self.sparse_balanced_registers, ) def _configure_parallel_reduction_topology(self) -> None: @@ -915,6 +1055,44 @@ def _configure_parallel_reduction_topology(self) -> None: self.parallel_reduction_topology = topology self.use_parallel_reduction = True + def _make_config(self): + """Share geometry and resource settings between launch and device code.""" + cfg = make_mla_decode_config( + mma_qk_tiler_mn=self.mma_qk_tiler_mn, + mma_pv_tiler_mn=self.mma_pv_tiler_mn, + rope_dim=self.rope_dim, + page_size=self.page_size, + qkv_dtype=self.qkv_dtype, + o_dtype=self.out_dtype, + max_active_clusters=self.max_active_clusters, + is_persistent=self.is_persistent, + is_var_seq=self.is_var_seq, + is_var_split_kv=self.is_var_split_kv, + mask_type=self.mask_type, + sparse_gather_warps=self.sparse_gather_warps, + sparse_kv_stages=self.sparse_kv_stages, + cache_uniform_sparse_pages=self.sparse_uniform_pages, + balance_sparse_registers=self.sparse_balanced_registers, + fuse_output_quant=self.output_quant is not None + and self.fuse_sparse_epilogue, + ) + + # One flattened sparse query supplies a complete M128 tile. Its two + # CTAs each stage 64 heads without row predicates. Warp 10 is otherwise + # padding in the sparse FP8 schedule; using it for stores preserves + # gather/MMA/softmax resources and overlaps output traffic with work. + # Split reductions and other geometries retain direct vector stores. + cfg.use_tma_output = bool( + self.output_quant is not None + and self.fuse_sparse_epilogue + and cfg.is_fp8_qkv() + and cfg.page_size == 1 + and self.num_heads == 128 + and self.seq_len_q == 1 + and cfg.mma_pv_tiler_mn == (128, 256) + ) + return cfg + @cute.jit def __call__( self, @@ -929,25 +1107,40 @@ def __call__( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, stream: object, + output_quant_args=None, ): """Execute the MLA decode TS kernel.""" - cfg = make_mla_decode_config( - mma_qk_tiler_mn=self.mma_qk_tiler_mn, - mma_pv_tiler_mn=self.mma_pv_tiler_mn, - rope_dim=self.rope_dim, - page_size=self.page_size, - qkv_dtype=self.qkv_dtype, - o_dtype=self.out_dtype, - max_active_clusters=self.max_active_clusters, - is_persistent=self.is_persistent, - is_var_seq=self.is_var_seq, - is_var_split_kv=self.is_var_split_kv, - mask_type=self.mask_type, - ) + if cutlass.const_expr(self.fuse_sparse_reduction): + if cutlass.const_expr( + self.is_persistent + or self.is_var_split_kv + or self.static_split_kv is None + or self.static_split_kv < 2 + ): + raise ValueError( + "fused sparse reduction requires fixed, nonpersistent split-KV" + ) + if cutlass.const_expr(self.fuse_sparse_epilogue): + # This CTA can finalize only an unsplit output. Split partials + # must first merge their statistics; this is not a gather4 limit. + if cutlass.const_expr(self.static_split_kv != 1): + raise ValueError("2CTA fused epilogue requires one split") + if cutlass.const_expr(self.direct_sparse): + if cutlass.const_expr( + not (self.fuse_sparse_reduction or self.fuse_sparse_epilogue) + or self.qkv_dtype not in ("e4m3", "bf16") + or self.page_size != 1 + or self.rope_dim != 0 + or self.seq_len_q != 1 + ): + raise ValueError( + "direct sparse 2CTA requires native FP8/BF16 D512 and fused output" + ) + cfg = self._make_config() physical_tile_rows = self.mma_qk_tiler_mn[0] num_query_tiles = self.num_q_tiles # Fixed public tensors retain [H,D,SQ,B]/[H,SQ,B]; variable-Q tensors @@ -1084,48 +1277,127 @@ def _flatten_query_rows_for_tma(t): c_latent.iterator, cute.select(c_latent.layout, mode=[1, 0, 2]), ) - tma_desc_c_latent = create_tensor_map_tiled_from_view( - c_latent_tma, - box_dims=(cfg.mma_qk_tiler[2], cfg.kc_page_tile_size, 1), - stride_order=(0, 1, 2), - swizzle=cuda.TensorMapSwizzle.s128b, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) - - if cutlass.const_expr(cfg.rope_dim > 0): - c_rope_tma = cute.make_tensor( + if cutlass.const_expr(cfg.page_size == 1): + # D512 has no separate Q/K RoPE maps. Reuse their ABI slots for + # extra-K and extra-V descriptors, keeping dense launch ABI stable. + primary_rows = cute.make_tensor( + c_latent.iterator, + cute.make_layout( + (c_latent.shape[1], c_latent.shape[2]), + stride=(c_latent.stride[1], c_latent.stride[2]), + ), + ) + extra_rows = cute.make_tensor( c_rope.iterator, - cute.select(c_rope.layout, mode=[1, 0, 2]), + cute.make_layout( + (c_rope.shape[1], c_rope.shape[2]), + stride=(c_rope.stride[1], c_rope.stride[2]), + ), + ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + primary_rows, + box_dims=(cfg.mma_qk_tiler[2], 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - c_rope_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): - c_rope_swizzle = cuda.TensorMapSwizzle.s64b tma_desc_c_rope = create_tensor_map_tiled_from_view( - c_rope_tma, - box_dims=(cfg.mma_qk_rope_tiler[2], cfg.kc_page_tile_size, 1), - stride_order=(0, 1, 2), - swizzle=c_rope_swizzle, + extra_rows, + box_dims=(cfg.mma_qk_tiler[2], 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + v_swizzle = cuda.TensorMapSwizzle.s128b + sparse_v_width = 128 if cfg.is_fp8_qkv() else V_TMA_LATENT_ELEMENTS + tma_desc_c_transpose = create_tensor_map_tiled_from_view( + primary_rows, + box_dims=(sparse_v_width, 1), + stride_order=(0, 1), + swizzle=v_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + tma_desc_q_rope = create_tensor_map_tiled_from_view( + extra_rows, + box_dims=(sparse_v_width, 1), + stride_order=(0, 1), + swizzle=v_swizzle, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) else: - tma_desc_c_rope = tma_desc_c_latent - - c_transpose_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.mma_pv_tiler[2] == 64): - c_transpose_swizzle = cuda.TensorMapSwizzle.s64b - c_latent_transpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) - c_latent_transpose = cute.make_tensor( - c_latent.iterator, c_latent_transpose_layout - ) - # The physical page controls only the GMEM coordinate. TMA assembles - # fixed K32 SMEM blocks for PV from one or more page-bounded copies. - tma_desc_c_transpose = create_tensor_map_tiled_from_view( - c_latent_transpose, - box_dims=(V_TMA_LATENT_ELEMENTS, cfg.v_tma_token_count, 1), - stride_order=(0, 1, 2), - swizzle=c_transpose_swizzle, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_latent_tma, + box_dims=(cfg.mma_qk_tiler[2], cfg.kc_page_tile_size, 1), + stride_order=(0, 1, 2), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + if cutlass.const_expr(cfg.rope_dim > 0): + c_rope_tma = cute.make_tensor( + c_rope.iterator, + cute.select(c_rope.layout, mode=[1, 0, 2]), + ) + c_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + c_rope_swizzle = cuda.TensorMapSwizzle.s64b + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_rope_tma, + box_dims=(cfg.mma_qk_rope_tiler[2], cfg.kc_page_tile_size, 1), + stride_order=(0, 1, 2), + swizzle=c_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_c_rope = tma_desc_c_latent + + c_transpose_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.mma_pv_tiler[2] == 64): + c_transpose_swizzle = cuda.TensorMapSwizzle.s64b + c_latent_transpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor( + c_latent.iterator, c_latent_transpose_layout + ) + # The physical page controls only the GMEM coordinate. TMA assembles + # fixed K32 SMEM blocks for PV from one or more page-bounded copies. + tma_desc_c_transpose = create_tensor_map_tiled_from_view( + c_latent_transpose, + box_dims=(V_TMA_LATENT_ELEMENTS, cfg.v_tma_token_count, 1), + stride_order=(0, 1, 2), + swizzle=c_transpose_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + tma_desc_o = tma_desc_q_latent + if cutlass.const_expr(cfg.use_tma_output): + flat_o = output_quant_args[0] + output_rows = cute.size(flat_o) // (128 * 512) + # Group-major [G, T, (H/G) * 512] output viewed as (channel, head in + # group, token, group). A box spans whole groups, or half of the + # single H128 group, retaining the same 64 shared-memory rows. + heads_per_group = ( + self.output_quant.num_heads // self.output_quant.num_groups + ) + box_heads = cutlass.const_expr(min(heads_per_group, 64)) + output_view = cute.make_tensor( + flat_o.iterator, + cute.make_layout( + (512, heads_per_group, output_rows, self.output_quant.num_groups), + stride=( + 1, + 512, + heads_per_group * 512, + output_rows * heads_per_group * 512, + ), + ), + ) + tma_desc_o = create_tensor_map_tiled_from_view( + output_view, + box_dims=(128, box_heads, 1, 64 // box_heads), + stride_order=(0, 1, 2, 3), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) softmax_scale_log2 = softmax_scale * LOG2_E @@ -1173,9 +1445,9 @@ def _flatten_query_rows_for_tma(t): kernel_split_kv, workspace, ) - # A one-wave producer can publish its dependent reducer launch while - # retiring, hiding launch latency without admitting reducer CTAs into - # an actively grid-striding persistent producer. + # Nonpersistent producers signal near retirement. The native reducer + # can overlap setup, then waits for the complete prerequisite grid + # before reading partials. Persistent schedules retain stream ordering. use_one_wave_reducer_pdl = acc_o is not None and not self.is_persistent self.split_kv_kernel( @@ -1199,6 +1471,8 @@ def _flatten_query_rows_for_tma(t): output_scale, tile_sched_params, clc_tile_sched_params, + output_quant_args, + tma_desc_o, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1210,8 +1484,10 @@ def _flatten_query_rows_for_tma(t): # Reduction kernel: combine per-split results when split_kv > 1 if cutlass.const_expr(acc_o is not None): - if cutlass.const_expr(self.use_parallel_reduction): - topology = self.parallel_reduction_topology + if cutlass.const_expr( + self.use_parallel_reduction or self.fuse_sparse_reduction + ): + topology = self._reduction_topology() self.parallel_reduction_kernel( o, lse, @@ -1221,14 +1497,22 @@ def _flatten_query_rows_for_tma(t): cache_seqs, cu_seqlens_q, block_split_kvs, + output_quant_args, ).launch( grid=( - physical_tile_rows * topology.cluster_size, + ( + min(self.num_heads, physical_tile_rows) + if self.fuse_sparse_reduction + else physical_tile_rows + ) + * topology.cluster_size, num_query_tiles, batch_size, ), block=[PARALLEL_REDUCTION_THREADS, 1, 1], - cluster=[topology.cluster_size, 1, 1], + cluster=None + if topology.cluster_size == 1 + else [topology.cluster_size, 1, 1], stream=stream, min_blocks_per_mp=1, use_pdl=use_one_wave_reducer_pdl, @@ -1280,26 +1564,52 @@ def split_kv_kernel( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, tile_sched_params: MLAStaticTileSchedulerParams, clc_tile_sched_params: object, + output_quant_args=None, + tma_desc_o: cutlass.GridConstant[cuda.TensorMap] = None, ) -> None: """MLA decode TS kernel: persistent tile-scheduled execution.""" - cfg = make_mla_decode_config( - mma_qk_tiler_mn=self.mma_qk_tiler_mn, - mma_pv_tiler_mn=self.mma_pv_tiler_mn, - rope_dim=self.rope_dim, - page_size=self.page_size, - qkv_dtype=self.qkv_dtype, - o_dtype=self.out_dtype, - max_active_clusters=self.max_active_clusters, - is_persistent=self.is_persistent, - is_var_seq=self.is_var_seq, - is_var_split_kv=self.is_var_split_kv, - mask_type=self.mask_type, - ) + if cutlass.const_expr(self.direct_sparse): + from ..helpers.sparse_indices import bind_sparse_kv_indices + + si, ci, sl, cl, sm, os, sinks = block_split_kvs + capacity = cutlass.const_expr( + max( + 256, + ( + (self.direct_sparse_capacities[0] + 127) // 128 + + (self.direct_sparse_capacities[1] + 127) // 128 + ) + * 128, + ) + ) + request = None + if cutlass.const_expr(not self.is_persistent): + _, row = divmod_constexpr_power_of_two_or_fdd( + cute.arch.block_idx()[1], + None, + tile_sched_params.problem_shape_b_fdd, + ) + request = Int64(row) + page_offsets, cache_seqs = bind_sparse_kv_indices( + si, + ci, + sl, + cl, + capacity, + query_row=request, + assume_valid_prefix=self.assume_valid_prefix, + has_extra=self.direct_sparse_capacities[1] > 0, + ) + if cutlass.const_expr(not self.direct_static_scales): + # Caller scales already include Q/K/V descales. + softmax_scale_log2 = Float32(sm[0]) * LOG2_E + output_scale = Float32(os[0]) + cfg = self._make_config() num_query_tiles = self.num_q_tiles use_clc_dynamic = self.is_persistent and not cfg.is_fp8_qkv() tiled_mma_qk = None @@ -1362,6 +1672,26 @@ def split_kv_kernel( space=cutlass.AddressSpace.smem, alignment=1024, ) + # One shared metadata ring feeds both K/V issuers. Each holds its + # stage through the last coordinate read; TMA consumes register + # operands and its payload buffers have independent lifetimes. + sparse_offsets_smem = None + if cutlass.const_expr(cfg.page_size == 1 and cfg.use_fp8_split_mma_schedule): + sparse_offsets_smem = cutlass.Array( + Int32, + cfg.sparse_offset_stages * cfg.tokens_per_k_tile, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + smem_o = None + if cutlass.const_expr(cfg.use_tma_output): + smem_o = cutlass.Array( + cutlass.Uint8, + OUTPUT_TMA_STAGES * OUTPUT_TMA_STAGE_BYTES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + softmax_exchange_arr = cutlass.Array( self.acc_dtype, cfg.softmax_exchange_elems, @@ -1452,6 +1782,26 @@ def split_kv_kernel( ) blk_coord = (cluster_idx, seq_q_idx, batch_idx, split_kv_idx) tile_cluster_idx, tile_seq_q_idx, tile_batch_idx, tile_split_kv_idx = blk_coord + if cutlass.const_expr(self.fuse_sparse_reduction): + # Final output uses a fixed-S reducer: neutral LSE slots let it + # avoid reloading source lengths and dividing split ranges in + # every output CTA. Active tiles publish their own statistics. + if tidx < Int32(cfg.mma_qk_tiler[0] // cfg.num_mma_ctas): + total_tiles = (Int32(cache_seqs[tile_batch_idx]) + Int32(127)) // Int32( + 128 + ) + _, local_tiles = runtime_split_tile_range( + total_tiles, split_kv, tile_split_kv_idx + ) + if local_tiles == Int32(0): + row_in_tile = ( + tile_cluster_idx + * Int32(cfg.mma_qk_tiler[0] // cfg.num_mma_ctas) + + tidx + ) + acc_lse[ + row_in_tile, tile_split_kv_idx, tile_seq_q_idx, tile_batch_idx + ] = Float32(-Float32.inf) del tile_cluster_idx fixed_nonempty_single_split = ( @@ -1636,6 +1986,7 @@ def split_kv_kernel( smem_kc_arr=smem_kc_arr, smem_vc_arr=smem_vc_arr, smem_p_arr=smem_p_arr, + sparse_offsets_smem=sparse_offsets_smem, tma_desc_q_latent=tma_desc_q_latent.get_ptr(), tma_desc_q_rope=tma_desc_q_rope.get_ptr(), tma_desc_c_latent=tma_desc_c_latent.get_ptr(), @@ -1645,6 +1996,12 @@ def split_kv_kernel( blk_coord=blk_coord, tidx=tidx, output=o, + atten_sinks=sinks if self.fuse_sparse_epilogue else None, + lse_in_natural_log=self.fuse_sparse_epilogue, + output_quant=self.output_quant if self.fuse_sparse_epilogue else None, + output_quant_args=output_quant_args, + smem_o=smem_o, + tma_o_desc=tma_desc_o.get_ptr(), acc_output=acc_o, lse=lse, acc_lse=acc_lse, @@ -1838,7 +2195,8 @@ def parallel_reduction_kernel( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, + output_quant_args=None, ): """Dispatch the high-split fixed-D512 cluster reducer.""" @@ -1853,7 +2211,11 @@ def parallel_reduction_kernel( ) if cutlass.const_expr(not self.is_persistent): prims.griddepcontrol(kind=prims.GridDepAction.WAIT) - topology = self.parallel_reduction_topology + atten_sinks = None + if cutlass.const_expr(self.fuse_sparse_reduction): + atten_sinks = block_split_kvs[-1] + block_split_kvs = None + topology = self._reduction_topology() run_parallel_reduction_kernel( self.num_heads, self.seq_len_q, @@ -1870,4 +2232,8 @@ def parallel_reduction_kernel( topology.actual_splits, topology.cluster_size, topology.slots_per_rank, + atten_sinks, + self.fuse_sparse_reduction, + self.output_quant, + output_quant_args, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py index 2512ef3dcfa1..41b61707a4c4 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py @@ -19,8 +19,12 @@ from cutlass import Float32, Int32, Int64 from cutlass.experimental import primitives as prims -from ...separate_reduction import finalize_log2_sum_exp, normalized_lse_weight -from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ...separate_reduction import ( + finalize_log2_sum_exp, + normalized_lse_weight, + attention_sink_scale, +) +from ..helpers.constants import LN_2, SPLIT_REDUCTION_SCALE_BARRIER_ID from ..helpers.mask import MaskType, mask_visible_k_length from ..helpers.math import ceil_div from ..helpers.ops import fmax_f32, warp_reduce_max_f32, warp_reduce_sum_f32 @@ -147,15 +151,30 @@ def _store_parallel_reduction_result( tidx, batch_idx, cu_seqlens_q, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Publish one normalized FP32 output fragment and final LSE.""" + if cutlass.const_expr(finalize_attention): + global_lse *= Float32(LN_2) if tidx == Int32(0) and query_is_valid: if cutlass.const_expr(cu_seqlens_q is not None): lse[logical_head_idx, storage_q_idx] = global_lse else: lse[logical_head_idx, logical_q_idx, batch_idx] = global_lse + if cutlass.const_expr(output_quant is not None): + if query_is_valid: + row = Int64(storage_q_idx) * output_quant.num_heads + logical_head_idx + if cutlass.const_expr(cu_seqlens_q is None): + row = ( + Int64(batch_idx) * output.shape[2] + logical_q_idx + ) * output_quant.num_heads + logical_head_idx + output_quant.store(output_quant_args, output_vals, row, element_idx) + return + out_element_dtype = output.element_type output_regs = cutlass.Array( out_element_dtype, @@ -193,6 +212,10 @@ def run_parallel_reduction_kernel( actual_splits: cutlass.Constexpr[int], cluster_size: cutlass.Constexpr[int], slots_per_rank: cutlass.Constexpr[int], + atten_sinks=None, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Reduce one D=512 row cooperatively across a padded CTA cluster. @@ -206,32 +229,46 @@ def run_parallel_reduction_kernel( signal before entering this body. """ + if cutlass.const_expr(finalize_attention and cluster_size != 1): + raise ValueError("final attention output uses a single-CTA reducer per row") block_idx_x, query_tile_idx, batch_idx = cute.arch.block_idx() tidx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) lane_idx = tidx % Int32(cfg.threads_per_warp) - cluster_rank = cute.arch.block_idx_in_cluster() + cluster_rank = Int32(0) if cluster_size == 1 else cute.arch.block_idx_in_cluster() row_in_tile = block_idx_x // Int32(cluster_size) - ( - logical_head_idx, - logical_q_idx, - storage_q_idx, - query_is_valid, - active_split_kv, - ) = _parallel_reduction_row_state( - cfg.mma_qk_tiler[0], - num_heads, - seq_len_q, - is_var_split_kv, - cache_seqs, - cu_seqlens_q, - block_split_kvs, - split_kv, - row_in_tile, - query_tile_idx, - batch_idx, - cfg, - ) + if cutlass.const_expr(finalize_attention): + # A sparse query is already one virtual request. The producer writes + # neutral LSE for skipped splits, so the compiled slot count suffices; + # live length loads and dynamic split division would repeat per head. + if cutlass.const_expr(seq_len_q != 1 or cu_seqlens_q is not None): + raise ValueError("final sparse reduction requires virtual query rows") + logical_head_idx = query_tile_idx * Int32(cfg.mma_qk_tiler[0]) + row_in_tile + logical_q_idx = Int32(0) + storage_q_idx = batch_idx + query_is_valid = logical_head_idx < Int32(num_heads) + active_split_kv = Int32(actual_splits) + else: + ( + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + active_split_kv, + ) = _parallel_reduction_row_state( + cfg.mma_qk_tiler[0], + num_heads, + seq_len_q, + is_var_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_kv, + row_in_tile, + query_tile_idx, + batch_idx, + cfg, + ) element_idx = tidx * Int32(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD) neg_inf = Float32(-Float32.inf) @@ -297,16 +334,25 @@ def run_parallel_reduction_kernel( ) local_sum_lse += lane_exp[lane_slot_i] local_sum_lse = warp_reduce_sum_f32(local_sum_lse) - local_lse_value = finalize_log2_sum_exp(local_exp_frame, local_sum_lse) + # Keep this warp-local value distinct from the post-barrier SMEM load. + stats_lse_value = finalize_log2_sum_exp(local_exp_frame, local_sum_lse) + sink_scale = Float32(1) + if cutlass.const_expr(finalize_attention and atten_sinks is not None): + sink_scale = Float32(0) + if query_is_valid: + sink_scale = attention_sink_scale( + stats_lse_value, Float32(atten_sinks[logical_head_idx]) + ) if lane_idx == Int32(0): - smem_local_lse[0] = local_lse_value + smem_local_lse[0] = stats_lse_value for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) if local_slot_idx < Int32(slots_per_rank): - smem_local_scale[local_slot_idx] = normalized_lse_weight( - lane_lse[lane_slot_i], local_lse_value - ) + weight = normalized_lse_weight(lane_lse[lane_slot_i], stats_lse_value) + if cutlass.const_expr(finalize_attention): + weight *= sink_scale + smem_local_scale[local_slot_idx] = weight prims.barrier_cta_sync(SPLIT_REDUCTION_SCALE_BARRIER_ID) @@ -331,23 +377,26 @@ def run_parallel_reduction_kernel( ) if active_slot: split_scale = Float32(smem_local_scale[slot_i]) - partial_offset = Int64( - row_in_tile * acc_output.stride[0] - + split_idx * acc_output.stride[1] - + element_idx * acc_output.stride[2] - + query_tile_idx * acc_output.stride[3] - + batch_idx * acc_output.stride[4] - ) - partial_output = ( - (acc_output_ptr + partial_offset) - .load( - count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, - alignment=8, + if not finalize_attention or split_scale > Float32(0): + partial_offset = Int64( + row_in_tile * acc_output.stride[0] + + split_idx * acc_output.stride[1] + + element_idx * acc_output.stride[2] + + query_tile_idx * acc_output.stride[3] + + batch_idx * acc_output.stride[4] ) - .to(Float32) - ) - for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): - local_output[j] += partial_output[j] * split_scale + partial_output = ( + (acc_output_ptr + partial_offset) + .load( + count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + alignment=8, + ) + .to(Float32) + ) + for j in cutlass.range_constexpr( + PARALLEL_REDUCTION_ELEMENTS_PER_THREAD + ): + local_output[j] += partial_output[j] * split_scale local_lse_value = smem_local_lse[0] @@ -367,6 +416,9 @@ def run_parallel_reduction_kernel( tidx, batch_idx, cu_seqlens_q, + finalize_attention, + output_quant, + output_quant_args, ) return diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py index 55523137fdea..4e3d24d11fd6 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py @@ -32,14 +32,16 @@ - TmemCorrResource : Async(2 stages), SoftmaxTask -> CorrectionTask - TmemOResource : UmmaProducerAsync(1 stage), MmaTask -> CorrectionTask -GMEM (no pipeline) ------------------- -- GmemOResource : No pipeline, Correction -> GMEM +Output +------ +- GmemOResource : Direct stores, or Async(2 stages) Correction -> TMA -> GMEM """ from dataclasses import dataclass, field from typing import Any, ClassVar +from ..helpers.constants import LN_2, LOG2_E + import cutlass import cutlass.cute as cute from cutlass.experimental import primitives as prims @@ -56,7 +58,19 @@ from ...mask import kv_tile_needs_right_mask from ...tensor_map import transform_ragged_coords +from ..helpers.sparse_indices import SparseKVIndices, SparseSourceTile +from ..helpers.gather import ( + gather4_cached, + load_sparse_rows, + invalid_sparse_token, + gather4_uniform_quad_slices, + load_sparse_quad, + cache_sparse_mask, + CachedSparseMask, +) + from .config import ( + OUTPUT_TMA_STAGE_BYTES, V_SMEM_K_BLOCK_TOKENS, V_TMA_LATENT_ELEMENTS, MlaDecodeConfig, @@ -72,6 +86,8 @@ TCGEN05_32B_REGS_PER_LOAD, TCGEN05_32B_SHAPE, WARP_LANE_SHIFT, + WARP_LANES, + TMA_GATHER_ROWS, ) from ..helpers.tile_scheduler import ( MLAStaticTileScheduler, @@ -623,10 +639,19 @@ def advance_tile(self, stage_info: StageInfo): @dataclass(kw_only=True) class PageOffsetWindowResource(HighThroughputMlaResource): - """GMEM page table cached as one register per lane of the TMA warp.""" + """Page-coordinate staging for paged and sparse KV loads. + + FP8 sparse loaders read uniform quads from the prefetched shared ring. + Other schedules retain page windows in registers. + """ + request_lengths: Any = field(init=False, default=None) page_offsets: Any = None # GMEM page-offset tensor + smem_sparse_offsets: Any = None cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cached_page_stage: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) cached_k_pages: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -639,7 +664,21 @@ class PageOffsetWindowResource(HighThroughputMlaResource): cached_window_page: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) + cached_mask_low: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + cached_mask_high: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("cached_page_stage", Int32, Int32(0), "Held shared-offset stage."), + ("cached_mask_low", Int32, Int32(0), "Low half of the prefetched sparse mask."), + ( + "cached_mask_high", + Int32, + Int32(0), + "High half of the prefetched sparse mask.", + ), ( "cached_k_pages", cutlass.Array, @@ -666,6 +705,76 @@ class PageOffsetWindowResource(HighThroughputMlaResource): ), ) + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_sparse_state(self, stage_info: StageInfo): + del stage_info + # Allocate outside skipped-work guards; refresh values for every tile. + self.request_lengths = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def bind_sparse_request(self, stage_info: StageInfo): + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): + row = stage_info.work_tile.tile_idx[2] + self.request_lengths[0] = Int32(self.page_offsets.primary_lengths[row]) + self.request_lengths[1] = Int32(self.page_offsets.extra_lengths[row]) + + @producer_work + @cute.jit + def load_sparse_offsets(self, stage_info: StageInfo): + lane = cute.arch.thread_idx()[0] % 32 + k_index = stage_info.work_tile.k_index_base + Int32(stage_info.loop_offset) + token = k_index * Int32(128) + lane * Int32(4) + row = stage_info.work_tile.tile_idx[2] + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): + routes = self.page_offsets.bind_query_row( + row, self.request_lengths[0], self.request_lengths[1] + ) + else: + routes = self.page_offsets + quad = load_sparse_quad(routes, row, token) + self.smem_sparse_offsets.data_ptr(stage_info.stage_idx * 128 + lane * 4).store( + quad.load(0, 4), alignment=16 + ) + cute.arch.fence_view_async_shared() + + @consumer_work(returns=(cached_mask_low, cached_mask_high)) + @cute.jit + def read_sparse_mask(self, stage_info: StageInfo): + lane = cute.arch.thread_idx()[0] % 32 + half = (cute.arch.thread_idx()[0] % 128) // 64 + base = stage_info.stage_idx * 128 + half * 64 + tile = stage_info.work_tile.k_index_base + Int32(stage_info.loop_offset) + token = tile * Int32(128) + half * Int32(64) + lane + length = Int32(stage_info.work_tile.k_len) + low = self.smem_sparse_offsets[base + lane] + high = self.smem_sparse_offsets[base + lane + 32] + low_bits = Int32( + cute.arch.vote_ballot_sync( + ((low & Int32(0x7FFFFFFF)) == Int32(0x7FFFFFFF)) | (token >= length) + ) + ) + high_bits = Int32( + cute.arch.vote_ballot_sync( + ((high & Int32(0x7FFFFFFF)) == Int32(0x7FFFFFFF)) + | (token + Int32(32) >= length) + ) + ) + return low_bits, high_bits + + @consumer_work(returns=cached_page_stage) + @cute.jit + def sparse_stage(self, stage_info: StageInfo): + """Return the shared-ring slot protected by the current consumer token.""" + return stage_info.stage_idx + @consumer_work( work_attrs=WorkAttr.AUXILIARY, returns=( @@ -680,20 +789,37 @@ def init_read_state(self, stage_info: StageInfo): """Create cached page-index arrays used by staged KV TMA loads.""" del stage_info cfg = self.cfg + # Uniform caches repeat each warp's K/current-V coordinates in all + # lanes, saving shuffles on later feature slices. Keep pending V in + # compact lane fragments until needed to limit live registers. return ( cutlass.Array( Int32, - cfg.pages_per_k_cta, + ( + cfg.tokens_per_k_cta // cfg.load_num_warps + if cfg.cache_uniform_sparse_pages + else cfg.tokens_per_k_cta // WARP_LANES + ) + if cfg.page_size == 1 + else cfg.pages_per_k_cta, space=cutlass.AddressSpace.rmem, ), cutlass.Array( Int32, - cfg.pages_per_v_tile, + ( + cfg.tokens_per_v_tile // cfg.load_num_warps + if cfg.cache_uniform_sparse_pages + else cfg.tokens_per_v_tile // WARP_LANES + ) + if cfg.page_size == 1 + else cfg.pages_per_v_tile, space=cutlass.AddressSpace.rmem, ), cutlass.Array( Int32, - cfg.pages_per_v_tile, + cfg.tokens_per_v_tile // WARP_LANES + if cfg.page_size == 1 + else cfg.pages_per_v_tile, space=cutlass.AddressSpace.rmem, ), Int32(0), @@ -731,6 +857,108 @@ def read_page_offset_window( blk_coord = work_tile.tile_idx local_k_index = Int32(stage_info.loop_offset) global_k_index = work_tile.k_index_base + local_k_index + if cutlass.const_expr(cfg.page_size == 1): + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + # K and pending V cover the same source-aligned BK128 tile. + # Bind its source once, not again for every lane/fragment. + routes = ( + self.page_offsets.select_source_tile( + global_k_index * cfg.tokens_per_k_tile, blk_coord[2] + ) + if cutlass.const_expr(isinstance(self.page_offsets, SparseKVIndices)) + else self.page_offsets + ) + if cutlass.const_expr(isinstance(routes, SparseSourceTile)): + new_v = load_sparse_rows( + routes, + blk_coord[2], + global_k_index * cfg.tokens_per_v_tile, + fragments=cfg.tokens_per_v_tile // WARP_LANES, + ) + # K is one CTA half of this freshly loaded pending-V window. + # cached_v_pages instead belongs to the previous K tile. + fragments_k = cfg.tokens_per_k_cta // WARP_LANES + new_k = cutlass.Array( + Int32, fragments_k, space=cutlass.AddressSpace.rmem + ) + for fragment in cutlass.range_constexpr(fragments_k): + new_k[fragment] = ( + new_v[fragment] + if cta == Int32(0) + else new_v[fragment + fragments_k] + ) + else: + new_k = load_sparse_rows( + routes, + blk_coord[2], + global_k_index * cfg.tokens_per_k_tile + cta * cfg.tokens_per_k_cta, + fragments=cfg.tokens_per_k_cta // WARP_LANES, + ) + new_v = load_sparse_rows( + routes, + blk_coord[2], + global_k_index * cfg.tokens_per_v_tile, + fragments=cfg.tokens_per_v_tile // WARP_LANES, + ) + # Expand only current K/V coordinates once per tile. "Uniform" + # means identical across lanes, not consecutive physical KV rows. + # Reuse across D slices trades register capacity for fewer shuffles. + if cutlass.const_expr(cfg.cache_uniform_sparse_pages): + warp = ( + cute.arch.make_warp_uniform(cute.arch.warp_idx()) + - cfg.load_tma_warp_id + ) + fragment = ( + new_k[0] if warp < Int32(cfg.load_num_warps // 2) else new_k[1] + ) + for i in cutlass.range_constexpr(64 // cfg.load_num_warps): + cached_k_pages[i] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=fragment, + offset=(warp % Int32(cfg.load_num_warps // 2)) + * Int32(64 // cfg.load_num_warps) + + Int32(i), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + else: + for i in cutlass.range_constexpr(2): + cached_k_pages[i] = new_k[i] + if cutlass.const_expr(cfg.cache_uniform_sparse_pages): + # Expand only the delayed V tile; retain the next V tile in + # lane fragments to keep the register footprint bounded. + for i in cutlass.range_constexpr(128 // cfg.load_num_warps): + if cutlass.const_expr(init_v_cache): + cached_v_pages[i] = Int32(0x7FFFFFFF) + else: + cached_v_pages[i] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_next_v_pages[ + i // (32 // cfg.load_num_warps) + ], + offset=warp * Int32(32 // cfg.load_num_warps) + + Int32(i % (32 // cfg.load_num_warps)), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + for i in cutlass.range_constexpr(4): + cached_next_v_pages[i] = new_v[i] + else: + for i in cutlass.range_constexpr(4): + cached_v_pages[i] = ( + Int32(0x7FFFFFFF) if init_v_cache else cached_next_v_pages[i] + ) + cached_next_v_pages[i] = new_v[i] + return ( + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + cached_window_page, + ) pages_per_k_tile = cutlass.const_expr(cfg.pages_per_k_tile) page_window_tiles = cutlass.const_expr(32 // pages_per_k_tile) page_window_mask = cutlass.const_expr(page_window_tiles - 1) @@ -864,6 +1092,15 @@ def _query_tma_coords( @producer_work @cute.jit def tma_load(self, stage_info: StageInfo) -> None: + """Let one warp issue Q while the full loader group signals readiness.""" + if cutlass.const_expr(self.cfg.load_num_warps > 1): + if cute.arch.thread_idx()[0] // 32 == self.cfg.load_k_warp_id: + self._tma_load(stage_info) + else: + self._tma_load(stage_info) + + @cute.jit + def _tma_load(self, stage_info: StageInfo) -> None: """TMA load Q latent + Q rope into SMEM (all sub-tiles in one commit).""" cfg = self.cfg work_tile = stage_info.work_tile @@ -963,6 +1200,7 @@ class SmemKVResource(HighThroughputMlaResource): tma_desc_c_latent: Any = None tma_desc_c_rope: Any = None tma_desc_c_transpose: Any = None + tma_desc_extra_v: Any = None cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) cta_rank: Any = field(init=False, default=None) is_leader: Any = field(init=False, default=None) @@ -1002,6 +1240,162 @@ def tma_load( cfg = self.cfg stage_idx = stage_info.stage_idx + if cutlass.const_expr(cfg.page_size == 1): + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + base = stage_idx * cfg.smem_k_stage_elems + warp = ( + cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_tma_warp_id + ) + # Each issuer owns disjoint four-row quads. K covers 64 rows per + # CTA (16 quads); V copies 32-token subtiles (8 quads). BF16 V is + # blocked as 32x64 panels: two 64-feature/128-byte slices cover + # each CTA's 128 features within a 256-feature PV panel. Source + # columns are elements; gather4 destination strides are bytes. + if cutlass.const_expr(cfg.load_num_warps in (4, 8)): + if cutlass.const_expr(not is_v): + if cutlass.const_expr(not cfg.cache_uniform_sparse_pages): + fragment = ( + cached_k_pages[0] + if warp < Int32(cfg.load_num_warps // 2) + else cached_k_pages[1] + ) + for owner in cutlass.range_constexpr(16 // cfg.load_num_warps): + raw = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + if cutlass.const_expr(cfg.cache_uniform_sparse_pages): + raw[j] = cached_k_pages[owner * 4 + j] + else: + lane_base = ( + warp % Int32(cfg.load_num_warps // 2) + ) * Int32(64 // cfg.load_num_warps) + Int32(owner * 4) + raw[j] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=fragment, + offset=lane_base + Int32(j), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + group = warp * Int32(16 // cfg.load_num_warps) + Int32(owner) + gather4_uniform_quad_slices( + self.smem_kv.data_ptr( + base + group * 4 * cfg.mma_qk_tiler_k + ), + self.tma_desc_c_latent, + self.tma_desc_c_rope, + Int32( + subtile_idx + * cfg.kv_subtiles_per_stage + * cfg.mma_qk_tiler_k + ), + raw, + stage_info.barrier, + num_slices=cfg.kv_subtiles_per_stage, + dst_stride=64 * cfg.mma_qk_tiler_k * cfg.qkv_dtype_bytes, + col_stride=cfg.mma_qk_tiler_k, + ) + else: + # PV consumes the delayed current-V window; the tail can + # instead request the compact pending-V coordinates. + cached = cached_next_v_pages if use_next_v_pages else cached_v_pages + pv_j = subtile_idx // cfg.kv_subtiles_per_stage + token_partition = subtile_idx % cfg.kv_subtiles_per_stage + for part in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + pv_i = token_partition * cfg.kv_subtiles_per_stage + part + for owner in cutlass.range_constexpr(8 // cfg.load_num_warps): + group = warp * Int32(8 // cfg.load_num_warps) + Int32(owner) + raw = cutlass.Array( + Int32, 4, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(4): + if cutlass.const_expr( + cfg.cache_uniform_sparse_pages + and not use_next_v_pages + ): + raw[j] = cached_v_pages[ + pv_i * (32 // cfg.load_num_warps) + + owner * 4 + + j + ] + else: + raw[j] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached[pv_i], + offset=group * Int32(4) + Int32(j), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + gather4_uniform_quad_slices( + self.smem_kv.data_ptr( + base + part * 32 * 128 + group * 4 * 64 + ), + self.tma_desc_c_transpose, + self.tma_desc_extra_v, + cta * Int32(128) + Int32(pv_j * 256), + raw, + stage_info.barrier, + num_slices=2, + dst_stride=32 * 64 * cfg.qkv_dtype_bytes, + col_stride=64, + ) + return + # Fewer issuers reconstruct quads from compact lane fragments. + if cutlass.const_expr(not is_v): + for part in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + column = Int32( + (subtile_idx * cfg.kv_subtiles_per_stage + part) + * cfg.mma_qk_tiler_k + ) + for local_group in cutlass.range( + 16 // cfg.load_num_warps, unroll_full=True + ): + group = local_group + gather4_cached( + self.smem_kv.data_ptr( + base + + part * 64 * cfg.mma_qk_tiler_k + + group * 4 * cfg.mma_qk_tiler_k + ), + self.tma_desc_c_latent, + self.tma_desc_c_rope, + column, + cached_k_pages, + group * 4, + stage_info.barrier, + cta_group=2, + ) + else: + cached = cached_next_v_pages if use_next_v_pages else cached_v_pages + pv_j = subtile_idx // cfg.kv_subtiles_per_stage + token_partition = subtile_idx % cfg.kv_subtiles_per_stage + for part in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + pv_i = token_partition * cfg.kv_subtiles_per_stage + part + for panel in cutlass.range_constexpr(2): + column = cta * 128 + Int32(pv_j * 256 + panel * 64) + for local_group in cutlass.range( + 8 // cfg.load_num_warps, unroll_full=True + ): + group = local_group + gather4_cached( + self.smem_kv.data_ptr( + base + + part * 32 * 128 + + panel * 32 * 64 + + group * 4 * 64 + ), + self.tma_desc_c_transpose, + self.tma_desc_extra_v, + column, + cached, + Int32(pv_i * 32) + group * 4, + stage_info.barrier, + cta_group=2, + ) + return + kc_page_smem_elems = cfg.kc_page_tile_size * cfg.mma_qk_tiler_k kv_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) kv_stage_stride_elems = cutlass.const_expr(cfg.smem_k_stage_elems) @@ -1259,11 +1653,12 @@ def v_desc(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] class SmemKResource(HighThroughputMlaResource): """SMEM K buffer for the FP8 split-MMA path. - Producer: LoadTma. Consumer: MmaQkTask. A single producer stage contains - all latent K sub-tiles plus the RoPE K sub-tile for one logical K tile. + Producer: LoadKTask. Consumer: MmaQkTask. A single producer stage contains + all latent K sub-tiles and optional RoPE channels for one logical K tile. """ smem_k: Any = None + smem_sparse_offsets: Any = None page_offsets: Any = None tma_desc_c_latent: Any = None tma_desc_c_rope: Any = None @@ -1284,8 +1679,42 @@ def init_load_state(self, stage_info: StageInfo) -> None: @producer_work @cute.jit - def tma_load_direct(self, stage_info: StageInfo) -> None: - """TMA load one full K tile using page offsets read directly from GMEM.""" + def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): + """Read one warp-uniform SMEM quad and reuse it across FP8 feature slices.""" + cfg = self.cfg + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_k_warp_id + base = ( + cute.arch.make_warp_uniform(stage_info.stage_idx) * cfg.smem_k_stage_elems + ) + for owner in cutlass.range_constexpr( + cfg.tokens_per_k_cta // TMA_GATHER_ROWS // cfg.load_num_warps + ): + # All lanes read the same quad: shared-memory broadcast replaces + # four register shuffles. Retain the quad across feature slices. + quad = warp + Int32(owner * cfg.load_num_warps) + offset = ( + page_offset_stage * Int32(128) + + quad * Int32(4) + + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + * Int32(64) + ) + raw = self.smem_sparse_offsets.data_ptr(offset).load(count=4, alignment=16) + gather4_uniform_quad_slices( + self.smem_k.data_ptr(base + quad * 4 * cfg.mma_qk_tiler_k), + self.tma_desc_c_latent, + self.tma_desc_c_rope, + Int32(0), + raw, + stage_info.barrier, + num_slices=cfg.iterations_qk_latent, + dst_stride=64 * cfg.mma_qk_tiler_k, + col_stride=cfg.mma_qk_tiler_k, + ) + + @producer_work + @cute.jit + def tma_load_paged(self, stage_info: StageInfo) -> None: + """Load a paged K tile; page-size-one inputs use the offset ring.""" cfg = self.cfg stage_idx = stage_info.stage_idx stage_base = stage_idx * cfg.smem_k_stage_elems @@ -1433,13 +1862,15 @@ def k_desc(self, stage_info: StageInfo, *, k_subtile_idx: cutlass.Constexpr[int] class SmemVResource(HighThroughputMlaResource): """SMEM V buffer for the FP8 split-MMA path. - Producer: LoadTma. Consumer: MmaPvTask. A single stage contains every V + Producer: LoadVTask. Consumer: MmaPvTask. A single stage contains every V sub-tile needed by the PV MMA for one logical K tile. """ smem_v: Any = None + smem_sparse_offsets: Any = None page_offsets: Any = None tma_desc_c_transpose: Any = None + tma_desc_extra_v: Any = None logical_seq_len_q: cutlass.Constexpr[int] = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) desc_v_base: cutlass.Constexpr[TaskLocalVariable] = ( @@ -1457,14 +1888,50 @@ def init_load_state(self, stage_info: StageInfo) -> None: @producer_work @cute.jit - def tma_load_direct(self, stage_info: StageInfo) -> None: - """TMA load one full V tile using page offsets read directly from GMEM.""" + def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): + """Read one warp-uniform SMEM quad and reuse it across FP8 V panels.""" + cfg = self.cfg + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_v_warp_id + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + base = ( + cute.arch.make_warp_uniform(stage_info.stage_idx) * cfg.smem_v_stage_elems + ) + for owner in cutlass.range_constexpr( + cfg.tokens_per_v_tile // TMA_GATHER_ROWS // cfg.load_num_warps + ): + # All lanes read the same quad: shared-memory broadcast replaces + # four register shuffles. Retain the quad across feature slices. + quad = warp + Int32(owner * cfg.load_num_warps) + offset = page_offset_stage * Int32(128) + quad * Int32(4) + raw = self.smem_sparse_offsets.data_ptr(offset).load(count=4, alignment=16) + call = (quad // Int32(16)) * cfg.iterations_pv_n + gather4_uniform_quad_slices( + self.smem_v.data_ptr( + base + + call * 128 * cfg.mma_pv_tiler[2] + + (quad % Int32(16)) * 4 * 128 + ), + self.tma_desc_c_transpose, + self.tma_desc_extra_v, + cta * 128, + raw, + stage_info.barrier, + num_slices=cfg.iterations_pv_n, + dst_stride=128 * cfg.mma_pv_tiler[2], + col_stride=256, + ) + + @producer_work + @cute.jit + def tma_load_paged(self, stage_info: StageInfo) -> None: + """Load a paged V tile; page-size-one inputs use the offset ring.""" cfg = self.cfg stage_idx = stage_info.stage_idx stage_base = stage_idx * cfg.smem_v_stage_elems work_tile = stage_info.work_tile blk_coord = work_tile.tile_idx k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) + page_row_idx = blk_coord[2] page_offsets_batch = self.page_offsets[None, page_row_idx] @@ -1567,6 +2034,10 @@ def v_desc(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] layout = cutlass.const_expr( qk_desc_layout_for_head_dim(cfg, cfg.mma_pv_tiler[2]) ) + if cutlass.const_expr(cfg.page_size == 1): + leading_byte_offset = 0 + stride_byte_offset = 1024 + layout = 2 svc_ptr = self.smem_v.data_ptr(subtile_offset) desc = Int64( prims.Tcgen05SmemDesc.build( @@ -1607,6 +2078,10 @@ def v_desc_n_major( layout = cutlass.const_expr( qk_desc_layout_for_head_dim(cfg, cfg.mma_pv_tiler[2]) ) + if cutlass.const_expr(cfg.page_size == 1): + leading_byte_offset = 0 + stride_byte_offset = 1024 + layout = 2 svc_ptr = self.smem_v.data_ptr(subtile_offset) desc = Int64( prims.Tcgen05SmemDesc.build( @@ -1639,6 +2114,7 @@ class TmemSResource(HighThroughputMlaResource): smem_p: Any = None # SMEM P array smem_exchange: Any = None # SMEM array for cross-warp max exchange softmax_scale_log2: Any = None # softmax_scale * log2(e) + page_offsets: Any = None cache_seqs: Any = None # per-batch valid K length cu_seqlens_q: Any = None # cumulative compact-Q offsets, or None for fixed Q split_kv: Any = None # per-work-tile split count @@ -1650,6 +2126,7 @@ class TmemSResource(HighThroughputMlaResource): is_leader: Any = field(init=False, default=None) row_max_state: Any = field(init=False, default=None) row_sum_state: Any = field(init=False, default=None) + request_lengths: Any = field(init=False, default=None) qk_acc_regs: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -1755,6 +2232,8 @@ def init_softmax_state(self, stage_info: StageInfo): self.is_leader = self.cta_rank == 0 self.row_max_state = Float32(-Float32.inf) self.row_sum_state = Float32(0) + # Allocate outside skipped-work guards; refresh values for every tile. + self.request_lengths = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) return ( cutlass.Array( Float32, @@ -1803,6 +2282,18 @@ def init_softmax_state_odd(self, stage_info: StageInfo): Int32(0), ) + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def bind_sparse_request(self, stage_info: StageInfo): + """Cache live source lengths once per query tile, including graph replay.""" + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): + row = stage_info.work_tile.tile_idx[2] + self.request_lengths[0] = Int32(self.page_offsets.primary_lengths[row]) + self.request_lengths[1] = Int32(self.page_offsets.extra_lengths[row]) + @producer_work @cute.jit def qk_mma( @@ -1981,6 +2472,8 @@ def _load_s_impl( row_max_new, correction_factor_out, no_correction_out, + cached_mask_low=None, + cached_mask_high=None, ): """Load S from TMEM and compute the local row max.""" del row_sum_out, row_max_new, correction_factor_out, no_correction_out @@ -2040,6 +2533,51 @@ def _load_s_impl( ) qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + if cutlass.const_expr(cfg.page_size == 1): + group_needs_mask = True + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): + mask_page_offsets = self.page_offsets.bind_query_row( + work_tile.tile_idx[2], self.request_lengths[0], self.request_lengths[1] + ) + else: + mask_page_offsets = self.page_offsets + if cutlass.const_expr(cached_mask_low is not None): + mask_routes = CachedSparseMask((cached_mask_low, cached_mask_high), False) + # The cached invalid bits already include the live K boundary. + # Whole valid halves need no per-score masking instructions. + if cutlass.const_expr(not needs_row_causal_mask): + group_needs_mask = (cached_mask_low | cached_mask_high) != Int32(0) + else: + raw_mask_routes = ( + mask_page_offsets.select_source_tile( + tile_offset_k, work_tile.tile_idx[2], for_mask=True + ) + if cutlass.const_expr(isinstance(mask_page_offsets, SparseKVIndices)) + else self.page_offsets + ) + mask_routes = ( + cache_sparse_mask( + raw_mask_routes, + tile_offset_k, + K, + work_tile.tile_idx[2], + swaps=False, + column_half=local_tidx >> EPILOGUE_COLUMN_GROUP_SHIFT, + ) + if cutlass.const_expr( + cfg.page_size == 1 and self.logical_seq_len_q == 1 + ) + else raw_mask_routes + ) + if cutlass.const_expr( + cfg.page_size == 1 + and self.logical_seq_len_q == 1 + and not needs_row_causal_mask + ): + group_needs_mask = (mask_routes.low | mask_routes.high) != Int32(0) if group_needs_mask: if cutlass.const_expr(needs_row_causal_mask): # Clamp padded physical tail rows to the request's last real @@ -2070,8 +2608,21 @@ def _load_s_impl( token_idx - K + logical_seq_len_q ) token_is_visible = flat_query_row >= mask_flat_row + elif cutlass.const_expr( + cached_mask_low is not None + or getattr(self.page_offsets, "assume_valid_prefix", False) + ): + token_is_visible = not mask_routes.is_invalid( + token_idx, work_tile.tile_idx[2] + ) else: - token_is_visible = token_idx < K + token_is_visible = not invalid_sparse_token( + token_idx, + K, + mask_routes, + work_tile.tile_idx[2], + sparse=cfg.page_size == 1, + ) qk_acc_regs[i] = qk_acc_regs[i] if token_is_visible else neg_inf max0 = neg_inf @@ -2134,6 +2685,46 @@ def load_s( no_correction_out=no_correction_out, ) + @consumer_work( + returns=( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ), + ) + @cute.jit + def load_s_cached( + self, + stage_info: StageInfo, + *, + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + cached_mask_low, + cached_mask_high, + ): + """Load S for the even softmax group.""" + return self._load_s_impl( + stage_info, + qk_acc_regs=qk_acc_regs, + row_max=row_max, + row_sum=row_sum, + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + cached_mask_low=cached_mask_low, + cached_mask_high=cached_mask_high, + ) + @consumer_work( returns=( qk_acc_regs_odd, @@ -2264,8 +2855,8 @@ def load_peer_state(): fma_b = self.softmax_scale_log2 fma_c = Float32(0) - safe_row_max_new * self.softmax_scale_log2 if cutlass.const_expr(cfg.is_fp8_qkv()): - # Match the 448-scaled E4M3 P convention used by the reference - # output and the 1CTA implementation. + # Use the same selected FP8 probability scale as output and LSE. + # Native FP8 uses exact maxima and scale 448. fma_c = fma_c + fp8_log2_quant_scale() for i in cutlass.range_constexpr(0, 64, 2): fma_result = fma_packed_f32x2( @@ -3034,8 +3625,10 @@ def rescale_o_slice( @dataclass(kw_only=True) class GmemOResource(HighThroughputMlaResource): - """GMEM output. Producer: Correction (epilogue store). No pipeline.""" + """Final output: direct stores or staged quantization with a TMA consumer.""" + smem_o: Any = None + tma_o_desc: Any = None output: Any = None partial_output: Any = None lse: Any = None @@ -3043,6 +3636,10 @@ class GmemOResource(HighThroughputMlaResource): tmem_o_ref: Any = None # Reference to TmemOResource for tmem_base_addr tmem_corr_ref: Any = None # Reference to TmemCorrResource for correction data output_scale: Any = None + atten_sinks: Any = None + output_quant: cutlass.Constexpr = None + output_quant_args: Any = None + lse_in_natural_log: cutlass.Constexpr[bool] = False softmax_scale_log2: Any = None smem_exchange: Any = None # SMEM for row_sum exchange (as Int32 base addr) split_kv: Any = None @@ -3051,6 +3648,66 @@ class GmemOResource(HighThroughputMlaResource): logical_seq_len_q: cutlass.Constexpr[int] = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + @cute.jit + def _prefetch_output_rope(self, flat_row, batch_idx, local_tidx, head, valid): + """Fetch rotary coefficients while the nonrotary output is finalized.""" + tile_h = self.cfg.mma_pv_tiler[0] // self.cfg.num_mma_ctas + # One lane per logical query in this CTA; the first row may start + # partway through a head group. The other column group duplicates it. + if local_tidx < tile_h and valid and (head == 0 or local_tidx == 0): + row = Int64(flat_row) + if cutlass.const_expr(self.cu_seqlens_q is None): + row += ( + Int64(batch_idx) * self.logical_num_heads_q * self.logical_seq_len_q + ) + _, positions, cos_sin, _ = self.output_quant_args + position = positions[row // self.logical_num_heads_q] + cache_row = cos_sin.iterator.raw_ptr() + Int64(position) * 64 + prims.prefetch_l1(cache_row) + prims.prefetch_l1(cache_row + 32) + + @cute.jit + def _normalization_scale(self, row_sum, row_max, local_tidx, blk_coord): + safe_sum = row_sum if row_sum > Float32(0) else Float32(1) + scale = self.output_scale * cute.math.rcp(safe_sum, approx=True) + if cutlass.const_expr(self.atten_sinks is not None): + head = ( + Int32(blk_coord[1]) * self.cfg.mma_qk_tiler[0] + + Int32(blk_coord[0]) + * (self.cfg.mma_qk_tiler[0] // self.cfg.num_mma_ctas) + + (local_tidx & Int32(63)) + ) + if cutlass.const_expr(self.logical_seq_len_q == 1): + # Preserve the cheap single-query mapping, including padded heads. + head = cute.math.min(head, Int32(self.logical_num_heads_q - 1)) + else: + row = Int32(blk_coord[0]) * ( + self.cfg.mma_qk_tiler[0] // self.cfg.num_mma_ctas + ) + (local_tidx & Int32(63)) + _, head, _, _, _ = self._query_row_state( + row, Int32(blk_coord[1]), Int32(blk_coord[2]) + ) + sink = Float32(self.atten_sinks[head]) + scale = Float32(0) + if row_sum > Float32(0) and sink != Float32(Float32.inf): + max_log2 = row_max * self.softmax_scale_log2 + sink_log2 = sink * Float32(LOG2_E) + norm_max = cute.math.max(max_log2, sink_log2) + rescale = cute.math.exp2(max_log2 - norm_max, fastmath=True) + inv_quant = Float32( + fp8_quant_scale_rcp() if self.cfg.is_fp8_qkv() else 1.0 + ) + denominator = row_sum * inv_quant * rescale + cute.math.exp2( + sink_log2 - norm_max, fastmath=True + ) + scale = ( + self.output_scale + * inv_quant + * rescale + * cute.math.rcp(denominator, approx=True) + ) + return scale + @cute.jit def _query_row_state(self, row_in_tile, query_tile_idx, batch_idx): """Map one physical flat-tile row to public storage.""" @@ -3064,6 +3721,94 @@ def _query_row_state(self, row_in_tile, query_tile_idx, batch_idx): batch_idx, ) + @cute.jit + def _store_quantized_slice( + self, + values, + flat_row, + batch_idx, + column, + stage_idx=0, + ): + row = Int64(flat_row) + if cutlass.const_expr(self.cu_seqlens_q is None): + row += Int64(batch_idx) * self.logical_num_heads_q * self.logical_seq_len_q + shared_output = None + if cutlass.const_expr(self.cfg.use_tma_output): + shared_output = self.smem_o.data_ptr( + stage_idx * OUTPUT_TMA_STAGE_BYTES + ).toint(Int32) + self.output_quant.store( + self.output_quant_args, + values, + row, + column, + rotary=False, + shared_output=shared_output, + ) + if cutlass.const_expr(self.cfg.use_tma_output): + cute.arch.fence_view_async_shared() + + @consumer_work + @cute.jit + def tma_store_output( + self, stage_info: StageInfo, *, iter_n: cutlass.Constexpr[int] + ): + # Sparse launch flattens queries into batch rows; each CTA owns 64 heads. + head_tile_idx, _, token, _ = stage_info.work_tile.tile_idx + # A tile can span groups or lie inside the single H128 group. + heads_per_group = cutlass.const_expr( + self.output_quant.num_heads // self.output_quant.num_groups + ) + head_start = head_tile_idx * Int32(64) + group = head_start // heads_per_group + head_in_group = head_start % heads_per_group + if prims.elect_sync(): + for half in cutlass.range_constexpr(2): + prims.cp_async_bulk_tensor_global_shared_cta( + self.tma_o_desc, + self.smem_o.subview( + stage_info.stage_idx * OUTPUT_TMA_STAGE_BYTES + half * 8192 + ), + (iter_n * 256 + half * 128, head_in_group, token, group), + ) + prims.cp_async_bulk_commit_group() + # Release the stage only after TMA has consumed its shared source. + # The store task waits for global completion before retiring the CTA. + prims.cp_async_bulk_wait_group(0, read=True) + + @cute.jit + def _load_rotary_slice( + self, tmem_addr, norm_scale, flat_row, batch_idx, column, valid + ): + """Rotate the upper half before the other 64 accumulators become live.""" + row = Int64(flat_row) + if cutlass.const_expr(self.cu_seqlens_q is None): + row += Int64(batch_idx) * self.logical_num_heads_q * self.logical_seq_len_q + result = cutlass.Array(Float32, 128, space=cutlass.AddressSpace.rmem) + # The correction warp's register budget cannot keep 128 accumulators + # and a batch of RoPE coefficients live together without serial loads. + for half in cutlass.range_constexpr(2): + offset = cutlass.const_expr(64 if half == 0 else 0) + values = cutlass.Array(Float32, 64, space=cutlass.AddressSpace.rmem) + for chunk in cutlass.range_constexpr(2): + ptr = prims.make_tmem_ptr(tmem_addr + offset + chunk * 32, Float32) + values.store( + prims.tcgen05_ld(TCGEN05_32B_SHAPE, ptr, num=32), chunk * 32 + ) + for i in cutlass.range_constexpr(0, 64, 2): + scaled = mul_packed_f32x2( + (values[i], values[i + 1]), (norm_scale, norm_scale) + ) + values[i], values[i + 1] = scaled[0], scaled[1] + if cutlass.const_expr(half == 0): + if valid: + self.output_quant.rotate( + self.output_quant_args, values, row, column + offset + ) + result.store(values.load(0, 64), offset) + return result + @producer_work @cute.jit def epilogue_store(self, stage_info: StageInfo) -> None: @@ -3146,12 +3891,22 @@ def epilogue_store(self, stage_info: StageInfo) -> None: query_is_valid, ) = self._query_row_state(row_in_tile, seq_q_idx, batch_idx) + if cutlass.const_expr( + self.output_quant is not None and self.partial_output is None + ): + self._prefetch_output_rope( + storage_flat_query_row, + batch_idx, + local_tidx, + logical_head_idx, + row_in_tile < physical_tile_rows and query_is_valid, + ) + # Fully masked split rows can occur when physical tail rows or earlier # causal query rows have no visible K values in this split. Store zero # O and -inf LSE so split-KV reduction gives those rows zero weight. row_has_values = row_sum > Float32(0) - safe_row_sum = row_sum if row_has_values else Float32(1) - norm_scale = self.output_scale * cute.math.rcp(safe_row_sum, approx=True) + norm_scale = self._normalization_scale(row_sum, row_max, local_tidx, blk_coord) for iter_n in cutlass.range_constexpr(cfg.iterations_pv_n): # Load O from TMEM @@ -3159,22 +3914,36 @@ def epilogue_store(self, stage_info: StageInfo) -> None: tmem_raw_addr_n = tmem_raw_addr + ( iter_n * (cfg.mma_pv_tiler[1] // cfg.warps_in_n) ) - for load_idx in cutlass.range_constexpr(num_tmem_loads): - curr_addr = tmem_raw_addr_n + load_idx * TCGEN05_32B_REGS_PER_LOAD - tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) - loaded = prims.tcgen05_ld( - t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + if cutlass.const_expr( + self.output_quant is not None + and self.partial_output is None + and iter_n + 1 == cfg.iterations_pv_n + ): + qk_acc_regs = self._load_rotary_slice( + tmem_raw_addr_n, + norm_scale, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + row_in_tile < physical_tile_rows and query_is_valid, ) - qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + else: + for load_idx in cutlass.range_constexpr(num_tmem_loads): + curr_addr = tmem_raw_addr_n + load_idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + loaded = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) - # Normalize: O = O * output_scale / row_sum - for i in cutlass.range_constexpr(0, 128, 2): - scaled = mul_packed_f32x2( - (qk_acc_regs[i], qk_acc_regs[i + 1]), - (norm_scale, norm_scale), - ) - qk_acc_regs[i] = scaled[0] - qk_acc_regs[i + 1] = scaled[1] + # Normalize: O = O * output_scale / row_sum + for i in cutlass.range_constexpr(0, 128, 2): + scaled = mul_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (norm_scale, norm_scale), + ) + qk_acc_regs[i] = scaled[0] + qk_acc_regs[i + 1] = scaled[1] # Store O to GMEM if row_in_tile < physical_tile_rows and query_is_valid: @@ -3219,6 +3988,13 @@ def epilogue_store(self, stage_info: StageInfo) -> None: vec_partial, evict="noallocate", ) + elif cutlass.const_expr(self.output_quant is not None): + self._store_quantized_slice( + qk_acc_regs, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + ) else: # 16-bit output (split_kv == 1, direct output) if cutlass.const_expr(self.cu_seqlens_q is not None): @@ -3303,6 +4079,8 @@ def epilogue_store(self, stage_info: StageInfo) -> None: else Float32(-Float32.inf) ) + if cutlass.const_expr(self.lse_in_natural_log): + lse *= Float32(LN_2) # Use local_tidx (0..127 within correction warpgroup) for LSE # indexing, not global tidx (which is 128..255 for correction warps). lse_tidx = local_tidx @@ -3386,24 +4164,8 @@ def epilogue_store_slice( t2r_shape = TCGEN05_32B_SHAPE num_tmem_loads = 4 qk_acc_regs = cutlass.Array(Float32, 128, space=cutlass.AddressSpace.rmem) - for load_idx in cutlass.range_constexpr(num_tmem_loads): - curr_addr = tmem_raw_addr + load_idx * TCGEN05_32B_REGS_PER_LOAD - tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) - loaded = prims.tcgen05_ld( - t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD - ) - qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) - row_has_values = row_sum > Float32(0) - safe_row_sum = row_sum if row_has_values else Float32(1) - norm_scale = self.output_scale * cute.math.rcp(safe_row_sum, approx=True) - for i in cutlass.range_constexpr(0, 128, 2): - scaled = mul_packed_f32x2( - (qk_acc_regs[i], qk_acc_regs[i + 1]), - (norm_scale, norm_scale), - ) - qk_acc_regs[i] = scaled[0] - qk_acc_regs[i + 1] = scaled[1] + norm_scale = self._normalization_scale(row_sum, row_max, local_tidx, blk_coord) tile_h = cfg.mma_pv_tiler[0] // cfg.num_mma_ctas tile_d = cfg.mma_pv_tiler[1] @@ -3427,6 +4189,48 @@ def epilogue_store_slice( query_is_valid, ) = self._query_row_state(row_in_tile, seq_q_idx, batch_idx) + if cutlass.const_expr( + self.output_quant is not None + and self.partial_output is None + and iter_n == 0 + ): + self._prefetch_output_rope( + storage_flat_query_row, + batch_idx, + local_tidx, + logical_head_idx, + row_in_tile < physical_tile_rows and query_is_valid, + ) + + if cutlass.const_expr( + self.output_quant is not None + and self.partial_output is None + and iter_n + 1 == cfg.iterations_pv_n + ): + qk_acc_regs = self._load_rotary_slice( + tmem_raw_addr, + norm_scale, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + row_in_tile < physical_tile_rows and query_is_valid, + ) + else: + for load_idx in cutlass.range_constexpr(num_tmem_loads): + curr_addr = tmem_raw_addr + load_idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + loaded = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + for i in cutlass.range_constexpr(0, 128, 2): + scaled = mul_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (norm_scale, norm_scale), + ) + qk_acc_regs[i] = scaled[0] + qk_acc_regs[i + 1] = scaled[1] + if row_in_tile < physical_tile_rows and query_is_valid: if cutlass.const_expr(self.partial_output is not None): S_q = ( @@ -3465,6 +4269,16 @@ def epilogue_store_slice( vec_partial, evict="noallocate", ) + elif cutlass.const_expr(self.output_quant is not None): + self._store_quantized_slice( + qk_acc_regs, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + stage_idx=stage_info.stage_idx + if cutlass.const_expr(cfg.use_tma_output) + else 0, + ) else: if cutlass.const_expr(self.cu_seqlens_q is not None): o_base_ptr = self.output.iterator.raw_ptr() + Int64( @@ -3543,6 +4357,8 @@ def epilogue_store_slice( if row_has_values else Float32(-Float32.inf) ) + if cutlass.const_expr(self.lse_in_natural_log): + lse *= Float32(LN_2) lse_tidx = local_tidx if lse_tidx < tile_h: lse_row_in_tile = head_tile_idx * tile_h + lse_tidx diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py index 0ae3fa0f8534..0a98fca8ea6a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py @@ -42,6 +42,7 @@ import cutlass import cutlass.cute as cute +from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.memory import ResourceContext from cutlass.experimental.task_scheduling.schedule_builder import ( @@ -468,30 +469,62 @@ def load_tma_schedule(page_offset_window, smem_q, smem_kv, work_queue=None): return task_class( src_resources=src, dst_resources=[smem_q, smem_kv], - warp_idx=9, - num_warps=1, + warp_idx=smem_kv.cfg.load_tma_warp_id, + num_warps=smem_kv.cfg.load_num_warps, schedule=captured_schedule, name="LoadTmaTask", **task_kwargs, ) +def create_sparse_offsets_task( + page_offsets, work_queue=None, task_class=MlaTask, **task_kwargs +): + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + + @schedule + def offsets_schedule(page_offsets, work_queue=None): + page_offsets.init_sparse_state() + page_offsets.bind_sparse_request() + with domain_loop(loop_start, loop_end, loop_step): + page_offsets.acquire() + page_offsets.load_sparse_offsets() + page_offsets.commit() + work_queue_tail(work_queue, advance_label="advance_tile") + + result = ( + offsets_schedule(page_offsets) + if work_queue is None + else offsets_schedule(page_offsets, work_queue) + ) + return task_class( + src_resources=[work_queue] if work_queue is not None else [], + dst_resources=[page_offsets], + warp_idx=9, + num_warps=1, + schedule=result, + name="SparseOffsetsTask", + **task_kwargs, + ) + + def create_load_k_task( smem_q: SmemQResource, smem_k: SmemKResource, + page_offsets=None, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, **task_kwargs, ) -> Task: - """Create the FP8 Q/K TMA task (warp 9). + """Create the FP8 Q/K TMA task on its configured loader warps. Q is loaded once before the loop. K uses one whole-tile pipeline stage per - logical K tile and reads page offsets directly from GMEM. + logical K tile. Sparse indices are prefetched through the shared offset + ring; ordinary paged inputs read page offsets directly from GMEM. """ loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) - @schedule - def load_k_schedule(smem_q, smem_k, work_queue=None): + def load_k_body(smem_q, smem_k, page_offsets=None, work_queue=None): """Load Q once and then publish one K stage per loop tile.""" smem_q.init_load_state() smem_k.init_load_state() @@ -500,23 +533,50 @@ def load_k_schedule(smem_q, smem_k, work_queue=None): smem_q.commit() with domain_loop(loop_start, loop_end, loop_step): - smem_k.acquire() - smem_k.tma_load_direct() - smem_k.commit() + if page_offsets is not None: + page_offsets.wait() + page_stage = page_offsets.sparse_stage() + smem_k.acquire() + smem_k.tma_load_shared(page_offset_stage=page_stage) + smem_k.commit() + # The last quad has been copied into TMA operands. Only now + # may the metadata producer overwrite this shared ring slot. + page_offsets.release() + else: + smem_k.acquire() + smem_k.tma_load_paged() + smem_k.commit() work_queue_tail(work_queue, advance_label="advance_tile") - schedule_result = ( - load_k_schedule(smem_q, smem_k) - if work_queue is None - else load_k_schedule(smem_q, smem_k, work_queue) - ) + @schedule + def load_k_schedule(smem_q, smem_k, page_offsets, work_queue=None): + load_k_body(smem_q, smem_k, page_offsets, work_queue) - src = [work_queue] if work_queue is not None else [] + @schedule + def load_k_dense_schedule(smem_q, smem_k, work_queue=None): + load_k_body(smem_q, smem_k, None, work_queue) + + if page_offsets is None: + schedule_result = ( + load_k_dense_schedule(smem_q, smem_k) + if work_queue is None + else load_k_dense_schedule(smem_q, smem_k, work_queue) + ) + else: + schedule_result = ( + load_k_schedule(smem_q, smem_k, page_offsets) + if work_queue is None + else load_k_schedule(smem_q, smem_k, page_offsets, work_queue) + ) + + src = ([page_offsets] if page_offsets is not None else []) + ( + [work_queue] if work_queue is not None else [] + ) return task_class( src_resources=src, dst_resources=[smem_q, smem_k], - warp_idx=9, - num_warps=1, + warp_idx=smem_k.cfg.load_k_warp_id, + num_warps=smem_k.cfg.load_num_warps, schedule=schedule_result, name="LoadKTask", **task_kwargs, @@ -525,35 +585,61 @@ def load_k_schedule(smem_q, smem_k, work_queue=None): def create_load_v_task( smem_v: SmemVResource, + page_offsets=None, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, **task_kwargs, ) -> Task: - """Create the FP8 V TMA task (warp 10).""" + """Create the FP8 V TMA task on its configured loader warps.""" loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) - @schedule - def load_v_schedule(smem_v, work_queue=None): + def load_v_body(smem_v, page_offsets=None, work_queue=None): """Publish one V stage per logical K tile.""" smem_v.init_load_state() with domain_loop(loop_start, loop_end, loop_step): - smem_v.acquire() - smem_v.tma_load_direct() - smem_v.commit() + if page_offsets is not None: + page_offsets.wait() + page_stage = page_offsets.sparse_stage() + smem_v.acquire() + smem_v.tma_load_shared(page_offset_stage=page_stage) + smem_v.commit() + # Keep the ring slot protected through every shared read. + page_offsets.release() + else: + smem_v.acquire() + smem_v.tma_load_paged() + smem_v.commit() work_queue_tail(work_queue, advance_label="advance_tile") - schedule_result = ( - load_v_schedule(smem_v) - if work_queue is None - else load_v_schedule(smem_v, work_queue) - ) + @schedule + def load_v_schedule(smem_v, page_offsets, work_queue=None): + load_v_body(smem_v, page_offsets, work_queue) - src = [work_queue] if work_queue is not None else [] + @schedule + def load_v_dense_schedule(smem_v, work_queue=None): + load_v_body(smem_v, None, work_queue) + + if page_offsets is None: + schedule_result = ( + load_v_dense_schedule(smem_v) + if work_queue is None + else load_v_dense_schedule(smem_v, work_queue) + ) + else: + schedule_result = ( + load_v_schedule(smem_v, page_offsets) + if work_queue is None + else load_v_schedule(smem_v, page_offsets, work_queue) + ) + + src = ([page_offsets] if page_offsets is not None else []) + ( + [work_queue] if work_queue is not None else [] + ) return task_class( src_resources=src, dst_resources=[smem_v], - warp_idx=10, - num_warps=1, + warp_idx=smem_v.cfg.load_v_warp_id, + num_warps=smem_v.cfg.load_num_warps, schedule=schedule_result, name="LoadVTask", **task_kwargs, @@ -897,6 +983,7 @@ def create_softmax_task( smem_p: SmemPResource, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, + page_offsets=None, warp_idx: int = 0, name: str = "SoftmaxTask", softmax_group_id: int = 0, @@ -920,9 +1007,15 @@ def softmax_prelude(tmem_s): ) return init_softmax_state() - def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): + def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state, page_offsets=None): """Consume S tiles, materialize P, and publish correction factors.""" - load_s = tmem_s.load_s_odd if softmax_group_id == 1 else tmem_s.load_s + load_s = ( + tmem_s.load_s_odd + if softmax_group_id == 1 + else tmem_s.load_s_cached + if page_offsets is not None + else tmem_s.load_s + ) finish_softmax = ( tmem_s.finish_softmax_odd if softmax_group_id == 1 @@ -947,7 +1040,17 @@ def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): no_correction_out, ) = softmax_state + tmem_s.bind_sparse_request() with domain_loop(loop_start, loop_end, loop_step): + mask_kwargs = {} + if page_offsets is not None: + page_offsets.wait() + mask_low, mask_high = page_offsets.read_sparse_mask() + page_offsets.release() + mask_kwargs = { + "cached_mask_low": mask_low, + "cached_mask_high": mask_high, + } tmem_s.wait() if softmax_group_id == 1: ( @@ -984,6 +1087,7 @@ def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): row_max_new=row_max_new, correction_factor_out=correction_factor_out, no_correction_out=no_correction_out, + **mask_kwargs, ) tmem_s.release() if softmax_group_id == 1: @@ -1057,29 +1161,44 @@ def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): ) tmem_corr.commit() - @schedule - def softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue=None): - """Capture one active softmax tile and unconditional queue progress.""" - + def capture_softmax(tmem_s, tmem_corr, smem_p, work_queue, page_offsets): _capture_clc_work_tile_body( work_queue, lambda softmax_state: softmax_body( - tmem_s, - tmem_corr, - smem_p, - softmax_state, + tmem_s, tmem_corr, smem_p, softmax_state, page_offsets ), lambda: softmax_prelude(tmem_s), use_clc_dynamic=use_clc_dynamic, ) - schedule_result = ( - softmax_schedule(tmem_s, tmem_corr, smem_p) - if work_queue is None - else softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue) - ) + @schedule + def softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue=None): + capture_softmax(tmem_s, tmem_corr, smem_p, work_queue, None) + + @schedule + def sparse_softmax_schedule( + tmem_s, tmem_corr, smem_p, page_offsets, work_queue=None + ): + capture_softmax(tmem_s, tmem_corr, smem_p, work_queue, page_offsets) + + if page_offsets is None: + schedule_result = ( + softmax_schedule(tmem_s, tmem_corr, smem_p) + if work_queue is None + else softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue) + ) + else: + schedule_result = ( + sparse_softmax_schedule(tmem_s, tmem_corr, smem_p, page_offsets) + if work_queue is None + else sparse_softmax_schedule( + tmem_s, tmem_corr, smem_p, page_offsets, work_queue + ) + ) src = [tmem_s] + if page_offsets is not None: + src.append(page_offsets) if work_queue is not None: src.append(work_queue) return task_class( @@ -1117,6 +1236,7 @@ def create_correction_task( """ loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 1) use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + use_tma_output = bool(gmem_o.cfg.use_tma_output) def correction_prelude(tmem_corr): """Create correction and epilogue register state before skip guards.""" @@ -1161,12 +1281,16 @@ def correction_body(tmem_corr, tmem_o, gmem_o, correction_state): tmem_corr.prepare_epilogue_slice_store() ) for iter_n in range(iterations_pv_n): + if use_tma_output: + gmem_o.acquire() tmem_o.wait() gmem_o.epilogue_store_slice( row_sum=epilogue_row_sum, row_max=epilogue_row_max, iter_n=iter_n, ) + if use_tma_output: + gmem_o.commit() tmem_o.release() else: tmem_o.wait() @@ -1297,3 +1421,42 @@ def scheduler_schedule(work_queue, work_throttle=None): run_only_on_cta_id=0, **task_kwargs, ) + + +class MlaOutputTmaTask(MlaTask): + """Drain global TMA writes once after the persistent work loop.""" + + @cute.jit + def _drain_mla_work_tile_tails(self): + MlaTask._drain_mla_work_tile_tails(self) + prims.cp_async_bulk_wait_group(0) + + +def create_output_tma_task(output, work_queue=None, **task_kwargs): + """Consume two quantized N slices on the otherwise idle sparse FP8 warp.""" + + @schedule + def output_schedule(output, work_queue=None): + # This task only has epilogue work; declare its empty K loop. + with domain_loop(0, 0, 1): + pass + for iter_n in range(2): + output.wait() + output.tma_store_output(iter_n=iter_n) + output.release() + work_queue_tail(work_queue, advance_label="advance_tile") + + captured = ( + output_schedule(output) + if work_queue is None + else output_schedule(output, work_queue) + ) + return MlaOutputTmaTask( + src_resources=[output] + ([work_queue] if work_queue is not None else []), + dst_resources=[], + warp_idx=10, + num_warps=1, + schedule=captured, + name="OutputTmaTask", + **task_kwargs, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py index dedc950131d2..7657de52608d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py @@ -26,7 +26,7 @@ from __future__ import annotations -from dataclasses import dataclass, replace +from dataclasses import dataclass, replace as dataclass_replace, replace from math import ceil from typing import Any @@ -165,6 +165,17 @@ class MlaConfig: correction_num_warps: int = 4 mma_num_warps: int = 1 load_num_warps: int = 1 + sparse_reuse_kv: bool = False + paired_sparse_correction: bool = False + defer_sparse_max_update: bool = False + cache_uniform_sparse_quads: bool = False + fuse_sparse_epilogue: bool = False + fuse_sparse_cluster_epilogue: bool = False + fuse_sparse_reduction: bool = False + sparse_direct: bool = False + sparse_static_scales: bool = False + sparse_valid_prefix: bool = False + sparse_direct_capacities: tuple[int, int] = (0, 0) page_offsets_num_warps: int = 1 scheduler_num_warps: int = 1 clc_padding_warp_idx: int = 14 @@ -193,6 +204,47 @@ class MlaConfig: use_sliding_window_causal: int = 0 attention_window_size: int = 0 + @property + def one_insts_kv_swap(self) -> bool: + """One KV stream with swapped MMA operands (P remains in SMEM).""" + return self.num_insts_kv == 1 and self.kernel_variant == "swaps_mma_ab" + + @property + def single_kv_pipe(self) -> bool: + """Use one softmax/P stream while retaining two score stages.""" + return self.num_insts_kv == 1 + + @property + def short_merged_softmax(self) -> bool: + """Process two independent KV streams with one softmax warp group. + + Each planned split fits one two-tile group, so no online-correction + loop runs. Keep both streams' final statistics and PV accumulators; + live lengths still mask partial/empty tiles. This is num_insts_kv=2, + not the one-stream schedule selected by one_insts_kv_swap. + """ + return ( + self.sparse_direct + and not self.one_insts_kv_swap + and self.kernel_variant == "swaps_mma_ab" + and self.fixed_sparse_kv_tiles > 0 + ) + + @property + def serial_swap_reuse(self): + return ( + self.sparse_reuse_kv + and self.num_insts_kv == 2 + and not self.is_fp8_qkv() + and self.kernel_variant == "swaps_mma_ab" + ) + + @property + def sparse_cache_words(self): + if self.cache_uniform_sparse_quads: + return 6 + self.tile_size_kv // self.load_num_warps + return 6 + @property def softmax0_num_warps(self) -> int: return self.softmax_num_warps @@ -369,6 +421,21 @@ def is_fp8_qkv(self) -> bool: return self.qkv_dtype == "e4m3" + @property + def fixed_sparse_kv_tiles(self) -> int: + """Return the minimum instruction group when bounded by the plan. + + Keep-AB uses one KV tile and Swap-AB uses two independent streams. + If the maximum planned split fits that group, every shorter valid + prefix already has the same partition. Live lengths still mask data. + """ + tiles = self.local_kv_tiles(self.total_kv_tiles) + return ( + self.num_insts_kv + if self.sparse_direct and tiles == self.num_insts_kv + else 0 + ) + def local_kv_tiles(self, total_kv_tiles: int) -> int: """Return per-CTA KV tiles after multi-CTA KV splitting.""" if self.use_multi_ctas_kv != 1: @@ -795,11 +862,9 @@ def resolve_runtime_cluster_reduction_mode( cfg: MlaConfig | None, *, reduction_mode: str | None, - hardware_info, - stream=None, log=None, ) -> str | None: - """Query cluster occupancy and return the final runtime reduction mode.""" + """Use cached device occupancy to resolve the runtime reduction mode.""" if ( cfg is None @@ -810,7 +875,9 @@ def resolve_runtime_cluster_reduction_mode( cluster_shape = cluster_reduction_cluster_shape(cfg) cluster_size = cluster_shape[0] * cluster_shape[1] * cluster_shape[2] - active_clusters = hardware_info.get_max_active_clusters(cluster_size, stream) + from ...._hardware import get_max_active_clusters + + active_clusters = get_max_active_clusters(cluster_size) cluster_count = cluster_reduction_cluster_count(cfg) if log is not None: log( @@ -1233,10 +1300,13 @@ def is_throughput_latency_mla_supported_shape( del batch_size return ( latent_dim == 512 - and rope_dim == 64 + and rope_dim in (0, 64) and 1 <= num_heads_q <= 128 and is_power_of_two(num_heads_q) - and num_tokens_per_page in SUPPORTED_MLA_PAGE_SIZES + and ( + num_tokens_per_page in SUPPORTED_MLA_PAGE_SIZES + or (num_tokens_per_page == 1 and rope_dim == 0) + ) and seq_len_q >= 1 and seq_len_kv >= 128 ) @@ -1479,7 +1549,9 @@ def make_throughput_latency_mla_config( raise ValueError("latent_dim must be positive") if rope_dim < 0: raise ValueError("rope_dim must be non-negative") - if num_tokens_per_page not in SUPPORTED_MLA_PAGE_SIZES: + if num_tokens_per_page not in SUPPORTED_MLA_PAGE_SIZES and not ( + num_tokens_per_page == 1 and rope_dim == 0 + ): raise ValueError( "num_tokens_per_page must be one of " f"{SUPPORTED_MLA_PAGE_SIZES}, got {num_tokens_per_page}" @@ -1491,7 +1563,10 @@ def make_throughput_latency_mla_config( f"num_tokens_per_page={num_tokens_per_page}" ) pages_per_kv_tile = MlaConfig.tile_size_kv // num_tokens_per_page - if pages_per_kv_tile > MlaConfig.page_offsets_entries_per_stage: + metadata_capacity = ( + 128 if num_tokens_per_page == 1 else MlaConfig.page_offsets_entries_per_stage + ) + if pages_per_kv_tile > metadata_capacity: raise ValueError( "page-offset staging capacity is smaller than one KV tile: " f"pages_per_kv_tile={pages_per_kv_tile}, " @@ -1528,6 +1603,11 @@ def make_throughput_latency_mla_config( selected_profile, explicit_persistent, ) + if selected_profile.use_persistent_scheduler and selected_profile.use_multi_ctas_kv: + # Persistent work queues enumerate Q/head-dimension tiles, not KV + # splits. Direct MlaProfile callers must obey the same restriction + # as persistent_override_profile, or only split zero gets produced. + raise ValueError("persistent scheduling is not supported with split-KV") tile_size_q = tile_size_q_for_profile( selected_profile, num_heads_q, seq_len_q, tile_size_q ) @@ -1612,6 +1692,7 @@ def make_throughput_latency_mla_config( softmax_regs=softmax_register_budget(tile_size_q), correction_regs=correction_register_budget(tile_size_q), num_tokens_per_page=num_tokens_per_page, + page_offsets_entries_per_stage=metadata_capacity, max_num_pages_per_seq_kv=max(1, ceil(seq_len_kv / num_tokens_per_page)), tmem_s_cols=tile_size_q, tmem_stats_cols=MlaConfig.tmem_stats_cols, @@ -1635,3 +1716,198 @@ def make_throughput_latency_mla_config( reduction_mode=reduction_mode, ) return cfg + + +def configure_sparse_mla( + cfg: MlaConfig, + profile, + *, + finalize_output=False, + direct_inputs=False, + assume_valid_prefix=False, + static_scales=False, + source_capacities=(0, 0), +) -> MlaConfig: + """Resolve one selected sparse profile without order-dependent overrides. + + Policy selects the schedule. This builder checks its implementation + constraints, fixes topology, then assigns buffers and warp/register roles. + """ + keep = cfg.kernel_variant == "keeps_mma_ab" + issuers = profile.gather_issue_warps + reuse = profile.reuse_kv + uniform = profile.uniform_offset_cache + single = profile.single_kv_stream + if cfg.num_tokens_per_page != 1 or cfg.rope_dim != 0: + raise ValueError("sparse profiles require native page-size-one MLA") + if issuers not in (1, 2, 4, 8): + raise ValueError("sparse profiles require one, two, four or eight issuers") + if (keep and cfg.tile_size_q != 64) or ( + not keep and cfg.tile_size_q not in (8, 16, 32) + ): + raise ValueError("unsupported sparse query tile") + if single and ( + keep + or cfg.tile_size_q not in (8, 16) + or cfg.num_ctas_per_head_dim != 1 + or issuers != 4 + or cfg.use_cluster_reduction + or (cfg.num_ctas_per_seq_kv > 1 and cfg.use_persistent_scheduler) + ): + raise ValueError( + "one-instance Swap-AB requires M8/M16, V1, four issuers and no cluster reduction" + ) + if reuse and cfg.num_ctas_per_head_dim != 1: + raise ValueError("retained KV requires one V owner") + if profile.kv_tile_size not in (64, 128) or ( + profile.kv_tile_size == 64 and (not keep or cfg.is_fp8_qkv() or not reuse) + ): + raise ValueError("BK64 requires native BF16 Keep-AB reuse") + if uniform and ( + not reuse + or issuers not in (4, 8) + or (cfg.is_fp8_qkv() and not keep) + or (issuers == 8 and not keep) + ): + raise ValueError("uniform coordinate caching requires a qualified reuse layout") + if profile.balanced_registers and (keep or cfg.tile_size_q not in (8, 16)): + raise ValueError("balanced registers require Swap-AB M8/M16") + if profile.defer_max_update and cfg.is_fp8_qkv(): + raise ValueError("deferred maxima require native BF16 probabilities") + if profile.paired_correction and not keep: + raise ValueError("paired correction requires Keep-AB") + if profile.page_pipeline_stages and (profile.page_pipeline_stages < 1 or not keep): + raise ValueError("page-stage overrides require positive stages and Keep-AB") + + # Final stream count and tile geometry must precede short-loop eligibility. + splits = cfg.num_ctas_per_seq_kv + cluster_finish = finalize_output and bool(cfg.use_cluster_reduction) + direct_finish = finalize_output and splits == 1 + separate_finish = finalize_output and splits > 1 and not cluster_finish + if cluster_finish and ( + keep or cfg.tile_size_q not in (16, 32) or not cfg.use_bf16_output + ): + raise ValueError("cluster finishing requires native Swap-AB M16/M32") + if direct_finish and cfg.logical_seq_len_q != 1: + raise ValueError("fused output requires one logical query per virtual request") + if direct_inputs and (not finalize_output or cfg.qkv_dtype not in ("bf16", "e4m3")): + raise ValueError("separate sparse indices require native BF16/FP8 final output") + cfg = dataclass_replace( + cfg, + num_insts_kv=1 if keep or single else 2, + tile_size_kv=profile.kv_tile_size, + page_offsets_entries_per_stage=profile.kv_tile_size, + sparse_reuse_kv=reuse, + cache_uniform_sparse_quads=uniform, + paired_sparse_correction=profile.paired_correction, + defer_sparse_max_update=profile.defer_max_update, + sparse_direct=direct_inputs, + sparse_static_scales=static_scales, + sparse_valid_prefix=assume_valid_prefix, + sparse_direct_capacities=source_capacities if direct_inputs else (0, 0), + fuse_sparse_epilogue=direct_finish, + fuse_sparse_cluster_epilogue=cluster_finish, + fuse_sparse_reduction=separate_finish, + load_num_warps=issuers, + ) + short = cfg.short_merged_softmax + if ( + issuers == 8 + and not keep + and ( + cfg.tile_size_q not in (8, 16) + or cfg.local_kv_tiles(cfg.total_kv_tiles) <= cfg.num_insts_kv + ) + ): + raise ValueError("eight Swap-AB issuers require a steady-state loop") + + # Four D128 slices make one D512 tile. FP8/BK64 reuse keeps two tiles; + # BF16 BK128 retains one, with additional stages used for prefetch. + kv_stages = ( + profile.reuse_kv_stages if reuse and profile.reuse_kv_stages else cfg.kv_stages + ) + if reuse and profile.reuse_kv_stages: + retained_tiles = 2 if cfg.is_fp8_qkv() or cfg.tile_size_kv == 64 else 1 + if kv_stages < retained_tiles * cfg.qk_head_dim_stages: + raise ValueError("insufficient stages to retain KV through PV") + if cfg.tile_size_kv == 64 and profile.reuse_kv_stages < 2 * cfg.qk_head_dim_stages: + raise ValueError("BK64 requires two retained K tiles") + page_stages = cfg.page_offsets_stages + if ( + reuse + and profile.reuse_kv_stages + and keep + and (kv_stages == (12 if cfg.is_fp8_qkv() else 5) or cfg.tile_size_kv == 64) + ): + page_stages = 1 + if profile.page_pipeline_stages: + page_stages = profile.page_pipeline_stages + if single and reuse and kv_stages >= 6 and not cfg.is_fp8_qkv(): + page_stages = 1 + + # Select the final layout once. Sharing a softmax group in a short split + # preserves two numerical streams; the one-instance layout is distinct. + if single: + roles = dict( + threads_per_cta=512, + correction_warp_idx=4, + mma_warp_idx=8, + load_warp_idx=12, + page_offsets_warp_idx=10, + scheduler_warp_idx=11, + softmax_regs=128, + correction_regs=160, + mma_load_regs=64, + scheduler_regs=64, + o_stages=1, + ) + elif short: + roles = dict( + threads_per_cta=512, + correction_warp_idx=4, + mma_warp_idx=8, + load_warp_idx=12, + page_offsets_warp_idx=9, + scheduler_warp_idx=10, + softmax_regs=96, + correction_regs=96, + mma_load_regs=96, + scheduler_regs=96, + ) + elif keep: + roles = dict( + threads_per_cta=640 if issuers == 8 else 512, + load_warp_idx=12, + softmax_regs=128 if uniform or issuers == 8 else 192, + correction_regs=160 if issuers == 8 else 192, + mma_load_regs=96 if uniform and issuers == 4 else 64, + ) + else: + balanced = profile.balanced_registers + roles = dict( + threads_per_cta=768 if issuers == 8 else 640, + load_warp_idx=16, + softmax_regs=96 + if balanced + else 128 + if uniform + else 144 + if cfg.tile_size_q == 32 + else 160, + correction_regs=96 if balanced else 192 if uniform else cfg.correction_regs, + mma_load_regs=(64 if issuers == 8 else 96) + if balanced + else 96 + if uniform + else 48, + scheduler_regs=(64 if issuers == 8 else 96) + if balanced + else cfg.scheduler_regs, + ) + return dataclass_replace( + cfg, + **roles, + kv_stages=kv_stages, + page_offsets_stages=page_stages, + q_stages=1 if single or (reuse and profile.reuse_kv_stages) else cfg.q_stages, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py index cab0923d5959..130297f3b4a9 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py @@ -36,7 +36,6 @@ ) from cutlass.experimental import cuda from cutlass.experimental import primitives as prims -from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo from cutlass.experimental.task_scheduling.enums import PipelineType, SignalingThreads from cutlass.experimental.task_scheduling.memory import ( @@ -49,12 +48,14 @@ PipelineConfig, TileSchedulerConfig, WorkQueue, + WorkTileInfo, ) from cutlass.experimental.task_scheduling.task_manager import TaskManager from .config import ( MlaConfig, make_throughput_latency_mla_config, + configure_sparse_mla, ) from .resources import ( SmemKvResource, @@ -74,7 +75,7 @@ MlaDecodeTask, create_throughput_latency_correction_task, create_keeps_mma_ab_correction_task, - create_keeps_mma_ab_mma_task, + create_single_kv_pipe_mma_task, create_keeps_mma_ab_softmax_task, create_load_page_offsets_task, create_padding_task, @@ -94,7 +95,7 @@ supports_parallel_gmem_reduction, ) from .reduction import gmem_reduction_launch_shape, run_gmem_reduction_kernel -from ..helpers.constants import TMEM_LIFECYCLE_BARRIER_ID +from ..helpers.constants import LOG2_E, TMEM_LIFECYCLE_BARRIER_ID, TMEM_READY_BARRIER_ID from ..helpers.mask import MaskType, normalize_mask_type from ..helpers.tile import ( runtime_query_tile_is_active, @@ -109,10 +110,6 @@ ) -# Softmax uses exp2, so natural-scale scores are multiplied by log2(e). -LOG2_E = 1.4426950408889634 - - @cute.jit def _publish_neutral_standalone_partial( cfg, @@ -331,6 +328,25 @@ def __init__( self.cache_seqs = cache_seqs self.cu_seqlens_q = cu_seqlens_q self.enable_runtime_skip = cu_seqlens_q is not None + self._clc_response_alloc = SmemAllocation( + self.name + "_clc_response", dtype=cutlass.Int128, count=2, alignment=16 + ) + self._response_allocator = None + + def get_smem_requirements(self): + return [self._clc_response_alloc] + + def create_tile_scheduler(self): + response = self._response_allocator.get(self._clc_response_alloc) + response_ptr = cute.make_ptr( + cutlass.Int128, response.data_ptr(), mem_space=cutlass.AddressSpace.smem + ) + return utils.ClcDynamicPersistentTileScheduler.create( + self.tile_scheduler_config.tile_scheduler_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + response_ptr, + ) @cute.jit def skip_work_tile_if(self, work_tile: WorkTileInfo): @@ -378,7 +394,7 @@ def _make_static_work_queue( tile_scheduler_params=tile_sched_params, ), cfg=cfg, - batch_size=cute.size(cache_seqs), + batch_size=cute.size(cache_seqs.shape), cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, name=name, @@ -457,6 +473,9 @@ def build_throughput_latency_mla_task_manager( cta_idx_head_dim_v=None, scale_softmax_log2=None, output_scale=None, + sparse_epilogue_params=None, + output_quant=None, + output_quant_args=None, o_tensor=None, lse_tensor=None, acc_o_tensor=None, @@ -485,8 +504,8 @@ def build_throughput_latency_mla_task_manager( use_clc_dynamic_scheduler, use_static_persistent_scheduler, ) - if cfg.kernel_variant == "keeps_mma_ab": - return _make_keeps_mma_ab_task_graph( + if cfg.single_kv_pipe: + return _make_single_kv_pipe_task_graph( cfg, total_kv_tiles=total_kv_tiles, use_page_offsets=use_page_offsets, @@ -506,6 +525,9 @@ def build_throughput_latency_mla_task_manager( cta_idx_head_dim_v=cta_idx_head_dim_v, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + sparse_epilogue_params=sparse_epilogue_params, + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -522,7 +544,7 @@ def build_throughput_latency_mla_task_manager( cta_layout_vmnk = (1, 1, 1, 1) tma_group = pipeline.CooperativeGroup(agent.Thread) umma_group = pipeline.CooperativeGroup(agent.Thread) - load_group = pipeline.CooperativeGroup(agent.Thread, 32) + load_group = pipeline.CooperativeGroup(agent.Thread, cfg.load_num_warps * 32) page_group = pipeline.CooperativeGroup(agent.Thread, 32) scheduler_group = pipeline.CooperativeGroup( agent.Thread, @@ -550,13 +572,27 @@ def build_throughput_latency_mla_task_manager( consumer_group=load_group, pipeline_type=PipelineType.AsyncAsync, cta_layout_vmnk=cta_layout_vmnk, - async_producer_op=pipeline.PipelineOp.AsyncLoad, + async_producer_op=( + pipeline.PipelineOp.AsyncThread + if cfg.sparse_direct + else pipeline.PipelineOp.AsyncLoad + ), advance_on_wait=True, ) smem_q_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.q_stages, num_bytes=cfg.qk_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.qk_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, @@ -564,11 +600,23 @@ def build_throughput_latency_mla_task_manager( smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.kv_stages, num_bytes=cfg.kv_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.kv_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, ) + if cfg.sparse_reuse_kv: + smem_kv_cfg = dataclass_replace(smem_kv_cfg, advance_on_wait=True) tmem_s_cfg = PipelineConfig.create_umma_async_pipeline_cfg( num_stages=1, producer_group=umma_group, @@ -674,6 +722,7 @@ def build_throughput_latency_mla_task_manager( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_v=tma_desc_v, + tma_desc_bulk_extra=tma_desc_q_rope, c_rope_tensor=c_rope_tensor, page_offsets_kv=smem_page_offsets if use_page_offsets else None, page_offsets=page_offsets, @@ -699,6 +748,7 @@ def build_throughput_latency_mla_task_manager( cta_idx_kv=cta_idx_kv, sync_barrier_id=0, name="ll_mla_tmem_s0", + page_offsets=page_offsets, ) tmem_s1 = TmemSResource( cfg=cfg, @@ -713,6 +763,7 @@ def build_throughput_latency_mla_task_manager( cta_idx_kv=cta_idx_kv, sync_barrier_id=1, name="ll_mla_tmem_s1", + page_offsets=page_offsets, ) order_p01_alloc = SmemAllocation( name="ll_mla_order_p01", @@ -770,6 +821,9 @@ def build_throughput_latency_mla_task_manager( inst_id=0, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + atten_sinks=sparse_epilogue_params, + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -788,6 +842,9 @@ def build_throughput_latency_mla_task_manager( inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + atten_sinks=sparse_epilogue_params, + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -824,44 +881,51 @@ def build_throughput_latency_mla_task_manager( **task_domain_kwargs, ) ) + tasks.append( + create_throughput_latency_load_task( + smem_q, + smem_kv, + work_queue, + schedule_token_throttle, + cfg, + domain=load_domain, + smem_page_offsets=smem_page_offsets if use_page_offsets else None, + use_page_offsets=use_page_offsets, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ) + ) + if not cfg.short_merged_softmax: + tasks.extend( + [ + create_throughput_latency_softmax0_task( + tmem_s0, + local0, + smem_p0, + global0, + work_queue, + cfg, + domain=softmax_domain, + task_class=MlaDecodeTask, + domain_bias=1, + **task_domain_kwargs, + ), + create_throughput_latency_softmax1_task( + tmem_s1, + local1, + smem_p1, + global1, + work_queue, + cfg, + domain=softmax_domain, + task_class=MlaDecodeTask, + domain_bias=1, + **task_domain_kwargs, + ), + ] + ) tasks.extend( [ - create_throughput_latency_load_task( - smem_q, - smem_kv, - work_queue, - schedule_token_throttle, - cfg, - domain=load_domain, - smem_page_offsets=smem_page_offsets if use_page_offsets else None, - use_page_offsets=use_page_offsets, - task_class=MlaDecodeTask, - **task_domain_kwargs, - ), - create_throughput_latency_softmax0_task( - tmem_s0, - local0, - smem_p0, - global0, - work_queue, - cfg, - domain=softmax_domain, - task_class=MlaDecodeTask, - domain_bias=1, - **task_domain_kwargs, - ), - create_throughput_latency_softmax1_task( - tmem_s1, - local1, - smem_p1, - global1, - work_queue, - cfg, - domain=softmax_domain, - task_class=MlaDecodeTask, - domain_bias=1, - **task_domain_kwargs, - ), create_throughput_latency_correction_task( local0, local1, @@ -918,6 +982,23 @@ def build_throughput_latency_mla_task_manager( ) ) + if cfg.load_warp_idx == 16: + idle = create_padding_task( + cfg, work_queue, warp_idx=15, num_warps=1, task_class=MlaDecodeTask + ) + idle.name = "FormerLoadWarpPadding" + tasks.append(idle) + if cfg.load_num_warps < 4: + idle = create_padding_task( + cfg, + work_queue, + warp_idx=16 + cfg.load_num_warps, + num_warps=4 - cfg.load_num_warps, + task_class=MlaDecodeTask, + ) + idle.name = "LoadGroupPadding" + tasks.append(idle) + deps = { smem_q: [], smem_kv: [], @@ -965,12 +1046,26 @@ def build_throughput_latency_mla_task_manager( deps[schedule_token_throttle] = [work_queue] dma_release_labels = { - (smem_kv, tmem_s0): {"k_desc_0"}, - (smem_kv, tmem_s1): {"k_desc_1"}, - (smem_kv, tmem_o): {"v_desc_0", "v_desc_1"}, + (smem_kv, tmem_s0): {"v_desc_reused" if cfg.sparse_reuse_kv else "k_desc_0"}, + (smem_kv, tmem_s1): {"v_desc_reused" if cfg.sparse_reuse_kv else "k_desc_1"}, + (smem_kv, tmem_o): ( + {"v_desc_reused"} if cfg.sparse_reuse_kv else {"v_desc_0", "v_desc_1"} + ), } + if cfg.serial_swap_reuse: + dma_release_labels = { + (smem_kv, tmem_s0): {"v_desc_reused_0"}, + (smem_kv, tmem_s1): {"v_desc_reused_1"}, + (smem_kv, tmem_o): {"v_desc_reused_0", "v_desc_reused_1"}, + } + smem_allocator = SmemAllocator() + if isinstance(work_queue, ThroughputLatencyMlaClcWorkQueue): + work_queue._response_allocator = smem_allocator + smem_allocator.add_resource(work_queue) + if schedule_token_throttle is not None: + smem_allocator.add_resource(schedule_token_throttle) smem_allocator.add_resource(smem_q) if use_page_offsets: smem_allocator.add_resource(smem_page_offsets) @@ -999,10 +1094,47 @@ def build_throughput_latency_mla_task_manager( tmem_allocator.add_resource(tmem_o) tmem_allocator.compute_layout() + if cfg.short_merged_softmax: + from .tasks import create_short_merged_softmax_task + + # Keep the two P/statistic streams, but create only their shared + # softmax worker. No steady-state correction exchanges are needed. + tasks.append( + create_short_merged_softmax_task( + tmem_s0, + tmem_s1, + local0, + local1, + smem_p0, + smem_p1, + global0, + global1, + cfg, + **task_domain_kwargs, + ) + ) + idle = create_padding_task( + cfg, work_queue, warp_idx=11, num_warps=1, task_class=MlaDecodeTask + ) + idle.name = "MergedSoftmaxPadding" + tasks.append(idle) + if cfg.load_num_warps < 4: + idle = create_padding_task( + cfg, + work_queue, + warp_idx=cfg.load_warp_idx + cfg.load_num_warps, + num_warps=4 - cfg.load_num_warps, + task_class=MlaDecodeTask, + ) + idle.name = "MergedLoadGroupPadding" + tasks.append(idle) task_manager = TaskManager( tasks=tasks, resource_dependency_graph=deps, dma_consumer_release_labels=dma_release_labels, + # DSL 4.7's bounded credit simulation can stop before delayed KV + # releases drain. Its warning mode still runs all schedule checks. + skip_validation=cfg.serial_swap_reuse, smem_allocator=smem_allocator, tmem_allocator=tmem_allocator, verbose=verbose, @@ -1011,7 +1143,7 @@ def build_throughput_latency_mla_task_manager( return task_manager, corr1 -def _make_keeps_mma_ab_task_graph( +def _make_single_kv_pipe_task_graph( cfg: MlaConfig, *, total_kv_tiles: int, @@ -1032,6 +1164,9 @@ def _make_keeps_mma_ab_task_graph( cta_idx_head_dim_v=None, scale_softmax_log2=None, output_scale=None, + sparse_epilogue_params=None, + output_quant=None, + output_quant_args=None, o_tensor=None, lse_tensor=None, acc_o_tensor=None, @@ -1043,10 +1178,10 @@ def _make_keeps_mma_ab_task_graph( verbose: bool = False, exhaustive_deadlock_race_check: bool = False, ) -> tuple[TaskManager, object]: - """Build the keeps-MMA-AB 1CTA task graph for the generic builder. + """Build the single-stream 1CTA task graph. - This variant uses one score/P pipe with P stored in TMEM. The swaps path - keeps its two SMEM-P pipes and is intentionally left untouched. + Keep stores P in TMEM; single-stream swap stores P in SMEM. Both use + one softmax group, two score stages, and one output accumulator. """ scale_softmax_log2, output_scale = _default_scales( @@ -1061,7 +1196,7 @@ def _make_keeps_mma_ab_task_graph( cta_layout_vmnk = (1, 1, 1, 1) tma_group = pipeline.CooperativeGroup(agent.Thread) umma_group = pipeline.CooperativeGroup(agent.Thread) - load_group = pipeline.CooperativeGroup(agent.Thread, 32) + load_group = pipeline.CooperativeGroup(agent.Thread, cfg.load_num_warps * 32) page_group = pipeline.CooperativeGroup(agent.Thread, 32) softmax_group = pipeline.CooperativeGroup( agent.Thread, @@ -1079,13 +1214,27 @@ def _make_keeps_mma_ab_task_graph( consumer_group=load_group, pipeline_type=PipelineType.AsyncAsync, cta_layout_vmnk=cta_layout_vmnk, - async_producer_op=pipeline.PipelineOp.AsyncLoad, + async_producer_op=( + pipeline.PipelineOp.AsyncThread + if cfg.sparse_direct + else pipeline.PipelineOp.AsyncLoad + ), advance_on_wait=True, ) smem_q_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.q_stages, num_bytes=cfg.qk_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.qk_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, @@ -1093,11 +1242,23 @@ def _make_keeps_mma_ab_task_graph( smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.kv_stages, num_bytes=cfg.kv_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.kv_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, ) + if cfg.sparse_reuse_kv: + smem_kv_cfg = dataclass_replace(smem_kv_cfg, advance_on_wait=True) tmem_s_cfg = PipelineConfig.create_umma_async_pipeline_cfg( num_stages=2, producer_group=umma_group, @@ -1111,8 +1272,8 @@ def _make_keeps_mma_ab_task_graph( consumer_group=corr_group, cta_layout_vmnk=cta_layout_vmnk, ) - tmem_p_cfg = PipelineConfig.create_async_umma_pipeline_cfg( - num_stages=2, + p_pipeline_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=1 if cfg.one_insts_kv_swap else 2, producer_group=softmax_group, consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, @@ -1180,6 +1341,7 @@ def _make_keeps_mma_ab_task_graph( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_v=tma_desc_v, + tma_desc_bulk_extra=tma_desc_q_rope, c_rope_tensor=c_rope_tensor, page_offsets_kv=smem_page_offsets if use_page_offsets else None, page_offsets=page_offsets, @@ -1192,8 +1354,9 @@ def _make_keeps_mma_ab_task_graph( cta_idx_head_dim_v=cta_idx_head_dim_v, name="ll_mla_q64_smem_kv", ) - tmem_s = TmemSKeepsResource( + tmem_s = (TmemSResource if cfg.one_insts_kv_swap else TmemSKeepsResource)( cfg=cfg, + page_offsets=page_offsets, pipeline_config=tmem_s_cfg, inst_id=0, scale_softmax_log2=scale_softmax_log2, @@ -1206,26 +1369,47 @@ def _make_keeps_mma_ab_task_graph( sync_barrier_id=0, name="ll_mla_q64_tmem_s", ) - tmem_p = TmemPResource( - cfg=cfg, - pipeline_config=tmem_p_cfg, - inst_id=0, - scale_softmax_log2=scale_softmax_log2, - tmem_alias_ref=tmem_s, - name="ll_mla_q64_tmem_p", - ) - tmem_s.p_ref = tmem_p + # Both one- and two-instance Swap-AB read P as the SMEM B operand. + # This builder's other branch is one-instance Keep-AB: P is the TMEM A + # operand. Stream count alone does not determine probability storage. + if cfg.one_insts_kv_swap: + p_resource = SmemPResource( + cfg=cfg, + pipeline_config=p_pipeline_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + tmem_s_ref=tmem_s, + name="ll_mla_single_smem_p", + ) + else: + p_resource = TmemPResource( + cfg=cfg, + pipeline_config=p_pipeline_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + tmem_alias_ref=tmem_s, + name="ll_mla_q64_tmem_p", + ) + tmem_s.p_ref = p_resource tmem_o = TmemOResource( cfg=cfg, pipeline_config=tmem_o_cfg, - p_tmem_ref=tmem_p, + p_tmem_ref=None if cfg.one_insts_kv_swap else p_resource, + p0_ref=p_resource if cfg.one_insts_kv_swap else None, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, batch_idx=batch_idx, cta_idx_q=cta_idx_q, name="ll_mla_q64_tmem_o", ) - local = TmemSoftmaxLocalResource( + from .resources.tmem_softmax_stats import SmemSoftmaxLocalResource + + local_resource = ( + SmemSoftmaxLocalResource + if cfg.num_tokens_per_page == 1 + else TmemSoftmaxLocalResource + ) + local = local_resource( cfg=cfg, pipeline_config=local_cfg, inst_id=0, @@ -1240,6 +1424,9 @@ def _make_keeps_mma_ab_task_graph( inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + atten_sinks=sparse_epilogue_params, + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -1287,10 +1474,14 @@ def _make_keeps_mma_ab_task_graph( task_class=MlaDecodeTask, **task_domain_kwargs, ), - create_keeps_mma_ab_softmax_task( + ( + create_throughput_latency_softmax0_task + if cfg.one_insts_kv_swap + else create_keeps_mma_ab_softmax_task + )( tmem_s, local, - tmem_p, + p_resource, global_softmax, work_queue, cfg, @@ -1309,11 +1500,11 @@ def _make_keeps_mma_ab_task_graph( task_class=MlaDecodeTask, **task_domain_kwargs, ), - create_keeps_mma_ab_mma_task( + create_single_kv_pipe_mma_task( smem_q, smem_kv, tmem_s, - tmem_p, + p_resource, tmem_o, work_queue, cfg, @@ -1353,14 +1544,31 @@ def _make_keeps_mma_ab_task_graph( ) ) + if cfg.load_warp_idx == 12: + idle = create_padding_task( + cfg, work_queue, warp_idx=9, num_warps=1, task_class=MlaDecodeTask + ) + idle.name = "FormerKeepLoadWarpPadding" + tasks.append(idle) + if cfg.load_num_warps < 4: + idle = create_padding_task( + cfg, + work_queue, + warp_idx=12 + cfg.load_num_warps, + num_warps=4 - cfg.load_num_warps, + task_class=MlaDecodeTask, + ) + idle.name = "KeepLoadGroupPadding" + tasks.append(idle) + deps = { smem_q: [], smem_kv: [], tmem_s: [smem_q, smem_kv], local: [tmem_s], global_softmax: [tmem_s], - tmem_p: [tmem_s], - tmem_o: [tmem_p, smem_kv], + p_resource: [tmem_s], + tmem_o: [p_resource, smem_kv], corr: [local, tmem_o], } if use_page_offsets: @@ -1373,8 +1581,8 @@ def _make_keeps_mma_ab_task_graph( tmem_s: [smem_q, smem_kv, work_queue], local: [tmem_s, work_queue], global_softmax: [tmem_s, work_queue], - tmem_p: [tmem_s, work_queue], - tmem_o: [tmem_p, smem_kv, work_queue], + p_resource: [tmem_s, work_queue], + tmem_o: [p_resource, smem_kv, work_queue], corr: [local, tmem_o, work_queue], work_queue: ( [work_queue, schedule_token_throttle] @@ -1389,15 +1597,32 @@ def _make_keeps_mma_ab_task_graph( deps[schedule_token_throttle] = [work_queue] dma_release_labels = { - (smem_kv, tmem_s): {"k_desc_0"}, - (smem_kv, tmem_o): {"v_desc_0"}, + (smem_kv, tmem_s): { + ( + ("v_desc_reused_0" if cfg.one_insts_kv_swap else "v_desc_reused") + if cfg.sparse_reuse_kv + else "k_desc_0" + ) + }, + (smem_kv, tmem_o): { + ( + ("v_desc_reused_0" if cfg.one_insts_kv_swap else "v_desc_reused") + if cfg.sparse_reuse_kv + else "v_desc_0" + ) + }, } smem_allocator = SmemAllocator() + if isinstance(work_queue, ThroughputLatencyMlaClcWorkQueue): + work_queue._response_allocator = smem_allocator + smem_allocator.add_resource(work_queue) + if schedule_token_throttle is not None: + smem_allocator.add_resource(schedule_token_throttle) smem_allocator.add_resource(smem_q) if use_page_offsets: smem_allocator.add_resource(smem_page_offsets) - for resource in (smem_kv, tmem_s, local, global_softmax, tmem_p, tmem_o, corr): + for resource in (smem_kv, tmem_s, local, global_softmax, p_resource, tmem_o, corr): smem_allocator.add_resource(resource) smem_allocator.add_tmem_ptr( SmemAllocation("ll_mla_q64_tmem_ptr_i32", dtype=Int32, alignment=4) @@ -1414,6 +1639,9 @@ def _make_keeps_mma_ab_task_graph( tasks=tasks, resource_dependency_graph=deps, dma_consumer_release_labels=dma_release_labels, + # See the delayed-release schedule above: keep the stock checks in + # warning mode for BF16 reuse instead of maintaining a second verifier. + skip_validation=cfg.sparse_reuse_kv and not cfg.is_fp8_qkv(), smem_allocator=smem_allocator, tmem_allocator=tmem_allocator, verbose=verbose, @@ -1423,7 +1651,7 @@ def _make_keeps_mma_ab_task_graph( class ThroughputLatencyMlaDecodeTs: - """Dense throughput-latency 1CTA MLA TS wrapper.""" + """Throughput-latency 1CTA MLA over dense pages or prepared sparse rows.""" def __init__( self, @@ -1449,6 +1677,8 @@ def __init__( explicit_split_kv: int | None = None, explicit_persistent: bool | None = None, mask_type: MaskType | str = MaskType.CAUSAL, + sparse_profile=None, + finalize_output: bool = False, ): """Initialize one selected physical tile profile over logical flat Q rows.""" import cutlass as _cutlass @@ -1477,6 +1707,15 @@ def __init__( self.explicit_split_kv = explicit_split_kv self.explicit_persistent = explicit_persistent self.mask_type = normalize_mask_type(mask_type) + self.sparse_profile = sparse_profile + self.finalize_output = finalize_output + self.output_quant = None + # Sparse attention binds separate source lists; dense attention uses + # the ordinary page table. Both reuse the staged KV loaders. + self.direct_sparse = False + self.assume_valid_prefix = False + self.direct_static_scales = False + self.direct_sparse_capacities = (0, 0) cfg = self._make_config() # Parallel standalone reduction requires a fixed split profile and a @@ -1527,18 +1766,36 @@ def __init__( self.acc_dtype = acc_dtype self.lse_dtype = lse_dtype + def _parallel_reduction_setup(self, cfg): + if cfg.fuse_sparse_reduction: + if cfg.logical_seq_len_q != 1: + raise ValueError( + "final sparse reduction requires one virtual query per request" + ) + # Partials already contain every V partition. Finalize one whole + # D512 head per CTA, sharing its statistics across four vec4 warps. + reducer_cfg = dataclass_replace( + cfg, num_ctas_per_head_dim=1, head_dim_per_cta_v=cfg.head_dim_v + ) + topology = make_balanced_parallel_reduction_topology( + cfg.num_ctas_per_seq_kv, cluster_size=1 + ) + return reducer_cfg, topology, cfg.head_dim_v + return ( + cfg, + self.parallel_reduction_topology, + self.parallel_reduction_elements_per_slice, + ) + def compile_topology_signature(self) -> tuple[object, ...]: """Describe batch-derived reducer choices without retaining batch.""" cfg = self._make_config() if cfg.use_multi_ctas_kv != 1 or cfg.use_cluster_reduction == 1: return ("no_separate_reducer",) - if self.use_parallel_reduction: - return ( - "parallel", - self.parallel_reduction_topology, - self.parallel_reduction_elements_per_slice, - ) + if self.use_parallel_reduction or cfg.fuse_sparse_reduction: + _, topology, elements = self._parallel_reduction_setup(cfg) + return ("parallel", topology, elements) grid, smem, threads, reduction_ctas = gmem_reduction_launch_shape( cfg, cfg.seq_len_q, @@ -1567,6 +1824,7 @@ def compile_signature(self) -> tuple[object, ...]: self.parallel_reduction_topology, self.parallel_reduction_elements_per_slice, self.compile_topology_signature(), + self.output_quant, ) def _make_config(self): @@ -1592,6 +1850,16 @@ def _make_config(self): explicit_persistent=self.explicit_persistent, mask_type=self.mask_type, ) + if self.sparse_profile is not None: + cfg = configure_sparse_mla( + cfg, + self.sparse_profile, + finalize_output=self.finalize_output, + direct_inputs=self.direct_sparse, + assume_valid_prefix=self.assume_valid_prefix, + static_scales=self.direct_static_scales, + source_capacities=self.direct_sparse_capacities, + ) return cfg def validate_split_kv_launch(self, split_kv: int, workspace) -> None: @@ -1662,10 +1930,11 @@ def __call__( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, stream: object, + output_quant_args=None, ): """Execute throughput-latency 1CTA MLA with fixed or compact ragged Q.""" cfg = self._make_config() @@ -1775,75 +2044,133 @@ def __call__( l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - q_rope_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): - q_rope_swizzle = cuda.TensorMapSwizzle.s64b - if cutlass.const_expr(cu_seqlens_q is not None): - q_rope_tma = cute.make_tensor( - q_rope.iterator, + if cutlass.const_expr(cfg.rope_dim > 0): + q_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + q_rope_swizzle = cuda.TensorMapSwizzle.s64b + if cutlass.const_expr(cu_seqlens_q is not None): + q_rope_tma = cute.make_tensor( + q_rope.iterator, + cute.make_layout( + (q_rope.shape[1], q_rope.shape[0] * q_rope.shape[2]), + stride=(q_rope.stride[1], q_rope.stride[0]), + ), + ) + tma_desc_q_rope = create_tensor_map_ragged_from_tensor( + q_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q), + ragged_dim=1, + stride_order=(0, 1), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + q_rope_tma = cute.make_tensor( + q_rope.iterator, + cute.make_layout( + ( + q_rope.shape[1], + q_rope.shape[0] * q_rope.shape[2], + q_rope.shape[3], + ), + stride=( + q_rope.stride[1], + q_rope.stride[0], + q_rope.stride[3], + ), + ), + ) + tma_desc_q_rope = create_tensor_map_tiled_from_view( + q_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q, 1), + stride_order=(0, 1, 2), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_q_rope = tma_desc_q_latent + + c_latent_tma = cute.make_tensor( + c_latent.iterator, + cute.select(c_latent.layout, mode=[1, 0, 2]), + ) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + # Sparse routes address individual rows; the unused D512 RoPE + # descriptor slot carries the independent compressed-pool map. + c_rows = cute.make_tensor( + c_latent.iterator, cute.make_layout( - (q_rope.shape[1], q_rope.shape[0] * q_rope.shape[2]), - stride=(q_rope.stride[1], q_rope.stride[0]), + (c_latent.shape[1], c_latent.shape[2]), + stride=(c_latent.stride[1], c_latent.stride[2]), ), ) - tma_desc_q_rope = create_tensor_map_ragged_from_tensor( - q_rope_tma, - box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q), - ragged_dim=1, + c_extra_rows = cute.make_tensor( + c_rope.iterator, + cute.make_layout( + (c_rope.shape[1], c_rope.shape[2]), + stride=(c_rope.stride[1], c_rope.stride[2]), + ), + ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_rows, + box_dims=(tma_box0, 1), stride_order=(0, 1), - swizzle=q_rope_swizzle, + swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - else: - q_rope_tma = cute.make_tensor( - q_rope.iterator, - cute.make_layout( - ( - q_rope.shape[1], - q_rope.shape[0] * q_rope.shape[2], - q_rope.shape[3], - ), - stride=( - q_rope.stride[1], - q_rope.stride[0], - q_rope.stride[3], - ), - ), + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_extra_rows, + box_dims=(tma_box0, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - tma_desc_q_rope = create_tensor_map_tiled_from_view( - q_rope_tma, - box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q, 1), + else: + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_latent_tma, + box_dims=(tma_box0, tma_page_tokens, 1), stride_order=(0, 1, 2), - swizzle=q_rope_swizzle, + swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - c_latent_tma = cute.make_tensor( - c_latent.iterator, - cute.select(c_latent.layout, mode=[1, 0, 2]), - ) - tma_desc_c_latent = create_tensor_map_tiled_from_view( - c_latent_tma, - box_dims=(tma_box0, tma_page_tokens, 1), - stride_order=(0, 1, 2), - swizzle=cuda.TensorMapSwizzle.s128b, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) - - c_rope_tma = cute.make_tensor( - c_rope.iterator, - cute.select(c_rope.layout, mode=[1, 0, 2]), - ) - c_rope_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): - c_rope_swizzle = cuda.TensorMapSwizzle.s64b - tma_desc_c_rope = create_tensor_map_tiled_from_view( - c_rope_tma, - box_dims=(min(tma_box0, cfg.rope_dim), tma_page_tokens, 1), - stride_order=(0, 1, 2), - swizzle=c_rope_swizzle, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) + if cutlass.const_expr(cfg.rope_dim > 0): + c_rope_tma = cute.make_tensor( + c_rope.iterator, + cute.select(c_rope.layout, mode=[1, 0, 2]), + ) + c_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + c_rope_swizzle = cuda.TensorMapSwizzle.s64b + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), tma_page_tokens, 1), + stride_order=(0, 1, 2), + swizzle=c_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + elif cutlass.const_expr(cfg.num_tokens_per_page != 1): + tma_desc_c_rope = tma_desc_c_latent + + tma_desc_v = tma_desc_c_latent + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + # D512 does not use Q-rope; reuse its descriptor slot for the + # second pool's contiguous-run map. Gather maps stay separate. + rows_per_issuer = 32 + tma_desc_v = create_tensor_map_tiled_from_view( + c_rows, + box_dims=(tma_box0, rows_per_issuer), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + tma_desc_q_rope = create_tensor_map_tiled_from_view( + c_extra_rows, + box_dims=(tma_box0, rows_per_issuer), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) softmax_scale_log2 = softmax_scale * LOG2_E use_gmem_reduction = cutlass.const_expr( @@ -1911,7 +2238,7 @@ def __call__( tma_desc_q_rope, tma_desc_c_latent, tma_desc_c_rope, - tma_desc_c_latent, + tma_desc_v, c_rope, page_offsets, o, @@ -1923,6 +2250,8 @@ def __call__( softmax_scale_log2, output_scale, tile_sched_params, + block_split_kvs, + output_quant_args, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1932,13 +2261,18 @@ def __call__( use_pdl=acc_o is not None, ) if cutlass.const_expr(acc_o is not None): - if cutlass.const_expr(self.use_parallel_reduction): - topology = self.parallel_reduction_topology + if cutlass.const_expr( + self.use_parallel_reduction or cfg.fuse_sparse_reduction + ): + reducer_cfg, topology, elements = self._parallel_reduction_setup(cfg) + atten_sinks = None + if cutlass.const_expr(cfg.fuse_sparse_reduction): + atten_sinks = block_split_kvs[-1] reduction_grid, reduction_cluster = ( parallel_gmem_reduction_launch_shape( - cfg, + reducer_cfg, topology, - self.parallel_reduction_elements_per_slice, + elements, ) ) reduction_grid = ( @@ -1946,9 +2280,9 @@ def __call__( reduction_grid[1], batch_size, ) - reduction_threads = parallel_gmem_reduction_threads( - self.parallel_reduction_elements_per_slice - ) + if cutlass.const_expr(cfg.fuse_sparse_reduction): + reduction_grid = (cfg.logical_num_heads_q, 1, batch_size) + reduction_threads = parallel_gmem_reduction_threads(elements) parallel_reducer = self.parallel_gmem_reduction_kernel( o, lse, @@ -1956,6 +2290,8 @@ def __call__( acc_lse, cache_seqs, cu_seqlens_q, + atten_sinks, + output_quant_args, ) if cutlass.const_expr(topology.cluster_size == 1): parallel_reducer.launch( @@ -2028,10 +2364,37 @@ def dense_kernel( softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, tile_sched_params: object, + sparse_inputs: object, + output_quant_args=None, ): """Execute one flat-Q, batch, KV-split, and V head-dimension tile.""" cfg = self._make_config() + if cutlass.const_expr(cfg.sparse_direct): + from ..helpers.sparse_indices import bind_sparse_kv_indices + + si, ci, sl, cl, sm, os, sinks = sparse_inputs + request = None + if cutlass.const_expr(cfg.use_persistent_scheduler != 1): + request = cutlass.Int64(cute.arch.block_idx()[2]) + page_offsets, cache_seqs = bind_sparse_kv_indices( + si, + ci, + sl, + cl, + cfg.total_kv_tiles * cfg.tile_size_kv, + query_row=request, + assume_valid_prefix=cfg.sparse_valid_prefix, + has_extra=cfg.sparse_direct_capacities[1] > 0, + ) + if cutlass.const_expr(not cfg.sparse_static_scales): + # Caller scales already include Q/K/V descales. + softmax_scale_log2 = cutlass.Float32(sm[0]) * LOG2_E + output_scale = cutlass.Float32(os[0]) + sparse_epilogue_params = sinks + else: + sparse_epilogue_params = None + # The grid is expressed in physical flat-Q tile coordinates. Decode it # once here; Q/O resources retain responsibility for logical-row # mapping, compact-ragged offsets, and padded-tail publication. @@ -2105,10 +2468,17 @@ def dense_kernel( prims.prefetch_tensormap(tma_desc_c_latent.get_ptr()) prims.prefetch_tensormap(tma_desc_c_rope.get_ptr()) prims.prefetch_tensormap(tma_desc_v.get_ptr()) + if cutlass.const_expr(cfg.sparse_direct): + lane = cute.arch.thread_idx()[0] % 32 + if lane < (cfg.logical_num_heads_q + 31) // 32: + address = sparse_epilogue_params.iterator.toint() + cutlass.Int64( + lane * 128 + ) + cute.arch.inline_ptx( + "prefetch.global.L2 [{$r0}];", read_only_args=[address] + ) clc_response_ptr = None - if cutlass.const_expr(use_clc_dynamic): - clc_response_ptr = cute.arch.alloc_smem(cutlass.Int128, 2) task_manager, cluster_corr_resource = build_throughput_latency_mla_task_manager( cfg, @@ -2130,6 +2500,11 @@ def dense_kernel( cta_idx_head_dim_v=cta_idx_head_dim_v, scale_softmax_log2=softmax_scale_log2, output_scale=output_scale, + sparse_epilogue_params=sparse_epilogue_params, + output_quant=self.output_quant + if (cfg.fuse_sparse_epilogue or cfg.fuse_sparse_cluster_epilogue) + else None, + output_quant_args=output_quant_args, o_tensor=o, lse_tensor=lse, acc_o_tensor=acc_o, @@ -2170,10 +2545,27 @@ def dense_kernel( ) prims.tcgen05_relinquish_alloc_permit(group=prims.CTAGroup.CTA_1) - prims.barrier_cta_sync( - barrier_id=TMEM_LIFECYCLE_BARRIER_ID, - thread_count=cfg.threads_per_cta, - ) + if cutlass.const_expr(cfg.sparse_direct): + # Loaders need the initialized pipelines, but not TMEM allocation. + # A distinct barrier keeps an early-finishing loader's final full- + # CTA rendezvous from joining the compute-only readiness phase. + if warp_idx < Int32( + cfg.correction_warp_idx + cfg.correction_num_warps + ) or warp_idx == Int32(cfg.mma_warp_idx): + prims.barrier_cta_sync( + barrier_id=TMEM_READY_BARRIER_ID, + thread_count=( + cfg.correction_warp_idx + + cfg.correction_num_warps + + cfg.mma_num_warps + ) + * 32, + ) + else: + prims.barrier_cta_sync( + barrier_id=TMEM_LIFECYCLE_BARRIER_ID, + thread_count=cfg.threads_per_cta, + ) if cutlass.const_expr(has_runtime_activity_guard): if run_task_graph: task_manager.run() @@ -2263,11 +2655,13 @@ def parallel_gmem_reduction_kernel( acc_lse: cute.Tensor, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor, + atten_sinks: cute.Tensor | None = None, + output_quant_args=None, ): """Dispatch the automatically selected parallel standalone reducer.""" cfg = self._make_config() - topology = self.parallel_reduction_topology + cfg, topology, elements = self._parallel_reduction_setup(cfg) run_parallel_gmem_reduction_kernel( output, lse, @@ -2279,5 +2673,9 @@ def parallel_gmem_reduction_kernel( topology.cluster_size, topology.slots_per_rank, topology.actual_splits, - self.parallel_reduction_elements_per_slice, + elements, + atten_sinks, + cfg.fuse_sparse_reduction, + self.output_quant, + output_quant_args, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py index 11d39bdec7e6..11a40b52bfd1 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py @@ -19,8 +19,12 @@ from cutlass import Float32, Int32, Int64 from cutlass.experimental import primitives as prims -from ...separate_reduction import finalize_log2_sum_exp, normalized_lse_weight -from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ...separate_reduction import ( + finalize_log2_sum_exp, + normalized_lse_weight, + attention_sink_scale, +) +from ..helpers.constants import LN_2, SPLIT_REDUCTION_SCALE_BARRIER_ID from ..helpers.math import ceil_div from ..helpers.ops import ( fmax_f32, @@ -197,6 +201,10 @@ def _run_parallel_gmem_reduction_g1_shared_stats( cu_seqlens_q, cfg, elements_per_slice: cutlass.Constexpr[int], + atten_sinks=None, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Reduce G1 partials with the compact row-shared schedule.""" @@ -204,8 +212,15 @@ def _run_parallel_gmem_reduction_g1_shared_stats( thread_idx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) lane_idx = thread_idx % Int32(GMEM_REDUCTION_WARP_LANES) - head_dim_cta_idx = block_idx_y % Int32(cfg.num_ctas_per_head_dim) - q_idx = block_idx_y // Int32(cfg.num_ctas_per_head_dim) + if cutlass.const_expr(finalize_attention): + # The final sparse grid is one CTA per logical head. Decode its + # physical producer tile without launching padded head/V slices. + q_idx = slice_idx // Int32(cfg.tile_size_q) + slice_idx = slice_idx % Int32(cfg.tile_size_q) + head_dim_cta_idx = Int32(0) + else: + head_dim_cta_idx = block_idx_y % Int32(cfg.num_ctas_per_head_dim) + q_idx = block_idx_y // Int32(cfg.num_ctas_per_head_dim) elements_per_slice = _parallel_reduction_effective_slice_elements( cfg, elements_per_slice ) @@ -271,6 +286,11 @@ def _run_parallel_gmem_reduction_g1_shared_stats( ) lse_sum = warp_reduce_sum_f32(lse_sum) global_lse = finalize_log2_sum_exp(lse_max, lse_sum) + stats_sink_scale = Float32(1) + if cutlass.const_expr(finalize_attention and atten_sinks is not None): + stats_sink_scale = attention_sink_scale( + global_lse, Float32(atten_sinks[stats_query_row[1]]) + ) if ( lane_idx == Int32(0) and head_dim_cta_idx == Int32(0) @@ -282,16 +302,23 @@ def _run_parallel_gmem_reduction_g1_shared_stats( batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + stats_output_query_row).store(global_lse) + (lse.iterator.raw_ptr() + stats_output_query_row).store( + global_lse * Float32(LN_2) if finalize_attention else global_lse + ) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): stats_split_idx = lane_idx + Int32( lane_slot_i * GMEM_REDUCTION_WARP_LANES ) if stats_split_idx < Int32(cfg.num_ctas_per_seq_kv): + stats_weight = normalized_lse_weight( + lane_lse[lane_slot_i], global_lse + ) + if cutlass.const_expr(finalize_attention): + stats_weight *= stats_sink_scale smem_scale[ stats_row * Int32(cfg.num_ctas_per_seq_kv) + stats_split_idx - ] = normalized_lse_weight(lane_lse[lane_slot_i], global_lse) + ] = stats_weight prims.barrier_cta_sync( barrier_id=SPLIT_REDUCTION_SCALE_BARRIER_ID, @@ -342,16 +369,20 @@ def _run_parallel_gmem_reduction_g1_shared_stats( scale = Float32( smem_scale[row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx] ) - split_ptr = acc_row_ptr + Int64(split_i * cfg.head_dim_v) - if cutlass.const_expr(output_elements_per_thread == 1): - output_acc[0] += Float32(split_ptr.load()) * scale - else: - partial_o = split_ptr.load( - count=output_elements_per_thread, - alignment=output_elements_per_thread * cfg.partial_o_dtype_bytes, - ).to(Float32) - for elem_i in cutlass.range_constexpr(output_elements_per_thread): - output_acc[elem_i] += partial_o[elem_i] * scale + # Sparse masked splits can leave their vector unused. Do not + # multiply an uninitialized/NaN partial by a zero statistical weight. + if not finalize_attention or scale > Float32(0): + split_ptr = acc_row_ptr + Int64(split_i * cfg.head_dim_v) + if cutlass.const_expr(output_elements_per_thread == 1): + output_acc[0] += Float32(split_ptr.load()) * scale + else: + partial_o = split_ptr.load( + count=output_elements_per_thread, + alignment=output_elements_per_thread + * cfg.partial_o_dtype_bytes, + ).to(Float32) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + output_acc[elem_i] += partial_o[elem_i] * scale output_query_row = public_query_flat_row( cfg, @@ -360,7 +391,9 @@ def _run_parallel_gmem_reduction_g1_shared_stats( cu_seqlens_q, ) out_elem_offset = _output_element_offset(cfg, output_query_row, dim_idx) - if cutlass.const_expr(output_elements_per_thread == 1): + if cutlass.const_expr(output_quant is not None): + output_quant.store(output_quant_args, output_acc, output_query_row, dim_idx) + elif cutlass.const_expr(output_elements_per_thread == 1): (output.iterator.raw_ptr() + out_elem_offset).store( output.element_type(output_acc[0]) ) @@ -769,9 +802,15 @@ def run_parallel_gmem_reduction_kernel( slots_per_rank: cutlass.Constexpr[int], actual_splits: cutlass.Constexpr[int], elements_per_slice: cutlass.Constexpr[int], + atten_sinks=None, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Reduce split-KV partials with row-shared local and peer statistics.""" + if cutlass.const_expr(finalize_attention and cluster_size != 1): + raise ValueError("final attention output uses a single-CTA reducer per row") prims.griddepcontrol(kind=prims.GridDepAction.WAIT) if cutlass.const_expr(cluster_size == 1): @@ -784,6 +823,10 @@ def run_parallel_gmem_reduction_kernel( cu_seqlens_q, cfg, elements_per_slice, + atten_sinks, + finalize_attention, + output_quant, + output_quant_args, ) return diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py index deae8aa5fed2..14fd5758621b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py @@ -194,7 +194,10 @@ def _init_order_p01_barriers(self): def _ordered_sequence_wait(self): """Wait for this P instance's ordered materialization turn.""" - if cutlass.const_expr(self.cfg.use_clc_dynamic_persistent_scheduler == 1): + if cutlass.const_expr( + self.cfg.use_clc_dynamic_persistent_scheduler == 1 + and self.cfg.num_insts_kv > 1 + ): cute.arch.mbarrier_wait( self._order_p01_barrier_ptr + self.inst_id, self._order_p01_phase, @@ -204,7 +207,10 @@ def _ordered_sequence_wait(self): def _ordered_sequence_arrive(self): """Signal the peer P instance after materialization completes.""" - if cutlass.const_expr(self.cfg.use_clc_dynamic_persistent_scheduler == 1): + if cutlass.const_expr( + self.cfg.use_clc_dynamic_persistent_scheduler == 1 + and self.cfg.num_insts_kv > 1 + ): signaling_id = 1 if self.inst_id == 0 else 0 cute.arch.mbarrier_arrive( self._order_p01_barrier_ptr + signaling_id, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py index 8e2544d61276..590173465f44 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py @@ -35,6 +35,8 @@ producer_work, ) +from ...helpers.gather import gather4, gather4_cached_warp, select_gather_map + from ...helpers.constants import ( CP_ASYNC_CACHE_CA, PAGE_OFFSET_BYTES, @@ -230,7 +232,11 @@ def _load_q_stage( local_flat_query_row, batch_idx, ) - if prims.elect_sync(): + if prims.elect_sync() and ( + self.cfg.load_num_warps == 1 + or cute.arch.thread_idx()[0] // 32 - self.cfg.load_warp_idx + == qk_stage_idx % self.cfg.load_num_warps + ): smem_offset = Int32(q_stage_smem_element_offset(self.cfg, qk_stage_idx)) prims.cp_async_bulk_tensor_shared_cta_global( stage_base.data_ptr(smem_offset), @@ -276,13 +282,14 @@ def load_q(self, stage_info: StageInfo): dim_offset, self.tma_desc_q_latent, ) - self._load_q_stage( - stage_info, - stage_base, - latent_stages, - Int32(0), - self.tma_desc_q_rope, - ) + if cutlass.const_expr(cfg.rope_dim > 0): + self._load_q_stage( + stage_info, + stage_base, + latent_stages, + Int32(0), + self.tma_desc_q_rope, + ) @consumer_work(returns=("q_desc", "q_desc_rope")) @cute.jit @@ -479,7 +486,15 @@ def init_read_state(self, stage_info: StageInfo): cache_slots = 2 if self.cfg.kernel_variant != "keeps_mma_ab" else 1 return cutlass.Array( Int32, - max(1, cache_slots * self.cfg.pages_per_kv_tile), + max( + 1, + cache_slots + * ( + self.cfg.sparse_cache_words + if self.cfg.num_tokens_per_page == 1 + else self.cfg.pages_per_kv_tile + ), + ), space=cutlass.AddressSpace.rmem, ) @@ -506,7 +521,7 @@ def _producer_load_page_offsets( cfg = self.cfg lane_idx = lane_idx_from_thread(cute.arch.thread_idx()[0]) pages_per_tile = Int32(cfg.pages_per_kv_tile) - if lane_idx < pages_per_tile: + if cutlass.const_expr(cfg.sparse_direct): batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) @@ -521,31 +536,73 @@ def _producer_load_page_offsets( cfg, stage_info, inst_id, is_v, section=section ) tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) - last_valid_page = cute.math.max( - (seq_len_kv + Int32(cfg.num_tokens_per_page - 1)) - // Int32(cfg.num_tokens_per_page) - - Int32(1), - Int32(0), - ) - logical_page_idx = cute.math.min( - tile_idx * pages_per_tile + lane_idx, - last_valid_page, - ) + tile_start = tile_idx * pages_per_tile + routes = self.page_offsets.select_source_tile(tile_start, batch_idx) + position = lane_idx * Int32(4) smem_offset = ( stage_info.stage_idx * Int32(cfg.page_offsets_entries_per_stage) - + lane_idx - ) - page_offsets_batch = self.page_offsets[None, batch_idx] - page_offsets_flat = cute.flat_divide(page_offsets_batch, (1,)) - gmem_ptr = page_offsets_flat[None, logical_page_idx].iterator.llvm_ptr - smem_ptr = cutlass.inttoptr( - self._smem_page_offsets.data_ptr(smem_offset).toint(cutlass.Int32), - 3, - cutlass.Int32, - ) - prims.cp_async_shared_global( - smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA + + position ) + if position < pages_per_tile: + # Source selection is tile-wide. Hoist the bounds decision out + # of individual quads so full tiles take one uniform branch; + # index values still mask holes even in an in-bounds tile. + if tile_start + pages_per_tile <= routes.origin + routes.length: + quad = routes.load_full_quad(tile_start + position) + self._smem_page_offsets.data_ptr(smem_offset).store( + quad.load(0, 4), alignment=16 + ) + else: + quad = routes.load_quad(tile_start + position) + self._smem_page_offsets.data_ptr(smem_offset).store( + quad.load(0, 4), alignment=16 + ) + cute.arch.fence_view_async_shared() + return + for fragment in cutlass.range_constexpr((cfg.pages_per_kv_tile + 31) // 32): + fragment_lane = lane_idx + Int32(fragment * 32) + if fragment_lane < pages_per_tile: + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + cfg, + decode_gen_task_cache(stage_info), + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + local_tile_idx = local_kv_tile_idx( + cfg, stage_info, inst_id, is_v, section=section + ) + tile_idx = global_kv_tile_idx( + cfg, local_tile_idx, seq_len_kv, cta_idx_kv + ) + last_valid_page = cute.math.max( + (seq_len_kv + Int32(cfg.num_tokens_per_page - 1)) + // Int32(cfg.num_tokens_per_page) + - Int32(1), + Int32(0), + ) + logical_page_idx = cute.math.min( + tile_idx * pages_per_tile + fragment_lane, + last_valid_page, + ) + smem_offset = ( + stage_info.stage_idx * Int32(cfg.page_offsets_entries_per_stage) + + fragment_lane + ) + page_offsets_batch = self.page_offsets[None, batch_idx] + page_offsets_flat = cute.flat_divide(page_offsets_batch, (1,)) + gmem_ptr = page_offsets_flat[None, logical_page_idx].iterator.llvm_ptr + smem_ptr = cutlass.inttoptr( + self._smem_page_offsets.data_ptr(smem_offset).toint(cutlass.Int32), + 3, + cutlass.Int32, + ) + prims.cp_async_shared_global( + smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA + ) @producer_work @cute.jit @@ -592,9 +649,90 @@ def read_offsets( # The load task waits on the page-offset stage, snapshots the page ids # into registers, and keeps swaps K0/K1 live until their delayed V use. del stage_info - cache_base = cache_slot * self.cfg.pages_per_kv_tile - for page_frag in cutlass.range_constexpr(self.cfg.pages_per_kv_tile): - cached_page_ids[cache_base + page_frag] = self.page_id(page_frag) + if cutlass.const_expr(self.cfg.num_tokens_per_page == 1): + lane = lane_idx_from_thread(cute.arch.thread_idx()[0]) + warp = cute.arch.thread_idx()[0] // 32 - self.cfg.load_warp_idx + groups_per_warp = self.cfg.tile_size_kv // (4 * self.cfg.load_num_warps) + for j in cutlass.range_constexpr(4): + cached_page_ids[cache_slot * self.cfg.sparse_cache_words + j] = Int32( + 0x7FFFFFFF + ) + if lane < groups_per_warp: + quad = warp * groups_per_warp + lane + offset = ( + self.consumer_work_stage * self.cfg.page_offsets_entries_per_stage + + quad * 4 + ) + values = self._smem_page_offsets.data_ptr(offset).load( + count=4, alignment=16 + ) + for j in cutlass.range_constexpr(4): + cached_page_ids[cache_slot * self.cfg.sparse_cache_words + j] = ( + values[j] + ) + base = cache_slot * self.cfg.sparse_cache_words + # Actual physical adjacency within one source is required. A + # full valid prefix can still contain random rows. Eight quads + # (32 rows) are needed for bulk; smaller issuer shares use gather4. + if cutlass.const_expr(groups_per_warp >= 8): + # A chunk stays within one source. Invalid coordinates retain + # its tag because the first coordinate selects the shared map. + first = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_page_ids[base], + offset=lane & Int32(-8), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + contiguous = (first & Int32(0x7FFFFFFF)) != Int32(0x7FFFFFFF) + empty = cutlass.Boolean(True) + for j in cutlass.range_constexpr(4): + raw = cached_page_ids[base + j] + contiguous = contiguous and raw == first + (lane % 8) * 4 + j + empty = empty and (raw & Int32(0x7FFFFFFF)) == Int32(0x7FFFFFFF) + cached_page_ids[base + 4] = Int32( + cute.arch.vote_ballot_sync(contiguous) + ) + cached_page_ids[base + 5] = Int32(cute.arch.vote_ballot_sync(empty)) + else: + cached_page_ids[base + 4] = Int32(0) + cached_page_ids[base + 5] = Int32(0) + + else: + cache_base = cache_slot * self.cfg.pages_per_kv_tile + for page_frag in cutlass.range_constexpr(self.cfg.pages_per_kv_tile): + cached_page_ids[cache_base + page_frag] = self.page_id(page_frag) + if cutlass.const_expr(self.cfg.cache_uniform_sparse_quads): + # Stage once per token tile, then reuse these uniform coordinates + # across all four head-dimension stages. Keep the raw lane cache + # for the contiguous/empty bulk path and source selection. + base = cache_slot * self.cfg.sparse_cache_words + bulk = (cached_page_ids[base + 4] & Int32(255)) == Int32(255) + empty = (cached_page_ids[base + 5] & Int32(255)) == Int32(255) + if self.cfg.load_num_warps == 8 or not (bulk or empty): + for owner in cutlass.range_constexpr( + self.cfg.tile_size_kv // (4 * self.cfg.load_num_warps) + ): + for j in cutlass.range_constexpr(4): + raw = prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_page_ids[base + j], + offset=owner, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + row = raw & Int32(0x7FFFFFFF) + row = row if row != Int32(0x7FFFFFFF) else Int32(-1) + cached_page_ids[base + 6 + owner * 4 + j] = ( + cute.arch.make_warp_uniform(row) + ) + else: + # No gather reads these coordinates on a bulk/empty tile. + # Constants avoid carrying the previous tile's uniform cache. + for j in cutlass.range_constexpr(self.cfg.sparse_cache_words - 6): + cached_page_ids[base + 6 + j] = Int32(0) return cached_page_ids @@ -633,6 +771,7 @@ class SmemKvResource(MlaResource): tma_desc_c_latent: object = None tma_desc_c_rope: object = None tma_desc_v: object = None + tma_desc_bulk_extra: object = None c_rope_tensor: object = None page_offsets_kv: object = None page_offsets: object = None @@ -800,7 +939,11 @@ def _producer_load_kv_stage( dim_offset = head_dim_cta_offset_v(cfg, cta_idx_head_dim_v) + Int32( qk_stage_idx * cfg.head_dim_per_stage_v ) - tma_desc = self.tma_desc_v + tma_desc = ( + self.tma_desc_c_latent + if cfg.num_tokens_per_page == 1 + else self.tma_desc_v + ) else: active_width = cfg.qk_head_stage_width(qk_stage_idx) if cutlass.const_expr( @@ -839,6 +982,101 @@ def _producer_load_kv_stage( and cfg.head_dim_per_stage_kv == 128 ) + # Each lane owns four adjacent index-list entries. Bulk TMA requires + # a verified 32-row physical run within one source; irregular rows + # use gather4. The separate empty-window path only zero-fills padding. + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + warp = cute.arch.thread_idx()[0] // 32 - cfg.load_warp_idx + groups_per_warp = cfg.tile_size_kv // (4 * cfg.load_num_warps) + base = page_id_slot * cfg.sparse_cache_words + for chunk in cutlass.range_constexpr( + max(1, cfg.tile_size_kv // (32 * cfg.load_num_warps)) + ): + mask = Int32(255) << Int32(chunk * 8) + contiguous = (cached_page_ids[base + 4] & mask) == mask + empty = (cached_page_ids[base + 5] & mask) == mask + first = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_page_ids[base], + offset=chunk * 8, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + if (contiguous or empty) and groups_per_warp >= 8: + row = Int32(-32) if empty else first & Int32(0x7FFFFFFF) + desc = select_gather_map( + self.tma_desc_v, self.tma_desc_bulk_extra, first < 0 + ) + if prims.elect_sync(): + for part in cutlass.range_constexpr( + cfg.head_dim_per_stage_kv // (128 // cfg.qkv_dtype_bytes) + ): + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr( + (warp * groups_per_warp + chunk * 8) + * 4 + * inner_width + + part * cfg.tile_size_kv * inner_width + ), + cute.make_ptr(cutlass.Int8, desc), + (dim_offset + part * inner_width, row), + stage_info.barrier, + ) + elif cutlass.const_expr(cfg.cache_uniform_sparse_quads): + desc = select_gather_map(tma_desc, self.tma_desc_c_rope, first < 0) + gather4_cached_warp( + stage_base.data_ptr(warp * groups_per_warp * 4 * inner_width), + desc, + dim_offset, + cached_page_ids, + stage_info.barrier, + cache_start=base + 6, + num_quads=min(8, groups_per_warp), + num_slices=cfg.head_dim_per_stage_kv + // (128 // cfg.qkv_dtype_bytes), + quad_stride=4 * inner_width * cfg.qkv_dtype_bytes, + slice_stride=cfg.tile_size_kv + * inner_width + * cfg.qkv_dtype_bytes, + col_stride=inner_width, + ) + else: + # Each active lane already owns a register-cached quad. + # Issue disjoint TMA transactions directly from those + # lanes, avoiding eight rounds of shuffles/election. + lane = cute.arch.thread_idx()[0] % 32 + if lane >= chunk * 8 and lane < min(chunk * 8 + 8, groups_per_warp): + group = warp * groups_per_warp + lane + coords = cutlass.Array( + Int32, 4, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(4): + raw = cached_page_ids[base + j] + row = raw & Int32(0x7FFFFFFF) + coords[j] = row if row != Int32(0x7FFFFFFF) else Int32(-1) + desc = select_gather_map( + tma_desc, self.tma_desc_c_rope, first < 0 + ) + for part in cutlass.range_constexpr( + cfg.head_dim_per_stage_kv // (128 // cfg.qkv_dtype_bytes) + ): + gather4( + stage_base.data_ptr( + group * 4 * inner_width + + part * cfg.tile_size_kv * inner_width + ), + desc, + dim_offset + part * inner_width, + coords[0], + coords[1], + coords[2], + coords[3], + stage_info.barrier, + ) + return + if cutlass.const_expr(cfg.use_paged_kv == 1): if cutlass.const_expr( self.page_offsets_kv is None and self.page_offsets is None @@ -1059,6 +1297,7 @@ def _set_desc( inst_id: int = 0, *, k_subtile_idx: cutlass.Constexpr[int] = 0, + reuse_k: cutlass.Constexpr[bool] = False, ): """Build a K or V SMEM descriptor for the current producer stage.""" # Descriptor offsets mirror _init_smem_state, but are rebuilt against the @@ -1070,22 +1309,30 @@ def _set_desc( qkv_major_k_stride_bytes_for(self.cfg, self.cfg.head_dim_per_stage_kv) ) if cutlass.const_expr(is_v): + stage_base = self._stage_base(stage_info) + if cutlass.const_expr(reuse_k): + # K waits advance independently of releases. Keep its payload + # live until the corresponding delayed PV has consumed it. + stage_base = self._smem_kv.subview( + self.consumer_release_state.index + * Int32(self.cfg.kv_smem_stage_elements) + ) v_leading_byte_offset = ( Int32(0) if self.cfg.head_dim_per_stage_v == TCGEN05_BF16_K_BLOCK_WIDTH - else Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + else Int32(self.cfg.tile_size_kv * 128) ) if cutlass.const_expr(self.cfg.is_fp8_qkv()): v_leading_byte_offset = Int32(0) desc = cprims.Tcgen05SmemDesc.build( - self._stage_base(stage_info), + stage_base, leading_byte_offset=v_leading_byte_offset, stride_byte_offset=stride_byte_offset, layout=qkv_smem_swizzle(self.cfg), ) return desc else: - k_leading_byte_offset = Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + k_leading_byte_offset = Int32(self.cfg.tile_size_kv * 128) if cutlass.const_expr(self.cfg.is_fp8_qkv()): k_leading_byte_offset = Int32(self.cfg.kv_smem_tile_bytes) qk_stage_idx = k_subtile_idx @@ -1139,6 +1386,31 @@ def v_desc_0(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[in del v_subtile_idx return self._set_desc(stage_info, 1, 0) + @consumer_work(returns=("v_desc_0",)) + @cute.jit + def v_desc_reused( + self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] + ): + """Read delayed V from the unreleased K payload in the same SMEM slot.""" + del v_subtile_idx + return self._set_desc(stage_info, 1, 0, reuse_k=True) + + @consumer_work(returns=("v_desc_0",)) + @cute.jit + def v_desc_reused_0( + self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] + ): + del v_subtile_idx + return self._set_desc(stage_info, 1, 0, reuse_k=True) + + @consumer_work(returns=("v_desc_1",)) + @cute.jit + def v_desc_reused_1( + self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] + ): + del v_subtile_idx + return self._set_desc(stage_info, 1, 1, reuse_k=True) + @consumer_work(returns=("v_desc_1",)) @cute.jit def v_desc_1(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int]): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py index f51036748809..cf547165e9e1 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py @@ -19,6 +19,8 @@ from cutlass.experimental import primitives as prims +from ...helpers.constants import LN_2, LOG2_E + import cutlass import cutlass.cute as cute from cutlass import Float32, Int8, Int32, Int64 @@ -58,6 +60,7 @@ smem_array, ) from ...helpers.math import ( + neg_max_f32, ceil_div, fadd2, ffma2, @@ -67,6 +70,7 @@ partial_output_dtype, ) from ...helpers.ops import ( + fp8_quant_scale, fp8_quant_scale_rcp, pack_float4_to_fp8_e4m3, store_transposed_smem8b_x2, @@ -112,6 +116,9 @@ class TmemCorrResource(MlaResource): inst_id: cutlass.Constexpr[int] = 0 scale_softmax_log2: Float32 = None output_scale: Float32 = None + atten_sinks: object = None + output_quant: cutlass.Constexpr = None + output_quant_args: object = None o_tensor: object = None lse_tensor: object = None acc_o_tensor: object = None @@ -133,9 +140,18 @@ class TmemCorrResource(MlaResource): _cluster_reduction_smem: object = None _cluster_reduction_barrier: object = None + def _keeps_direct_output(self): + return ( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + and self.cfg.use_cluster_reduction != 1 + ) + def get_smem_requirements(self): """Return correction/output SMEM scratch and optional cluster buffers.""" - if self.inst_id != 1: + if self.inst_id != 1 or self._keeps_direct_output(): + # Keeps-AB normalizes through warp shuffles and stores from TMEM + # directly to GMEM. Its output and sum scratch are never used. return [] if self._alloc is None: self._alloc = SmemAllocation( @@ -170,7 +186,7 @@ def get_smem_requirements(self): @cute.jit def _init_store_state_from_context(self, context) -> None: """Create correction/output SMEM views and initialize cluster state.""" - if cutlass.const_expr(self.inst_id != 1): + if cutlass.const_expr(self.inst_id != 1 or self._keeps_direct_output()): return self._smem_o = smem_array( context, @@ -239,7 +255,7 @@ def init_store_state(self, stage_info: StageInfo) -> None: def _o_base_col(self): """Return the first TMEM column owned by the O accumulator.""" - if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(self.cfg.single_kv_pipe): return Int32(2 * self.cfg.tmem_s_cols) return Int32(2 * self.cfg.tmem_s_cols + 2 * self.cfg.tmem_stats_cols) @@ -268,6 +284,10 @@ def _init_cluster_reduction_barrier(self): if warp_idx == Int32(cfg.correction_warp_idx): with cute.arch.elect_one(): prims.mbarrier_init(self._cluster_reduction_barrier, 1) + # Init and arrive can lower to different elected lanes. Order + # the warp's initialization before registering transactions. + prims.bar_warp_sync(cute.arch.FULL_MASK) + with cute.arch.elect_one(): prims.mbarrier_arrive_expect_tx( self._cluster_reduction_barrier, expected_bytes ) @@ -500,10 +520,13 @@ def _store_lse_to_gmem( self.cfg, col_group_idx, scale_idx ) head_idx = head_base_idx + local_row_idx + if cutlass.const_expr( + self.cfg.num_tokens_per_page == 1 and self.cfg.use_cluster_reduction != 1 + ): + should_store_lse = should_store_lse and cta_idx_head_dim_v == Int32(0) if should_store_lse: - # Swaps keeps at least four scale groups, so TileQ8 has eight - # padded rows in the correction register footprint. Do not let - # those rows publish split LSE into the next logical head tile. + # Bound physical rows before publishing statistics so a padded + # query row cannot write into the next logical head tile. if local_row_idx < Int32(self.cfg.tile_size_q) and head_idx < Int32( self.cfg.num_heads_q ): @@ -518,16 +541,38 @@ def _store_lse_to_gmem( batch_idx=batch_idx, ) ) - lse_sum = reduced_sum[scale_idx] + # Keep the four swap-M16 statistics in registers. A dynamic + # Array index here otherwise creates a local-memory round trip. + selected_sum = Float32(0) + selected_max = Float32(0) + for group in cutlass.range_constexpr( + num_softmax_scale_groups(self.cfg) + ): + if scale_idx == Int32(group): + selected_sum = reduced_sum[group] + selected_max = final_max[group] + lse_sum = selected_sum if cutlass.const_expr(self.cfg.is_fp8_qkv()): - # FP8 P is quantized as 448 * softmax(P) for BMM2. The - # online sum tracks the same scale, so undo it for the - # externally visible log-sum-exp value. + # P and the online sum use the same FP8 quantization scale. + # Undo that scale for the externally visible log-sum-exp. lse_sum = lse_sum * fp8_quant_scale_rcp() lse_val = ( cute.math.log2(lse_sum, fastmath=True) - + self.scale_softmax_log2 * final_max[scale_idx] + + self.scale_softmax_log2 * selected_max ) + if cutlass.const_expr(self.cfg.sparse_direct): + if selected_max == neg_max_f32(): + lse_val = Float32(-Float32.inf) + if cutlass.const_expr(self.cfg.fuse_sparse_epilogue): + count = Float32(selected_max != neg_max_f32()) + lse_val = ( + lse_val * Float32(LN_2) + if count > Float32(0) + else Float32(-Float32.inf) + ) + valid_output_row = valid_output_row and cta_idx_head_dim_v == Int32( + 0 + ) if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): self._store_partial_lse_to_cluster_smem( local_row_idx, cta_idx_kv, lse_val @@ -796,6 +841,24 @@ def _reduce_cluster_partials_and_store_for_tile( Float32(-Float32.inf) if is_fully_masked else global_lse ) + sink_scale = Float32(1) + if cutlass.const_expr(cfg.fuse_sparse_cluster_epilogue): + head = storage_flat_query_row % Int32( + cfg.logical_num_heads_q + ) + sink = Float32(self.atten_sinks[head]) + sink_scale = Float32(0) + if has_finite_mass and sink != Float32(Float32.inf): + delta = sink * Float32(LOG2_E) - global_lse + z = cute.math.exp2( + -cute.math.abs(delta), fastmath=True + ) + inverse = cute.math.rcp(Float32(1) + z, approx=True) + sink_scale = ( + z * inverse if delta > Float32(0) else inverse + ) + published_lse *= Float32(LN_2) + if ( valid_output_row and dim_idx == Int32(0) @@ -842,6 +905,9 @@ def _reduce_cluster_partials_and_store_for_tile( ).to(Float32) acc_vec = acc_vec + partial_vec * scale + if cutlass.const_expr(cfg.fuse_sparse_cluster_epilogue): + acc_vec = acc_vec * sink_scale + if cutlass.const_expr(self.o_tensor is not None): if valid_output_row: output_query_row = public_query_flat_row( @@ -858,7 +924,16 @@ def _reduce_cluster_partials_and_store_for_tile( out_byte_offset = out_elem_offset * Int64( cfg.o_dtype_bytes ) - if cutlass.const_expr(cfg.use_fp8_output == 1): + if cutlass.const_expr( + self.output_quant is not None + ): + self.output_quant.store( + self.output_quant_args, + acc_vec, + output_query_row, + head_dim_offset + dim_idx, + ) + elif cutlass.const_expr(cfg.use_fp8_output == 1): packed_o = cutlass.Array( Int32, 2, space=cutlass.AddressSpace.rmem ) @@ -992,8 +1067,20 @@ def _apply_loop_correction_in_tmem( cfg.head_dim_per_stage_v // 2, space=cutlass.AddressSpace.rmem, ) - for reg_idx in cutlass.range_constexpr(cfg.head_dim_per_stage_v // 2): - scaled[reg_idx] = loaded[reg_idx] * scale_vals[0] + if cutlass.const_expr(cfg.paired_sparse_correction): + for reg_idx in cutlass.range_constexpr( + 0, cfg.head_dim_per_stage_v // 2, 2 + ): + pair = fmul2( + (loaded[reg_idx], loaded[reg_idx + 1]), + (scale_vals[0], scale_vals[0]), + ) + scaled[reg_idx], scaled[reg_idx + 1] = pair[0], pair[1] + else: + for reg_idx in cutlass.range_constexpr( + cfg.head_dim_per_stage_v // 2 + ): + scaled[reg_idx] = loaded[reg_idx] * scale_vals[0] scaled_vec = vector_from_scalars( tuple( scaled[reg_idx] @@ -1203,6 +1290,7 @@ def _rescale_loop_o( def _compute_tail_softmax_scales( self, task_cache, + stage_info, *, new_max_arr, sum_arr, @@ -1262,9 +1350,17 @@ def _compute_tail_softmax_scales( reduced_sum[scale_idx] = final_sum[scale_idx] else: inst0_sum = inst0_sum_arr[scale_idx] - inst1_sum = inst1_sum_arr[scale_idx] inst0_max = inst0_new_max_arr[scale_idx] - inst1_max = inst1_new_max_arr[scale_idx] + inst1_sum = ( + Float32(0.0) + if cutlass.const_expr(cfg.one_insts_kv_swap) + else inst1_sum_arr[scale_idx] + ) + inst1_max = ( + inst0_max + if cutlass.const_expr(cfg.one_insts_kv_swap) + else inst1_new_max_arr[scale_idx] + ) final_max[scale_idx] = cute.math.max(inst0_max, inst1_max, ftz=True) exp_scale0[scale_idx] = cute.math.exp2( self.scale_softmax_log2 * (inst0_max - final_max[scale_idx]), @@ -1274,6 +1370,8 @@ def _compute_tail_softmax_scales( self.scale_softmax_log2 * (inst1_max - final_max[scale_idx]), fastmath=True, ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + exp_scale1[scale_idx] = Float32(0.0) final_sum[scale_idx] = ( inst0_sum * exp_scale0[scale_idx] + inst1_sum * exp_scale1[scale_idx] @@ -1306,7 +1404,7 @@ def _compute_tail_softmax_scales( ) ) - if not cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(cfg.kernel_variant != "keeps_mma_ab"): # The two softmax pipes first reduce within their owning warp, then # publish one value per column group so the warpgroup can form the # final normalization denominator before O is stored. @@ -1346,6 +1444,65 @@ def _compute_tail_softmax_scales( norm_scale = Float32(0.0) if reduced_sum[scale_idx] != Float32(0.0): norm_scale = self.output_scale / reduced_sum[scale_idx] + if cutlass.const_expr(cfg.fuse_sparse_epilogue): + # Four lanes compute the distinct row factors. Broadcast to + # their peers instead of repeating sink SFU work eight times. + norm_lane = ( + lane_idx & Int32(15) + if cfg.kernel_variant == "keeps_mma_ab" + else col_group_idx + ) + norm_lanes = 16 if cfg.kernel_variant == "keeps_mma_ab" else 4 + if lane_idx < Int32(norm_lanes): + logical_head = ( + cta_idx_q_for_stage(self.cta_idx_q, stage_info) + * cfg.tile_size_q + + head_idx_for_stage(self.head_idx, cfg, stage_info) + + ( + warp_idx * Int32(16) + norm_lane + if cfg.kernel_variant == "keeps_mma_ab" + else local_q_head_idx_for_scale( + cfg, col_group_idx, scale_idx + ) + ) + ) + count = Float32(final_max[scale_idx] != neg_max_f32()) + norm_scale = Float32(0) + if logical_head < cfg.logical_num_heads_q and count > Float32(0): + sink = Float32(self.atten_sinks[logical_head]) + if sink == Float32(-Float32.inf): + norm_scale = self.output_scale / reduced_sum[scale_idx] + elif sink != Float32(Float32.inf): + delta = ( + sink * Float32(LOG2_E) + - self.scale_softmax_log2 * final_max[scale_idx] + ) + attention_rescale = cute.math.exp2( + -cute.math.max(delta, Float32(0)), approx=True + ) + sink_mass = cute.math.exp2( + cute.math.min(delta, Float32(0)), approx=True + ) + if cutlass.const_expr(cfg.is_fp8_qkv()): + sink_mass *= fp8_quant_scale() + norm_scale = ( + self.output_scale + * attention_rescale + * cute.math.rcp( + reduced_sum[scale_idx] * attention_rescale + + sink_mass, + approx=True, + ) + ) + norm_scale = Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=norm_scale, + offset=norm_lane, + mask_and_clamp=0x1F, + kind=cprims.Shfl.IDX, + ) + ) final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] return final_max, reduced_sum, final_scale0, final_scale1 @@ -1386,6 +1543,40 @@ def _store_keeps_tail_o_stage( cfg.head_dim_per_stage_v // 2 ) + if cutlass.const_expr(self.output_quant is not None): + # Each lane owns half a head's D stage. K128 can span the two + # lanes separated by 16; K32 blocks remain entirely lane-local. + width = min(self.output_quant.block_size, cfg.head_dim_per_stage_v // 2) + for stripe in cutlass.range_constexpr( + 0, cfg.head_dim_per_stage_v // 2, width + ): + values = cutlass.Array(Float32, width, space=cutlass.AddressSpace.rmem) + loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr0 + Int32(stripe), Float32), + num=width, + offset=Int32(cfg.head_dim_per_stage_v // 2), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + for i in cutlass.range_constexpr(width): + values[i] = loaded[i] * final_scale0[0] + if global_head_idx < Int32(cfg.num_heads_q) and valid_output_row: + row = public_query_flat_row( + cfg, storage_flat_query_row, batch_idx, self.cu_seqlens_q + ) + self.output_quant.store( + self.output_quant_args, + values, + row, + head_dim_offset + + v_stage_idx * cfg.head_dim_per_stage_v + + half_warp_col_offset + + stripe, + lane_stride=16, + rotary=v_stage_idx + 1 == cfg.v_head_dim_stages, + ) + return + # ``tcgen05.ld.sync.aligned`` must remain convergent across the warp. # Padded query rows still own valid physical TMEM rows, so load and # normalize every row and predicate only the externally visible store. @@ -1695,13 +1886,16 @@ def _stage_generic_tail_o_to_smem( ), num=1, ) - o1_loaded_lo = prims.tcgen05_ld( - shape, - prims.make_tmem_ptr( - tcgen05_panel_addr(base_addr1, chunk_idx), Float32 - ), - num=1, - ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + o1_loaded_lo = o0_loaded_lo + else: + o1_loaded_lo = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr1, chunk_idx), Float32 + ), + num=1, + ) o0_loaded_hi = prims.tcgen05_ld( shape, prims.make_tmem_ptr( @@ -1712,16 +1906,19 @@ def _stage_generic_tail_o_to_smem( ), num=1, ) - o1_loaded_hi = prims.tcgen05_ld( - shape, - prims.make_tmem_ptr( - tcgen05_second_panel_addr( - tcgen05_panel_addr(base_addr1, chunk_idx) + if cutlass.const_expr(cfg.one_insts_kv_swap): + o1_loaded_hi = o0_loaded_hi + else: + o1_loaded_hi = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_second_panel_addr( + tcgen05_panel_addr(base_addr1, chunk_idx) + ), + Float32, ), - Float32, - ), - num=1, - ) + num=1, + ) else: o0_loaded = prims.tcgen05_ld( shape, @@ -1730,13 +1927,16 @@ def _stage_generic_tail_o_to_smem( ), num=q_repeats, ) - o1_loaded = prims.tcgen05_ld( - shape, - prims.make_tmem_ptr( - tcgen05_panel_addr(base_addr1, chunk_idx), Float32 - ), - num=q_repeats, - ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + o1_loaded = o0_loaded + else: + o1_loaded = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr1, chunk_idx), Float32 + ), + num=q_repeats, + ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) cute.arch.fence_view_async_tmem_load() for pair_idx in cutlass.range_constexpr(o_reg_pair_count): @@ -1770,20 +1970,28 @@ def _stage_generic_tail_o_to_smem( o1_loaded[reg_base], o1_loaded[reg_base + 1], ) - final_pair = ffma2( - ( - final_scale0[scale_base], - final_scale0[scale_base + 1], - ), - o0_pair, - fmul2( + if cutlass.const_expr(cfg.one_insts_kv_swap): + # Only O0 exists: normalize it directly, without a second + # TMEM load or the two-stream weighted merge. + final_pair = fmul2( + (final_scale0[scale_base], final_scale0[scale_base + 1]), + o0_pair, + ) + else: + final_pair = ffma2( ( - final_scale1[scale_base], - final_scale1[scale_base + 1], + final_scale0[scale_base], + final_scale0[scale_base + 1], ), - o1_pair, - ), - ) + o0_pair, + fmul2( + ( + final_scale1[scale_base], + final_scale1[scale_base + 1], + ), + o1_pair, + ), + ) if cutlass.const_expr( cfg.use_fp8_output == 1 and self.acc_o_tensor is None @@ -1841,6 +2049,97 @@ def _stage_swaps_tail_o_to_smem( final_scale1=final_scale1, ) + @cute.jit + def _store_swap_quantized_stage( + self, + stage_info, + task_cache, + base_addr0, + base_addr1, + final_scale0, + final_scale1, + v_stage_idx, + head_dim_offset, + ): + """Transpose FP32 O in two half-Q passes using the existing O scratch. + + Half as many FP32 rows fit in the same buffer as the ordinary BF16 + stage. After the transpose, each lane owns a contiguous vector, so + rotation pairs stay local and quantization needs only warp reductions. + """ + cfg = self.cfg + tid = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + loaded0 = prims.tcgen05_ld( + "32x32b", prims.make_tmem_ptr(base_addr0, Float32), num=cfg.tile_size_q + ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + loaded1 = loaded0 + else: + loaded1 = prims.tcgen05_ld( + "32x32b", prims.make_tmem_ptr(base_addr1, Float32), num=cfg.tile_size_q + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + batch = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + head_base = head_idx_for_stage(self.head_idx, cfg, stage_info) + q_tile = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + rows_per_pass = cfg.tile_size_q // 2 + width = rows_per_pass # 128 channels * rows / 128 correction threads. + threads_per_row = 128 // width + for half in cutlass.range_constexpr(2): + for i in cutlass.range_constexpr(rows_per_pass): + h = half * rows_per_pass + i + scale_idx = (h // 8) * 2 + h % 2 + source_lane = (h % 8) // 2 + s0 = cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_scale0[scale_idx], + offset=source_lane, + mask_and_clamp=0x1F, + kind=cprims.Shfl.IDX, + ) + value = loaded0[h] * s0 + if cutlass.const_expr(not cfg.one_insts_kv_swap): + s1 = cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_scale1[scale_idx], + offset=source_lane, + mask_and_clamp=0x1F, + kind=cprims.Shfl.IDX, + ) + value += loaded1[h] * s1 + self._smem_o_i32[i * 128 + tid] = value.bitcast(Int32) + prims.barrier_cta_sync( + self.store_barrier_id, thread_count=WARPGROUP_THREADS + ) + h = head_base + half * rows_per_pass + tid // threads_per_row + flat_row, _, _, _, valid = flat_query_row_state( + h, + q_tile, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + self.cu_seqlens_q, + batch, + ) + values = self._smem_o_i32.load( + tid * width, vector_size=width, alignment=width * 4 + ).bitcast(Float32) + if h < cfg.num_heads_q and valid: + row = public_query_flat_row(cfg, flat_row, batch, self.cu_seqlens_q) + self.output_quant.store( + self.output_quant_args, + values, + row, + head_dim_offset + + v_stage_idx * cfg.head_dim_per_stage_v + + (tid % threads_per_row) * width, + rotary=v_stage_idx + 1 == cfg.v_head_dim_stages, + ) + prims.barrier_cta_sync( + self.store_barrier_id, thread_count=WARPGROUP_THREADS + ) + @cute.jit def _normalize_and_store_tail_o( self, @@ -1867,6 +2166,7 @@ def _normalize_and_store_tail_o( final_max, reduced_sum, final_scale0, final_scale1 = ( self._compute_tail_softmax_scales( task_cache, + stage_info, new_max_arr=new_max_arr, sum_arr=sum_arr, inst0_new_max_arr=inst0_new_max_arr, @@ -1897,6 +2197,19 @@ def _normalize_and_store_tail_o( v_stage_idx=v_stage_idx, head_dim_offset=head_dim_offset, ) + elif cutlass.const_expr( + self.output_quant is not None and cfg.fuse_sparse_epilogue + ): + self._store_swap_quantized_stage( + stage_info, + task_cache, + base_addr0, + base_addr1, + final_scale0, + final_scale1, + v_stage_idx, + head_dim_offset, + ) else: self._stage_swaps_tail_o_to_smem( task_cache, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py index eebeb725d40f..6e66442172e4 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py @@ -127,7 +127,7 @@ def init_stage_state(self, stage_info: StageInfo): # Consumer aux work creates the stage-tracking variables used by the # correction task after it waits for O. self._init_tmem_state(stage_info) - return Int32(0), Int32(0), Int32(1) + return Int32(0), Int32(0), Int32(0 if self.cfg.one_insts_kv_swap else 1) @consumer_work( work_attrs=WorkAttr.AUXILIARY, @@ -138,7 +138,7 @@ def init_stage_work_tile_state(self, stage_info: StageInfo): """Initialize O-stage index variables for a persistent work tile.""" # Reset O stage bookkeeping for each persistent work tile. del stage_info - return Int32(0), Int32(0), Int32(1) + return Int32(0), Int32(0), Int32(0 if self.cfg.one_insts_kv_swap else 1) @cute.jit def _tail_has_prior_o(self, stage_info: StageInfo): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py index 91dd9802cbff..e867ba4e7c90 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py @@ -43,6 +43,7 @@ ) from ...helpers.math import ( fadd2, + ffma2, fmul2, neg_max_f32, pack_float2_to_bf16, @@ -71,7 +72,7 @@ class TmemPResource(MlaResource): """TMEM probability tile exchanged from softmax to PV MMA. - The keeps-MMA-AB schedule uses one P pipe: softmax writes BF16 P directly + The keeps-MMA-AB schedule uses one P pipe: softmax writes native BF16/FP8 P directly into the score TMEM stage, and MmaTask consumes that TMEM stage with ``tcgen05.mma`` A-from-TMEM. """ @@ -147,6 +148,9 @@ def materialize_p( neg_scaled_max[idx] = -self.scale_softmax_log2 * safe_new_max local_sums[idx] = Float32(0.0) + sums = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(4): + sums[i] = Float32(0) packed_p_reg_count = num_packed_p_regs(cfg) regs_p = cutlass.Array( Int32, packed_p_reg_count, space=cutlass.AddressSpace.rmem @@ -167,34 +171,44 @@ def materialize_p( p2 = Float32(0.0) p3 = Float32(0.0) if has_finite_max: - scaled01 = fadd2( - fmul2( - ( - s_arr[s_base + 0], - s_arr[s_base + 1], - ), + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + scaled01 = ffma2( + (s_arr[s_base + 0], s_arr[s_base + 1]), log2_scale_pair, - ), - neg_scaled_pair, - ) - scaled23 = fadd2( - fmul2( - ( - s_arr[s_base + 2], - s_arr[s_base + 3], + neg_scaled_pair, + ) + else: + scaled01 = fadd2( + fmul2( + (s_arr[s_base + 0], s_arr[s_base + 1]), log2_scale_pair ), + neg_scaled_pair, + ) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + scaled23 = ffma2( + (s_arr[s_base + 2], s_arr[s_base + 3]), log2_scale_pair, - ), - neg_scaled_pair, - ) + neg_scaled_pair, + ) + else: + scaled23 = fadd2( + fmul2( + (s_arr[s_base + 2], s_arr[s_base + 3]), log2_scale_pair + ), + neg_scaled_pair, + ) p0 = cute.math.exp2(scaled01[0], fastmath=True) p1 = cute.math.exp2(scaled01[1], fastmath=True) p2 = cute.math.exp2(scaled23[0], fastmath=True) p3 = cute.math.exp2(scaled23[1], fastmath=True) - local_sums[0] += p0 - local_sums[0] += p1 - local_sums[0] += p2 - local_sums[0] += p3 + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + sums[0], sums[1] = fadd2((sums[0], sums[1]), (p0, p1)) + sums[2], sums[3] = fadd2((sums[2], sums[3]), (p2, p3)) + else: + local_sums[0] += p0 + local_sums[0] += p1 + local_sums[0] += p2 + local_sums[0] += p3 regs_p[packed_idx] = pack_float4_to_fp8_e4m3(p0, p1, p2, p3) elif cutlass.const_expr(cfg.is_fp8_qkv()): for packed_idx in cutlass.range_constexpr(packed_p_reg_count): @@ -228,16 +242,31 @@ def materialize_p( if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): new_max = new_max_arr[0] if new_max != neg_max_f32(): - p0 = cute.math.exp2( - s_arr[s0] * self.scale_softmax_log2 + neg_scaled_max[0], - fastmath=True, - ) - p1 = cute.math.exp2( - s_arr[s1] * self.scale_softmax_log2 + neg_scaled_max[0], - fastmath=True, + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + scaled = ffma2( + (s_arr[s0], s_arr[s1]), + (self.scale_softmax_log2, self.scale_softmax_log2), + (neg_scaled_max[0], neg_scaled_max[0]), + ) + p0 = cute.math.exp2(scaled[0], fastmath=True) + p1 = cute.math.exp2(scaled[1], fastmath=True) + else: + p0 = cute.math.exp2( + s_arr[s0] * self.scale_softmax_log2 + neg_scaled_max[0], + fastmath=True, + ) + p1 = cute.math.exp2( + s_arr[s1] * self.scale_softmax_log2 + neg_scaled_max[0], + fastmath=True, + ) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + offset = (pair_idx % 2) * 2 + sums[offset], sums[offset + 1] = fadd2( + (sums[offset], sums[offset + 1]), (p0, p1) ) - local_sums[0] += p0 - local_sums[0] += p1 + else: + local_sums[0] += p0 + local_sums[0] += p1 else: scale0 = ((pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2) * 2 scale1 = scale0 + 1 @@ -259,11 +288,18 @@ def materialize_p( local_sums[scale1] += p1 regs_p[pair_idx] = pack_float2_to_bf16(p0, p1) + if cutlass.const_expr( + cfg.kernel_variant == "keeps_mma_ab" and cfg.num_tokens_per_page == 1 + ): + local_sums[0] = (sums[0] + sums[1]) + (sums[2] + sums[3]) + for scale_idx in cutlass.range_constexpr(num_scale_groups): self.tmem_alias_ref._p_local_sum_arr[scale_idx] = local_sums[scale_idx] local_sum_arr[scale_idx] = local_sums[scale_idx] - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr( + cfg.kernel_variant == "keeps_mma_ab" and cfg.num_tokens_per_page != 1 + ): state_idx = cute.arch.thread_idx()[0] state_ptr = self.tmem_alias_ref._softmax_scratch.data_ptr(state_idx) old_max = u32_bits_to_float(state_ptr.load(is_volatile=True, alignment=4)) @@ -295,15 +331,31 @@ def materialize_p( tcgen05_second_panel_addr(p_base), regs_p ) else: - tcgen05_store_p_16x32bx2_x16(p_base, regs_p, 0) - tcgen05_store_p_16x32bx2_x16(p_base + Int32(16), regs_p, 16) - if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): - tcgen05_store_p_16x32bx2_x16( - tcgen05_second_panel_addr(p_base), regs_p, 0 - ) - tcgen05_store_p_16x32bx2_x16( - tcgen05_second_panel_addr(p_base) + Int32(16), regs_p, 16 - ) + if cutlass.const_expr(cfg.tile_size_kv == 128): + tcgen05_store_p_16x32bx2_x16(p_base, regs_p, 0) + tcgen05_store_p_16x32bx2_x16(p_base + Int32(16), regs_p, 16) + if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base), regs_p, 0 + ) + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base) + Int32(16), regs_p, 16 + ) + else: + for part in cutlass.range_constexpr(packed_p_reg_count // 16): + tcgen05_store_p_16x32bx2_x16( + p_base + Int32(part * 16), + regs_p, + part * 16, + stride=cfg.tile_size_kv // 4, + ) + if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base) + Int32(part * 16), + regs_p, + part * 16, + stride=cfg.tile_size_kv // 4, + ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) cute.arch.fence_view_async_tmem_store() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py index 88c7aa259d4f..9831ab71e5d3 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py @@ -74,6 +74,7 @@ u32_to_float_for_atomic_max, ) from ...helpers.mask import MaskType +from ...helpers.gather import invalid_sparse_token, cache_sparse_mask from ...helpers.ops import ( float_to_u32_bits, freeze_smem_descriptor, @@ -118,6 +119,7 @@ class TmemSResource(MlaResource): p_ref: Optional[MemoryResource] = None global_ref: Optional[MemoryResource] = None cache_seqs: object = None + page_offsets: object = None head_idx: object = None batch_idx: object = None cta_idx_q: object = None @@ -162,6 +164,11 @@ def _reset_softmax_state(self): def get_smem_requirements(self): """Return the softmax scratch SMEM allocation.""" + if ( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + ): + return [] if self._scratch_alloc is None: self._scratch_alloc = SmemAllocation( name=f"{self.name}_softmaxScratch", @@ -175,8 +182,7 @@ def get_tmem_requirements(self): if self._tmem_alloc is None: num_stages = ( self.pipeline_config.num_stages - if self.cfg.kernel_variant == "keeps_mma_ab" - and self.pipeline_config is not None + if self.cfg.single_kv_pipe and self.pipeline_config is not None else 1 ) self._tmem_alloc = TmemAllocation( @@ -233,15 +239,24 @@ def init_softmax_state(self, stage_info: StageInfo): # materialize_p(), and update_softmax_sum(). context = stage_info.context self._init_tmem_state(stage_info) - self._softmax_scratch = smem_array( - context, - self._scratch_alloc, - Uint32, - softmax_scratch_words(self.cfg), - ) + if cutlass.const_expr( + not ( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + ) + ): + self._softmax_scratch = smem_array( + context, + self._scratch_alloc, + Uint32, + softmax_scratch_words(self.cfg), + ) self.q_desc_current = Int64(0) self.q_desc_rope_current = Int64(0) - if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page != 1 + ): thread_idx = cute.arch.thread_idx()[0] state_ptr = self._softmax_scratch.data_ptr(thread_idx) state_ptr.store( @@ -322,7 +337,7 @@ def qk_mma( task_cache = decode_gen_task_cache(stage_info) stage_col_offset = Int32(0) - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(cfg.single_kv_pipe): stage_col_offset = stage_info.stage_idx * Int32(cfg.tmem_s_cols) tmem_ptr = prims.make_tmem_ptr( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] @@ -373,7 +388,11 @@ def qk_mma( scale_d = Boolean(True) if cutlass.const_expr(k_block + 1 < k_block_count): if cutlass.const_expr(k_block == 3): - kv_desc = kv_desc + Int32(TCGEN05_DESC_WRAPPED_K_BLOCK_UNITS) + kv_desc = kv_desc + Int32( + cfg.tile_size_kv * 8 - 6 + if cfg.tile_size_kv != 128 + else TCGEN05_DESC_WRAPPED_K_BLOCK_UNITS + ) q_desc = q_desc + Int32(q_p_desc_k_block_wrap_units(cfg)) else: kv_desc = kv_desc + Int32(TCGEN05_DESC_NEXT_K_BLOCK_UNITS) @@ -420,7 +439,9 @@ def update_softmax( ) for idx in cutlass.range_constexpr(num_scale_groups): - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr( + cfg.kernel_variant == "keeps_mma_ab" and cfg.num_tokens_per_page != 1 + ): state_ptr = self._softmax_scratch.data_ptr(warp_grp_thread_idx) old_max_vals[idx] = u32_bits_to_float( state_ptr.load(is_volatile=True, alignment=4) @@ -436,11 +457,47 @@ def update_softmax( for idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): s_vals[idx] = neg_max_f32() + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + cfg, + task_cache, + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + local_tile_idx = softmax_kv_tile_idx(cfg, stage_info, self.inst_id) + tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) + tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) + next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) + raw_mask_routes = ( + self.page_offsets.select_source_tile( + tile_offset_k, batch_idx, for_mask=True + ) + if cutlass.const_expr(cfg.sparse_direct) + else self.page_offsets + ) + mask_routes = ( + cache_sparse_mask( + raw_mask_routes, + tile_offset_k, + seq_len_kv, + batch_idx, + swaps=cfg.kernel_variant != "keeps_mma_ab", + tile_size_kv=cfg.tile_size_kv, + ) + if cutlass.const_expr( + cfg.num_tokens_per_page == 1 and cfg.logical_seq_len_q == 1 + ) + else raw_mask_routes + ) + # Normalize both score layouts into one local S register array before # masking: keeps-MMA-AB reads a single TMEM panel, swaps-MMA-AB reads # two 16x256b panels. stage_col_offset = Int32(0) - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(cfg.single_kv_pipe): stage_col_offset = stage_info.stage_idx * Int32(cfg.tmem_s_cols) base_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] @@ -479,20 +536,6 @@ def update_softmax( s_vals[q_repeats * 4 + ld_base + 2] = loaded1[ld_base + 2] s_vals[q_repeats * 4 + ld_base + 3] = loaded1[ld_base + 3] - batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) - cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) - cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) - seq_len_kv = runtime_seq_len_kv_from_task_cache( - cfg, - task_cache, - cta_idx_q, - self.cu_seqlens_q, - batch_idx, - ) - local_tile_idx = softmax_kv_tile_idx(cfg, stage_info, self.inst_id) - tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) - tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) - next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) should_apply_dense_mask = ( (seq_len_kv % Int32(cfg.tile_size_kv)) != Int32(0) ) or (next_tile_offset_k > seq_len_kv) @@ -511,6 +554,14 @@ def update_softmax( should_apply_dense_mask = should_apply_dense_mask or ( next_tile_offset_k > min_seq_len_kv ) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + if cutlass.const_expr(cfg.logical_seq_len_q == 1): + # Cached route bits include the live prefix boundary. + should_apply_dense_mask = (mask_routes.low | mask_routes.high) != Int32( + 0 + ) + else: + should_apply_dense_mask = True if should_apply_dense_mask: # The CTA domain follows the latest logical query row in the flat # tile. Earlier rows can have a narrower bottom-right causal limit. @@ -531,7 +582,13 @@ def update_softmax( ) for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): token_idx = tile_offset_k + local_col_base + Int32(reg_idx) - if token_idx >= row_seq_len_kv: + if invalid_sparse_token( + token_idx, + row_seq_len_kv, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[reg_idx] = neg_max_f32() else: local_idx_k0 = warp_idx * Int32(WARP_LANES) + ( @@ -571,42 +628,106 @@ def update_softmax( s_base = repeat_idx * 4 s_second_panel_base = q_repeats * 4 + s_base token_idx = tile_offset_k + local_idx_k0 - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 0] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 1] = neg_max_f32() token_idx = ( tile_offset_k + local_idx_k0 + Int32(SCORE_TOKENS_PER_QK_GROUP) ) - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 2] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 3] = neg_max_f32() token_idx = ( tile_offset_k + local_idx_k0 + Int32(2 * SCORE_TOKENS_PER_QK_GROUP) ) - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 0] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 1] = neg_max_f32() token_idx = ( tile_offset_k + local_idx_k0 + Int32(3 * SCORE_TOKENS_PER_QK_GROUP) ) - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 2] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 3] = neg_max_f32() if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): # Keeps-MMA-AB keeps the softmax state in scratch words indexed by # CTA thread, so only a warp-level max is needed here. local_max = old_max_vals[0] - for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): - local_max = cute.math.max(local_max, s_vals[reg_idx], ftz=True) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + max0 = local_max + max1 = neg_max_f32() + max2 = neg_max_f32() + max3 = neg_max_f32() + for i in cutlass.range_constexpr(0, num_s_regs_per_thread(cfg), 4): + max0 = cute.math.max(max0, s_vals[i], ftz=True) + max1 = cute.math.max(max1, s_vals[i + 1], ftz=True) + max2 = cute.math.max(max2, s_vals[i + 2], ftz=True) + max3 = cute.math.max(max3, s_vals[i + 3], ftz=True) + local_max = cute.math.max( + cute.math.max(max0, max1, ftz=True), + cute.math.max(max2, max3, ftz=True), + ftz=True, + ) + else: + for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + local_max = cute.math.max(local_max, s_vals[reg_idx], ftz=True) local_max = cute.math.max( local_max, Float32( @@ -718,6 +839,21 @@ def update_softmax( ) self.new_max_state[scale_idx] = new_max_vals[scale_idx] for scale_idx in cutlass.range_constexpr(num_scale_groups): + if cutlass.const_expr(cfg.defer_sparse_max_update): + old_anchor = old_max_vals[scale_idx] + candidate = new_max_vals[scale_idx] + # Any common normalization anchor gives the same attention. + # Match FlashMLA's six-log2 bound (weights up to 64). + # Native FP8 retains exact maxima and scale 448. + # Initialize normally and update on larger score increases. + if old_anchor != neg_max_f32() and self.scale_softmax_log2 > Float32( + 0.0 + ): + increase_log2 = (candidate - old_anchor) * self.scale_softmax_log2 + limit = Float32(6.0) + if increase_log2 <= limit: + new_max_vals[scale_idx] = old_anchor + self.new_max_state[scale_idx] = new_max_vals[scale_idx] old_max_arr[scale_idx] = old_max_vals[scale_idx] sum_arr[scale_idx] = sum_vals[scale_idx] new_max_arr[scale_idx] = new_max_vals[scale_idx] @@ -744,6 +880,23 @@ def update_softmax_sum( # This consumer aux step runs after SmemP/TmemP has produced P and # published local_sum_arr. It updates the running denominator for the # next score tile without consuming a new score payload. + if cutlass.const_expr( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + ): + old_max = old_max_arr[0] + new_max = new_max_arr[0] + old_sum = sum_arr[0] + local_sum = self._p_local_sum_arr[0] + exp_scale = cute.math.exp2( + self.scale_softmax_log2 * (old_max - new_max), + fastmath=True, + ) + updated_sum = exp_scale * old_sum + local_sum + self.new_max_state[0] = new_max + self.sum_state[0] = updated_sum + sum_arr[0] = updated_sum + return old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): state_idx = cute.arch.thread_idx()[0] state_ptr = self._softmax_scratch.data_ptr(state_idx) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py index d9c83631fbec..f6bec6761305 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py @@ -23,7 +23,7 @@ import cutlass.cute as cute from cutlass import Float32, Int32 from cutlass.experimental.task_scheduling.enums import WorkAttr -from cutlass.experimental.task_scheduling.memory import TmemAllocation +from cutlass.experimental.task_scheduling.memory import TmemAllocation, SmemAllocation from cutlass.experimental.task_scheduling.resources import ( StageInfo, TaskLocalVariable, @@ -39,6 +39,7 @@ _TASK_CACHE_TMEM_BASE_OFFSET, decode_gen_task_cache, num_softmax_scale_groups, + smem_array, ) from ...helpers.math import ( neg_max_f32, @@ -51,6 +52,7 @@ MlaResource, ) + # ===================================================================== # TmemSoftmaxLocalResource — Local softmax stats in TMEM # ===================================================================== @@ -561,3 +563,75 @@ def track_global(self, stage_info: StageInfo): # This resource has no payload. Its producer work gives TS a named # dependency edge between the two softmax instances and correction. del stage_info + + +@dataclass(kw_only=True) +class SmemSoftmaxLocalResource(TmemSoftmaxLocalResource): + """Keep correction statistics alive independently of recycled QK/P slots. + + Sparse loads can change task pacing enough for a future QK to overwrite + the score TMEM before correction reads aliased statistics. The stats pipe + owns these small shared-memory slots until its consumer releases them. + """ + + def get_tmem_requirements(self): + return [] + + def get_smem_requirements(self): + if self._alloc is None: + self._alloc = SmemAllocation( + name=self.name + "_stats", + size_bytes=self.pipeline_config.num_stages + * 128 + * 2 + * num_softmax_scale_groups(self.cfg) + * 4, + alignment=16, + ) + return [self._alloc] + + @cute.jit + def _stats_pointer(self, stage_info): + storage = smem_array( + stage_info.context, + self._alloc, + Float32, + self.pipeline_config.num_stages + * 128 + * 2 + * num_softmax_scale_groups(self.cfg), + ) + local_thread = cute.arch.thread_idx()[0] & Int32(127) + return storage.data_ptr( + (stage_info.stage_idx * Int32(128) + local_thread) + * Int32(2 * num_softmax_scale_groups(self.cfg)) + ) + + @producer_work + @cute.jit + def store_loop_stats( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst_idx: cutlass.Constexpr[int], + ): + n = num_softmax_scale_groups(self.cfg) + values = tuple( + old_max_arr[i] if inst_idx == 0 else sum_arr[i] for i in range(n) + ) + values += tuple(new_max_arr[i] for i in range(n)) + self._stats_pointer(stage_info).store( + vector_from_scalars(values, Float32), + alignment=16 if n > 1 else 8, + ) + cute.arch.fence_view_async_shared() + + @cute.jit + def _load_stats_payload(self, stage_info: StageInfo): + return self._stats_pointer(stage_info).load( + count=2 * num_softmax_scale_groups(self.cfg), + alignment=16 if num_softmax_scale_groups(self.cfg) > 1 else 8, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py index bf7715ce4ddf..1d75ee237638 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py @@ -80,7 +80,7 @@ def init_variables(self, context=None): self._local_warp_idx = self._warp_grp_thread_idx >> Int32(WARP_LANE_SHIFT) self._lane_idx = self._warp_grp_thread_idx & Int32(WARP_LANE_MASK) if context is not None and context.tmem_ptr_i32 is not None: - loaded = Int32(context.tmem_ptr_i32.load()) + loaded = self._load_allocated_tmem_base(context.tmem_ptr_i32) self._tmem_base_offset = cute.arch.make_warp_uniform( prims.shfl_sync( thread_mask=0xFFFFFFFF, @@ -91,6 +91,22 @@ def init_variables(self, context=None): ) ) + @cute.jit + def _load_allocated_tmem_base(self, pointer): + # TaskManager initializes every task before dispatch. Sparse loaders + # may already be running while the MMA warp allocates TMEM, so only + # the warps participating in the readiness barrier may read its result. + if cutlass.const_expr(self.cfg is not None and self.cfg.sparse_direct): + warp = cute.arch.warp_idx() + base = Int32(0) + if warp < Int32( + self.cfg.correction_warp_idx + self.cfg.correction_num_warps + ) or warp == Int32(self.cfg.mma_warp_idx): + base = Int32(pointer.load(is_volatile=True)) + return base + else: + return Int32(pointer.load()) + @cute.jit def make_task_cache(self): """Return the compact tuple of cached per-task thread values.""" @@ -151,6 +167,8 @@ def get_domain(self, tile_coord): self.cu_seqlens_q, batch_idx, ) + if cutlass.const_expr(self.cfg.fixed_sparse_kv_tiles > 0): + return Int32(self.domain_bias) total_kv_tiles = runtime_local_kv_tiles(self.cfg, seq_len_kv) remaining_kv_tiles = cute.math.max( total_kv_tiles - Int32(self.cfg.num_insts_kv), Int32(0) @@ -309,6 +327,63 @@ def init_work_tile_state(): ) +def create_short_merged_softmax_task(s0, s1, l0, l1, p0, p1, g0, g1, cfg, **kw): + @schedule + def merged(s0, s1, l0, l1, p0, p1, g0, g1): + def process(s, local, p, global_stats): + old, total, maximum, local_sum, scores = s.init_softmax_state() + p.init_materialize_state() + s.wait() + old, total, maximum, local_sum, scores = s.update_softmax( + old_max_arr=old, + sum_arr=total, + new_max_arr=maximum, + local_sum_arr=local_sum, + s_arr=scores, + section=MlaStage.Loop, + ) + s.release() + p.acquire() + p.materialize_p(new_max_arr=maximum, s_arr=scores, local_sum_arr=local_sum) + p.commit() + global_stats.track_global() + old, total, maximum, local_sum, scores = s.update_softmax_sum( + old_max_arr=old, + sum_arr=total, + new_max_arr=maximum, + local_sum_arr=local_sum, + s_arr=scores, + ) + return old, maximum, total + + def publish_final(local, state): + old, maximum, total = state + local.acquire() + local.store_loop_stats( + old_max_arr=old, new_max_arr=maximum, sum_arr=total, inst_idx=1 + ) + local.commit() + + with domain_loop(0, 1, 1): + state0 = process(s0, l0, p0, g0) + state1 = process(s1, l1, p1, g1) + publish_final(l0, state0) + publish_final(l1, state1) + + return MlaDecodeTask( + src_resources=[s0, s1], + dst_resources=[l0, l1, p0, p1, g0, g1], + cfg=cfg, + warp_idx=0, + num_warps=4, + num_registers=cfg.softmax_regs, + schedule=merged(s0, s1, l0, l1, p0, p1, g0, g1), + name="ShortMergedSoftmaxTask", + domain_bias=1, + **kw, + ) + + def create_keeps_mma_ab_softmax_task( tmem_s, tmem_softmax_local, @@ -830,8 +905,15 @@ def init_work_tile_state(): tail_o_stage_idx_0, tail_o_stage_idx_1, ) = work_tile_state - local0_state = load_initial_local_stats(tmem_softmax_local0, local0_state) - local1_state = load_initial_local_stats(tmem_softmax_local1, local1_state) + # With only one tile per stream there is no correction loop; + # consume only final statistics, published after both P tiles. + if not cfg.short_merged_softmax: + local0_state = load_initial_local_stats( + tmem_softmax_local0, local0_state + ) + local1_state = load_initial_local_stats( + tmem_softmax_local1, local1_state + ) with domain_loop(0, domain, 1): ( @@ -1014,7 +1096,7 @@ def page_offsets_schedule(smem_page_offsets, work_queue=None): smem_page_offsets.init_load_state() with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): - if cfg.kernel_variant == "keeps_mma_ab": + if cfg.single_kv_pipe: page_offsets_produce( smem_page_offsets, "load_k0", section=MlaStage.Head ) @@ -1022,11 +1104,27 @@ def page_offsets_schedule(smem_page_offsets, work_queue=None): page_offsets_produce( smem_page_offsets, "load_k0", section=MlaStage.Loop ) + if not cfg.sparse_reuse_kv: + page_offsets_produce( + smem_page_offsets, "load_v0", section=MlaStage.Loop + ) + if not cfg.sparse_reuse_kv: page_offsets_produce( - smem_page_offsets, "load_v0", section=MlaStage.Loop + smem_page_offsets, "load_v0", section=MlaStage.Tail ) + elif cfg.serial_swap_reuse: page_offsets_produce( - smem_page_offsets, "load_v0", section=MlaStage.Tail + smem_page_offsets, "load_k0", section=MlaStage.Head + ) + with domain_loop(0, domain, 1): + page_offsets_produce( + smem_page_offsets, "load_k1", section=MlaStage.Loop + ) + page_offsets_produce( + smem_page_offsets, "load_k0", section=MlaStage.Loop + ) + page_offsets_produce( + smem_page_offsets, "load_k1", section=MlaStage.Tail ) else: page_offsets_produce( @@ -1132,7 +1230,7 @@ def load_schedule_body( else: cached_page_ids = None reuse_delayed_v_page_ids = ( - smem_page_offsets is not None and cfg.kernel_variant != "keeps_mma_ab" + smem_page_offsets is not None and not cfg.single_kv_pipe ) def load_kv( @@ -1171,75 +1269,106 @@ def load_kv( smem_q.acquire() smem_q.load_q() smem_q.commit() - load_kv( - head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k0", - section=MlaStage.Head, - page_id_slot=0, - ) - if cfg.kernel_variant != "keeps_mma_ab": + if cfg.serial_swap_reuse: load_kv( head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k1", + producer_label="load_k0", section=MlaStage.Head, - page_id_slot=1, + page_id_slot=0, ) - - with domain_loop(0, domain, 1): - if cfg.kernel_variant == "keeps_mma_ab": + with domain_loop(0, domain, 1): load_kv( head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k0", + producer_label="load_k1", section=MlaStage.Loop, + page_id_slot=1, ) load_kv( - head_dim_stages=cfg.v_head_dim_stages, - producer_label="load_v0", + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", section=MlaStage.Loop, + page_id_slot=0, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Tail, + page_id_slot=1, + ) + else: + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Head, + page_id_slot=0, + ) + if not cfg.single_kv_pipe: + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Head, + page_id_slot=1, ) - else: + + with domain_loop(0, domain, 1): + if cfg.single_kv_pipe: + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Loop, + ) + if not cfg.sparse_reuse_kv: + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v0", + section=MlaStage.Loop, + ) + else: + if not cfg.sparse_reuse_kv: + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v0", + section=MlaStage.Loop, + page_id_slot=0, + reuse_cached_page_ids=True, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Loop, + page_id_slot=0, + ) + if not cfg.sparse_reuse_kv: + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v1", + section=MlaStage.Loop, + page_id_slot=1, + reuse_cached_page_ids=True, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Loop, + page_id_slot=1, + ) + + if not cfg.sparse_reuse_kv: load_kv( head_dim_stages=cfg.v_head_dim_stages, producer_label="load_v0", - section=MlaStage.Loop, + section=MlaStage.Tail, page_id_slot=0, reuse_cached_page_ids=True, ) - load_kv( - head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k0", - section=MlaStage.Loop, - page_id_slot=0, - ) + if not cfg.single_kv_pipe and not cfg.sparse_reuse_kv: load_kv( head_dim_stages=cfg.v_head_dim_stages, producer_label="load_v1", - section=MlaStage.Loop, + section=MlaStage.Tail, page_id_slot=1, reuse_cached_page_ids=True, ) - load_kv( - head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k1", - section=MlaStage.Loop, - page_id_slot=1, - ) - - load_kv( - head_dim_stages=cfg.v_head_dim_stages, - producer_label="load_v0", - section=MlaStage.Tail, - page_id_slot=0, - reuse_cached_page_ids=True, - ) - if cfg.kernel_variant != "keeps_mma_ab": - load_kv( - head_dim_stages=cfg.v_head_dim_stages, - producer_label="load_v1", - section=MlaStage.Tail, - page_id_slot=1, - reuse_cached_page_ids=True, - ) @schedule def load_schedule(smem_q, smem_kv): @@ -1514,67 +1643,150 @@ def mma_schedule( q_desc, q_desc_rope = smem_q.q_desc() tmem_s0.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) tmem_s1.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) - staged_qk_mma( - smem_kv, - tmem_s0, - head_dim_stages=cfg.qk_head_dim_stages, - consumer_label="k_desc_0", - ) - staged_qk_mma( - smem_kv, - tmem_s1, - head_dim_stages=cfg.qk_head_dim_stages, - consumer_label="k_desc_1", - ) - with domain_loop(0, domain, 1): - tmem_s0.acquire() + if cfg.serial_swap_reuse: + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=False, + ) + with domain_loop(0, domain, 1): + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_reused_0", + reuse_k=True, + producer_label="pv_mma_loop_0", + ) + staged_qk_mma( + smem_kv, + tmem_s1, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_1", + release_k=False, + ) + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_reused_1", + reuse_k=True, + producer_label="pv_mma_loop_1", + ) + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=False, + ) staged_pv_mma( smem_kv, smem_p0, tmem_o, head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_loop_0", + consumer_label="v_desc_reused_0", + reuse_k=True, + producer_label="pv_mma_tail_0", ) staged_qk_mma( smem_kv, - tmem_s0, + tmem_s1, head_dim_stages=cfg.qk_head_dim_stages, - consumer_label="k_desc_0", - include_acquire=False, + consumer_label="k_desc_1", + release_k=False, ) - tmem_s1.acquire() staged_pv_mma( smem_kv, smem_p1, tmem_o, head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_1", - producer_label="pv_mma_loop_1", + consumer_label="v_desc_reused_1", + reuse_k=True, + producer_label="pv_mma_tail_1", + ) + else: + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, ) staged_qk_mma( smem_kv, tmem_s1, head_dim_stages=cfg.qk_head_dim_stages, consumer_label="k_desc_1", - include_acquire=False, + release_k=not cfg.sparse_reuse_kv, + ) + with domain_loop(0, domain, 1): + tmem_s0.acquire() + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_0" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_loop_0", + ) + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, + include_acquire=False, + ) + tmem_s1.acquire() + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_1" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_loop_1", + ) + staged_qk_mma( + smem_kv, + tmem_s1, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_1", + release_k=not cfg.sparse_reuse_kv, + include_acquire=False, + ) + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_0" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_tail_0", + ) + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_1" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_tail_1", ) - staged_pv_mma( - smem_kv, - smem_p0, - tmem_o, - head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_tail_0", - ) - staged_pv_mma( - smem_kv, - smem_p1, - tmem_o, - head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_1", - producer_label="pv_mma_tail_1", - ) smem_q.release() captured_schedule = ( @@ -1615,11 +1827,11 @@ def mma_schedule( ) -def create_keeps_mma_ab_mma_task( +def create_single_kv_pipe_mma_task( smem_q, smem_kv, tmem_s, - tmem_p, + p_resource, tmem_o, work_queue: WorkQueue | None, cfg, @@ -1628,33 +1840,49 @@ def create_keeps_mma_ab_mma_task( task_class=MlaDecodeTask, **kw, ) -> Task: - """Create the keeps-MMA-AB single-pipe MMA task. + """Schedule one QK/PV stream with SMEM P (swap) or TMEM P (keep). - HEAD computes QK for K[0]. LOOP computes QK for K[n] before PV for - K[n-1], and TAIL drains PV for K[last]. + HEAD computes the first QK; TAIL drains the last PV. Retaining one BF16 + BK128 tile requires PV-before-next-QK; two retained tiles permit overlap. """ work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + swap = cfg.one_insts_kv_swap + pv_mma = staged_pv_mma if swap else staged_pv_mma_tmem_p + pv_loop = "pv_mma_loop_0" if swap else "pv_mma_loop_tmem_p" + pv_tail = "pv_mma_tail_0" if swap else "pv_mma_tail_tmem_p" + v_desc = "v_desc_reused_0" if swap else "v_desc_reused" + if not cfg.sparse_reuse_kv: + v_desc = "v_desc_0" + pv_before_qk = ( + cfg.sparse_reuse_kv + and not cfg.is_fp8_qkv() + and (swap or cfg.tile_size_kv == 128) + ) @schedule def mma_schedule( smem_q, smem_kv, tmem_s, - tmem_p, + p_resource, tmem_o, work_queue=None, ): - """Captured keeps-MMA-AB schedule with TMEM P.""" + """Capture identical pipeline ownership for the two P layouts.""" smem_q.init_descriptor_state() smem_kv.init_descriptor_state() tmem_o.init_mma_state() - tmem_p.init_stage_state() + if swap: + p_resource.init_descriptor_state() + else: + p_resource.init_stage_state() with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): if work_queue is not None: tmem_s.reset_softmax_work_tile_state() - tmem_p.init_stage_work_tile_state() + if not swap: + p_resource.init_stage_work_tile_state() smem_q.wait() q_desc, q_desc_rope = smem_q.q_desc() tmem_s.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) @@ -1663,38 +1891,53 @@ def mma_schedule( tmem_s, head_dim_stages=cfg.qk_head_dim_stages, consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, ) with domain_loop(0, domain, 1): + if pv_before_qk: + pv_mma( + smem_kv, + p_resource, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=v_desc, + reuse_k=True, + producer_label=pv_loop, + ) staged_qk_mma( smem_kv, tmem_s, head_dim_stages=cfg.qk_head_dim_stages, consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, ) - staged_pv_mma_tmem_p( - smem_kv, - tmem_p, - tmem_o, - head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_loop_tmem_p", - ) - staged_pv_mma_tmem_p( + if not pv_before_qk: + pv_mma( + smem_kv, + p_resource, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=v_desc, + reuse_k=cfg.sparse_reuse_kv, + producer_label=pv_loop, + ) + pv_mma( smem_kv, - tmem_p, + p_resource, tmem_o, head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_tail_tmem_p", + consumer_label=v_desc, + reuse_k=cfg.sparse_reuse_kv, + producer_label=pv_tail, ) smem_q.release() schedule_result = ( - mma_schedule(smem_q, smem_kv, tmem_s, tmem_p, tmem_o) + mma_schedule(smem_q, smem_kv, tmem_s, p_resource, tmem_o) if work_queue is None - else mma_schedule(smem_q, smem_kv, tmem_s, tmem_p, tmem_o, work_queue) + else mma_schedule(smem_q, smem_kv, tmem_s, p_resource, tmem_o, work_queue) ) - src = [smem_q, smem_kv, tmem_p] + src = [smem_q, smem_kv, p_resource] if work_queue is not None: src.append(work_queue) return task_class( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py index 01bd9e49152d..152e956a40a5 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py @@ -28,6 +28,8 @@ import cutlass.cute as cute from cutlass import BFloat16, Float16, Float32, Int32 +from .mla_decode.helpers.constants import LOG2_E + @cute.jit def finalize_log2_sum_exp(frame: Float32, exp_sum: Float32) -> Float32: @@ -98,3 +100,15 @@ def unpack_normalized_vec8( if cutlass.const_expr(use_bf16_partial): return regs_vec.bitcast(BFloat16).to(Float32) return regs_vec.bitcast(Float16).to(Float32) + + +@cute.jit +def attention_sink_scale(log2_lse, sink): + """Denominator-only sink mass; keep the KV LSE independent of the sink.""" + scale = Float32(0) + if log2_lse != Float32(-Float32.inf) and sink != Float32(Float32.inf): + delta = sink * Float32(LOG2_E) - log2_lse + z = cute.math.exp2(-cute.math.abs(delta), fastmath=True) + inverse = cute.math.rcp(Float32(1) + z, approx=True) + scale = z * inverse if delta > Float32(0) else inverse + return scale diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 3d5fb159a329..7a305b1ec6f2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -16,16 +16,19 @@ """Task-scheduled paged MLA decode with a plan/run lifecycle.""" from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +import copy +import math +import struct import functools from typing import Any, Literal, Optional, cast +from ._compile_cache import build_and_load_attention_kernel +from ._hardware import get_max_active_clusters + import torch from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api -from flashinfer.trace.templates.attention import ( - prims_ts_decode_mla_one_shot_trace_dispatch, -) from .decode import ( _WorkspaceSection, @@ -44,6 +47,8 @@ ) +from .kernels.mla_decode.sparse_policy import select_sparse_mla_profile + _COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 2" _MLA_LATENT_DIM = 512 _MLA_ROPE_DIM = 64 @@ -136,6 +141,9 @@ class _MLARuntime: num_physical_pages: int bmm1_scale: float bmm2_scale: float + extra_cache: Optional[torch.Tensor] = None + sparse_inputs: Optional[tuple] = None + output_quant_args: Optional[tuple] = None def _make_mla_workspace_layout( @@ -533,21 +541,24 @@ def _validate_query( num_heads: Optional[int] = None, max_seq_len_q: Optional[int] = None, q_dtype: Optional[torch.dtype] = None, + query_dim: int = _MLA_QUERY_DIM, ) -> None: if not isinstance(query, torch.Tensor): raise TypeError("query must be a torch.Tensor") expected_rank = 3 if packed_query else 4 if query.ndim != expected_rank: - expected_shape = "[total_q, H, 576]" if packed_query else "[B, SQ, H, 576]" + expected_shape = ( + f"[total_q, H, {query_dim}]" if packed_query else f"[B, SQ, H, {query_dim}]" + ) raise ValueError(f"query must have shape {expected_shape}") if packed_query: if int(query.shape[1]) <= 0: raise ValueError("query head extent must be positive") elif any(int(extent) <= 0 for extent in query.shape[:-1]): raise ValueError("query row and head extents must be positive") - if query.shape[-1] != _MLA_QUERY_DIM: + if query.shape[-1] != query_dim: raise ValueError( - f"query last dimension must be {_MLA_QUERY_DIM}, got {query.shape[-1]}" + f"query last dimension must be {query_dim}, got {query.shape[-1]}" ) if query.dtype not in _SUPPORTED_INPUT_DTYPES: raise NotImplementedError( @@ -596,7 +607,9 @@ def _validate_query( f"query dtype must match the plan ({q_dtype}), got {query.dtype}" ) if not query.is_contiguous(): - layout = "[total_q, H, 576]" if packed_query else "[B, SQ, H, 576]" + layout = ( + f"[total_q, H, {query_dim}]" if packed_query else f"[B, SQ, H, {query_dim}]" + ) raise ValueError(f"query must be compact in {layout} layout") _validate_16byte_alignment(query, "query") @@ -731,8 +744,6 @@ def _resolve_mla_decode_launch_spec( max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") import cutlass - import cutlass.utils as cutlass_utils - from cuda.bindings import driver as cuda_drv from .kernels.mla_decode.kernel_policy import ( resolve_mla_kernel_policy, @@ -768,16 +779,8 @@ def _resolve_mla_decode_launch_spec( output_dtype_name = _kernel_dtype_name(output_dtype_key) with torch.cuda.device(device_index): - plan_stream = cuda_drv.CUstream( - torch.cuda.current_stream(device_index).cuda_stream - ) - hardware_info = cutlass_utils.HardwareInfo(device_index) - max_active_one_cta_clusters = hardware_info.get_max_active_clusters( - 1, plan_stream - ) - max_active_two_cta_clusters = hardware_info.get_max_active_clusters( - 2, plan_stream - ) + max_active_one_cta_clusters = get_max_active_clusters(1) + max_active_two_cta_clusters = get_max_active_clusters(2) one_cta_launch_shape = resolve_auto_flat_query_launch_shape( num_heads_q=num_heads, seq_len_q=seq_len_q, @@ -843,8 +846,6 @@ def _resolve_mla_decode_launch_spec( reduction_mode = resolve_runtime_cluster_reduction_mode( decision.config, reduction_mode=None, - hardware_info=hardware_info, - stream=plan_stream, ) kernel = ThroughputLatencyMlaDecodeTs( batch_size=batch_size, @@ -1043,11 +1044,35 @@ def _make_mla_decode_compile_spec( ) -@functools.cache -def _get_compiled_mla_decode( - compile_spec: _MLADecodeCompileSpec, -): - """Compile and cache one batch-dynamic TS MLA topology.""" +_COMPILED_MLA_DECODE: dict = {} + + +def _get_compiled_mla_decode(compile_spec: _MLADecodeCompileSpec): + """Compile and cache one batch-dynamic TS MLA topology. + + A manual memo rather than ``functools.cache``: under the warm-up precompile + modes of ``_compile_cache`` the build returns a placeholder (probe / persist + only), which must not be remembered for the serving launches that follow. + """ + compiled = _COMPILED_MLA_DECODE.get(compile_spec) + if compiled is None: + compiled = _compile_mla_decode(compile_spec) + if not _compile_cache_is_placeholder(compiled): + _COMPILED_MLA_DECODE[compile_spec] = compiled + return compiled + + +_get_compiled_mla_decode.cache_clear = _COMPILED_MLA_DECODE.clear # type: ignore[attr-defined] + + +def _compile_cache_is_placeholder(compiled) -> bool: + from ._compile_cache import is_placeholder + + return is_placeholder(compiled) + + +def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): + """Build (or load from the on-disk cache) one batch-dynamic TS MLA topology.""" import cutlass import cutlass.cute as cute @@ -1056,6 +1081,10 @@ def _get_compiled_mla_decode( num_heads = compile_spec.num_heads kv_lora_rank = compile_spec.kv_lora_rank qk_rope_head_dim = compile_spec.qk_rope_head_dim + query_dim = kv_lora_rank + qk_rope_head_dim + # The D512 path has no separate RoPE stage. Use an unused full-width + # alias in the callable ABI rather than a zero-extent tensor map. + rope_view_dim = qk_rope_head_dim or kv_lora_rank page_size = compile_spec.page_size max_seq_len_q = compile_spec.max_seq_len_q packed_query = compile_spec.packed_query @@ -1072,21 +1101,21 @@ def _get_compiled_mla_decode( # These fake tensors pin the public ABI while allowing runtime page counts, # table widths/row strides, and batch metadata pointers to vary. - q_stride_h = _MLA_QUERY_DIM - q_stride_q = num_heads * _MLA_QUERY_DIM + q_stride_h = query_dim + q_stride_q = num_heads * query_dim q_latent_shape: tuple[int, ...] q_rope_shape: tuple[int, ...] q_stride: tuple[int, ...] if packed_query: q_latent_shape = (num_heads, kv_lora_rank, runtime_total_q) - q_rope_shape = (num_heads, qk_rope_head_dim, runtime_total_q) + q_rope_shape = (num_heads, rope_view_dim, runtime_total_q) q_stride = (q_stride_h, 1, q_stride_q) else: q_stride_batch = max_seq_len_q * q_stride_q q_latent_shape = (num_heads, kv_lora_rank, max_seq_len_q, batch_size) q_rope_shape = ( num_heads, - qk_rope_head_dim, + rope_view_dim, max_seq_len_q, batch_size, ) @@ -1097,8 +1126,8 @@ def _get_compiled_mla_decode( q_rope_fake = cute.runtime.make_fake_tensor( qkv_dtype, q_rope_shape, stride=q_stride, assumed_align=16 ) - cache_token_stride = _MLA_QUERY_DIM - cache_page_stride = page_size * _MLA_QUERY_DIM + cache_token_stride = query_dim + cache_page_stride = page_size * query_dim c_latent_fake = cute.runtime.make_fake_tensor( qkv_dtype, (page_size, kv_lora_rank, physical_pages), @@ -1107,7 +1136,11 @@ def _get_compiled_mla_decode( ) c_rope_fake = cute.runtime.make_fake_tensor( qkv_dtype, - (page_size, qk_rope_head_dim, physical_pages), + ( + page_size, + rope_view_dim, + cute.sym_int() if page_size == 1 else physical_pages, + ), stride=(cache_token_stride, 1, cache_page_stride), assumed_align=16, ) @@ -1165,12 +1198,40 @@ def _get_compiled_mla_decode( stride_order=(0,), assumed_align=4, ) + sparse_inputs_fake = None + if getattr(kernel, "direct_sparse", False): + + def sparse_indices(capacity): + return cute.runtime.make_fake_tensor( + cutlass.Int32, + (batch_size, max(1, capacity)), + stride=(cute.sym_int64(), 1), + assumed_align=4, + ) + + scalar = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (1,), stride_order=(0,), assumed_align=4 + ) + sink = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (num_heads,), stride_order=(0,), assumed_align=4 + ) + sparse_inputs_fake = ( + sparse_indices(kernel.direct_sparse_capacities[0]), + sparse_indices(kernel.direct_sparse_capacities[1]), + cache_seqs_fake, + cache_seqs_fake, + scalar, + scalar, + sink, + ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + output_quant_args = _fake_output_quant_args(getattr(kernel, "output_quant", None)) + # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. - with torch.cuda.device(device_index): - compiled = cute.compile[cute.FrontendNext]( + def compile_kernel(): + return cute.compile[cute.FrontendNext]( kernel, q_latent_fake, q_rope_fake, @@ -1183,12 +1244,23 @@ def _get_compiled_mla_decode( cutlass.Int32(compile_spec.split_kv), cache_seqs_fake, qo_indptr_fake, - None, + sparse_inputs_fake, cutlass.Float32(1.0), cutlass.Float32(1.0), stream_fake, + output_quant_args=output_quant_args, options=_COMPILE_OPTIONS, ) + + with torch.cuda.device(device_index): + compiled = build_and_load_attention_kernel( + "mla_decode", + "main", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(), + ) return compiled @@ -1387,6 +1459,11 @@ def _launch_mla_decode( lse_kernel = workspace.lse.permute(2, 1, 0) c_latent = runtime.normalized_cache[..., :kv_lora_rank].permute(1, 2, 0) c_rope = runtime.normalized_cache[..., kv_lora_rank:].permute(1, 2, 0) + if runtime.query.shape[-1] == kv_lora_rank: + q_rope = q_latent + c_rope = c_latent + if runtime.extra_cache is not None: + c_rope = runtime.extra_cache.permute(1, 2, 0) page_offsets = block_tables.transpose(0, 1) compiled( q_latent, @@ -1400,9 +1477,15 @@ def _launch_mla_decode( split_kv, seq_lens, qo_indptr, - None, + runtime.sparse_inputs, runtime.bmm1_scale, runtime.bmm2_scale, + # Passed positionally and also when None. The kernel was compiled with + # ``output_quant_args`` in its signature (``None`` for unquantized + # output), and a kernel reloaded from the persistent compile cache is + # a positional TVM-FFI entry with no Python defaults: it rejects both + # a missing parameter and a keyword argument. + runtime.output_quant_args, ) return runtime.out @@ -1681,7 +1764,7 @@ def run( ) -@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) +@flashinfer_experimental_api def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, @@ -1885,3 +1968,1096 @@ def batch_mla_decode_with_paged_kv_cache( "batch_mla_decode_with_paged_kv_cache", "get_prims_ts_batch_mla_decode_workspace_size", ] + + +def _make_sparse_output_quant(heads, block_size, scale_format, groups): + if block_size is None: + if groups is not None or scale_format != "ue8m0": + raise ValueError( + "output quantization options require output_quant_block_size" + ) + return None + block_size = _validate_positive_int(block_size, "output_quant_block_size") + if block_size not in (32, 128): + raise ValueError("output_quant_block_size must be 32, 128, or None") + if scale_format not in ("fp32", "ue8m0"): + raise ValueError("out_scale_format must be 'fp32' or 'ue8m0'") + groups = _validate_positive_int(groups, "num_output_groups") + if heads % groups: + raise ValueError("num_output_groups must divide num_heads") + from .kernels.mla_decode.helpers.output_quant import MlaOutputQuant + + return MlaOutputQuant(block_size, scale_format == "ue8m0", heads, groups) + + +def _fake_output_quant_args(quant): + if quant is None: + return None + import cutlass + import cutlass.cute as cute + + rows = cute.sym_int() + padded_rows = cute.sym_int(divisibility=4) + columns = quant.num_heads // quant.num_groups * 512 // quant.block_size + if quant.ue8m0: + columns //= 4 + return ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (cute.sym_int64(),), + stride_order=(0,), + assumed_align=16, + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (rows,), stride_order=(0,), assumed_align=4 + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 64), stride_order=(1, 0), assumed_align=16 + ), + cute.runtime.make_fake_tensor( + cutlass.Int32 if quant.ue8m0 else cutlass.Float32, + (rows, quant.num_groups, columns), + stride=(1, columns * padded_rows, padded_rows), + assumed_align=16, + ), + ) + + +def _prepare_sparse_output_quant( + quant, query, positions, cos_sin, out, scales, validate +): + """Bind the token-major output and column-major block scales.""" + prefix = tuple(query.shape[:-2]) + rows, device = math.prod(prefix), query.device + channels = quant.num_heads // quant.num_groups * 512 + columns = channels // quant.block_size // (4 if quant.ue8m0 else 1) + padded_rows = (rows + 3) // 4 * 4 + if ( + positions is None + or positions.shape != prefix + or positions.dtype != torch.int32 + or positions.device != device + or not positions.is_contiguous() + ): + raise ValueError( + "token_positions must be contiguous CUDA INT32 matching the query prefix" + ) + if ( + cos_sin is None + or cos_sin.ndim != 2 + or cos_sin.shape[0] == 0 + or cos_sin.shape[1] != 64 + or cos_sin.dtype != torch.float32 + or cos_sin.device != device + or not cos_sin.is_contiguous() + or cos_sin.data_ptr() % 16 + ): + raise ValueError( + "cos_sin_cache must be 16-byte aligned contiguous CUDA FP32[P,64]" + ) + if ( + validate + and rows + and ((positions < 0) | (positions >= cos_sin.shape[0])).any().item() + ): + raise ValueError("token_positions must index cos_sin_cache") + shape = (rows, quant.num_groups, channels) + # Physical storage is group-major [G, T, C] (as in vLLM and TensorRT-LLM); + # the public tensor is its token-major [T, G, C] view. + if out is None: + out = torch.empty( + (quant.num_groups, rows, channels), dtype=torch.float8_e4m3fn, device=device + ).transpose(0, 1) + if ( + out.shape != shape + or out.dtype != torch.float8_e4m3fn + or out.device != device + or out.data_ptr() % 16 + ): + raise ValueError( + f"out must be the [T, G, C] view of a contiguous group-major CUDA " + f"E4M3 [G, T, C] tensor ({shape}), aligned to 16 bytes" + ) + physical_out = out.transpose(0, 1) + if rows and quant.num_groups > 1: + # A live prefix can retain a larger graph bucket's group stride. + capacity, remainder = divmod(physical_out.stride(0), channels) + if remainder or capacity < rows or capacity * quant.num_heads >= 2**31: + raise ValueError("out must have valid group-major token capacity") + try: + physical_out = physical_out.as_strided( + (quant.num_groups, capacity, channels), physical_out.stride() + ) + except RuntimeError as exc: + raise ValueError( + "out must view complete group-major backing storage" + ) from exc + if not physical_out.is_contiguous(): + raise ValueError("out must have contiguous physical group-major storage") + scale_dtype = torch.int32 if quant.ue8m0 else torch.float32 + scale_shape = (rows, quant.num_groups, columns) + scale_stride = (1, columns * padded_rows, padded_rows) + if scales is None: + scales = torch.empty( + (quant.num_groups, columns, padded_rows), dtype=scale_dtype, device=device + ).permute(2, 0, 1)[:rows] + if rows == 0: + scales = scales.as_strided(scale_shape, scale_stride) + if ( + scales.shape != scale_shape + or scales.stride() != scale_stride + or scales.dtype != scale_dtype + or scales.device != device + or scales.data_ptr() % 16 + ): + raise ValueError( + f"out_scale must be 16-byte aligned CUDA {scale_dtype}{scale_shape} " + f"with strides {scale_stride}" + ) + return ( + out, + scales, + (physical_out.view(-1), positions.view(-1), cos_sin, scales), + ) + + +@functools.lru_cache(maxsize=1024) +def _resolve_sparse_mla_plan( + device, + batch_size, + num_heads, + *, + max_topk, + max_extra_topk=0, + max_seq_len_q=1, + packed_query=False, + q_data_type=torch.bfloat16, + kv_layout="NHD", + has_sinks=False, + return_lse=False, + assume_valid_prefix=False, + output_quant_block_size=None, + out_scale_format="ue8m0", + num_output_groups=None, +): + """Resolve sparse geometry and scratch sizing without allocation/compilation.""" + import cutlass + from .kernels.mla_decode.throughput_latency_1cta.config import ( + MlaProfile, + compute_workspace_size, + ) + from .kernels.mla_decode.throughput_latency_1cta.kernel import ( + ThroughputLatencyMlaDecodeTs, + ) + from .kernels.mla_decode.throughput_2cta.config import ( + compute_workspace_size as workspace_2cta, + ) + from .kernels.mla_decode.throughput_2cta.kernel import ( + MlaDecodeTs, + ) + + if not isinstance(assume_valid_prefix, bool): + raise TypeError("assume_valid_prefix must be bool") + device, _ = _resolve_cuda_device(device) + if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3)): + raise NotImplementedError("sparse TS MLA requires SM100 or SM103") + for name, value in ( + ("batch_size", batch_size), + ("num_heads", num_heads), + ("max_seq_len_q", max_seq_len_q), + ): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if num_heads > 128: + raise ValueError("num_heads must be at most 128") + output_quant = _make_sparse_output_quant( + num_heads, output_quant_block_size, out_scale_format, num_output_groups + ) + for value in (max_topk, max_extra_topk): + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise ValueError("source capacities must be nonnegative integers") + if max_topk == 0: + raise ValueError("the primary selected-slot capacity must be positive") + if q_data_type not in (torch.bfloat16, torch.float8_e4m3fn): + raise ValueError("native BF16 or E4M3 is required") + if kv_layout not in ("NHD", "HND"): + raise ValueError("kv_layout must be NHD or HND") + max_rows = batch_size * max_seq_len_q + capacity = max( + 256, + ((max_topk + 127) // 128 + (max_extra_topk + 127) // 128) * 128, + ) + if max_rows * num_heads >= 2**31 or capacity >= 2**31 - 32768: + raise ValueError("query/route extent exceeds the int32 kernel domain") + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + profile = select_sparse_mla_profile( + rows=max_rows, + heads=num_heads, + query_length=max_seq_len_q, + capacity=capacity, + dtype="bf16" if q_data_type == torch.bfloat16 else "fp8", + sm_count=sm_count, + ) + family, tile = profile.family, profile.tile_size_q + splits, head_dim_ctas = profile.split_kv, profile.head_dim_ctas + dtype_name = "bf16" if q_data_type == torch.bfloat16 else "e4m3" + if family == "2cta": + kernel = MlaDecodeTs( + page_size=1, + max_active_clusters=sm_count // 2, + is_persistent=profile.scheduler != "nonpersistent", + is_var_seq=False, + is_var_split_kv=False, + static_split_kv=splits, + qkv_dtype=dtype_name, + out_dtype="bf16", + rope_dim=0, + num_heads=num_heads, + seq_len_q=1, + batch_size=max_rows, + mask_type="dense", + ) + kernel.sparse_gather_warps = profile.gather_issue_warps + kernel.sparse_kv_stages = profile.kv_pipeline_stages + kernel.sparse_uniform_pages = profile.uniform_offset_cache + kernel.sparse_balanced_registers = profile.balanced_registers + core_bytes = workspace_2cta( + tile_size_q=128, + num_q_tiles=1, + latent_dim=512, + batch_size=max_rows, + split_kv=splits, + partial_o_dtype=cutlass.BFloat16, + lse_dtype=cutlass.Float32, + ) + else: + kernel_profile = MlaProfile( + name="sparse", + kernel_variant="swaps_mma_ab" if family == "swap" else "keeps_mma_ab", + tile_size_q=tile, + num_ctas_per_head_dim=head_dim_ctas, + num_ctas_per_seq_kv=splits, + use_multi_ctas_kv=int(splits > 1), + use_cluster_reduction=int(profile.reduction == "cluster"), + use_persistent_scheduler=int(profile.scheduler != "nonpersistent"), + use_clc_dynamic_persistent_scheduler=int(profile.scheduler == "clc"), + ) + kernel = ThroughputLatencyMlaDecodeTs( + batch_size=max_rows, + num_heads=tile, + seq_len_q=(num_heads + tile - 1) // tile, + seq_len_k=capacity, + rope_dim=0, + page_size=1, + max_active_clusters=sm_count, + qkv_dtype=dtype_name, + out_dtype="bf16", + profile=kernel_profile, + reduction_mode=profile.reduction, + logical_num_heads=num_heads, + logical_seq_len_q=1, + tile_size_q=tile, + sparse_profile=profile, + mask_type="dense", + ) + core_bytes = compute_workspace_size( + cfg=kernel._make_config(), + partial_o_dtype=cutlass.BFloat16, + lse_dtype=cutlass.Float32, + ) + spec = _make_mla_decode_compile_spec( + _MLADecodeLaunchSpec(kernel, (), core_bytes, splits), + device_index=device.index, + num_heads=num_heads, + kv_lora_rank=512, + qk_rope_head_dim=0, + page_size=1, + q_dtype_key=str(q_data_type).removeprefix("torch."), + output_dtype_key="float8_e4m3fn" if output_quant is not None else "bfloat16", + max_seq_len_q=1, + packed_query=False, + ) + sections = {} + byte_end = 0 + for name, shape, dtype in ( + ("core", (core_bytes,), torch.int8), + ("public_lse", (max_rows, num_heads), torch.float32), + ): + sections[name], byte_end = _append_workspace_section(byte_end, shape, dtype) + return device, profile, spec, sections, byte_end, capacity, output_quant + + +# Prepared launches kept per plan (one per layer sharing the plan; DeepSeek-V4 has 61). +_PREPARED_LAUNCH_LIMIT = 256 + +class BatchSparseMLADecodePagedTSWrapper: + """Plan native BF16/E4M3 attention, then bind one or two source index lists. + + Q is [B,Sq,H,512] or packed [total_q,H,512]. Each source supplies int32 + [rows, capacity] storage-row indices and int32 [rows] live lengths directly + to run(). Indices include physical page strides; -1 is masked. Both source + lists participate in one attention distribution. The caller folds the + shared KV descale into the effective softmax and output scales. + + ``output_quant_block_size=32`` or ``128`` fuses conjugate output RoPE on + D[448:512] and E4M3 quantization after the final attention reduction. Q/KV + remain prepared native tensors. Supply the local ``num_output_groups`` G + and ``out_scale_format="ue8m0"`` (packed exponent bytes) or ``"fp32"``. + This option does not change attention dispatch or softmax correction. + + Planning with ``assume_valid_prefix=True`` promises that every index before + each live source length is valid (no -1 holes). Kernels can then derive + masks from lengths and omit per-index hole checks. This does not imply + that the physical KV rows are consecutive. + ``run(validate=True)`` checks the promise; graph replay with validation + disabled must preserve it. Tail entries beyond the lengths are ignored. + + A wrapper/workspace permits one in-flight run. Graphs require a completed + warmup, stable addresses, preallocated outputs, and validate=False. + Input/output/workspace storage must not overlap; this is an unchecked + caller precondition. Separate wrappers/workspaces are needed per stream. + """ + + @flashinfer_experimental_api + def __init__(self, workspace_buffer=None): + self._workspace_buffer = workspace_buffer + self._state = None + self._scalar_cache = {} + # Prepared launches for the unfused validate=False path: everything that + # is static for one (plan, pools, scales, sinks) tuple, keyed by the pools' + # storage and the sinks address. One plan serves every layer of a model, + # and each layer brings its own sinks, so one entry per layer. + self._prepared: dict[tuple, dict] = {} + self.workspace_size_bytes = 0 + + def plan( + self, + device, + batch_size, + num_heads, + *, + max_topk, + max_extra_topk=0, + max_seq_len_q=1, + packed_query=False, + q_data_type=torch.bfloat16, + kv_layout="NHD", + has_sinks=False, + return_lse=False, + assume_valid_prefix=False, + output_quant_block_size=None, + out_scale_format="ue8m0", + num_output_groups=None, + ): + """Plan D512 attention over one primary and an optional extra pool. + + ``max_topk`` describes arbitrary selected primary entries, not an SWA + window. Prepared indices already include the physical page stride, so + the core uses page size one regardless of the pools' external pages. + Planning compiles no metadata-preparation kernel. + + Output quantization is disabled by default. When enabled, G must divide + the local query-head count H. Blocks of 32 with UE8M0 scales implement + MXFP8; FP32 scales use unrounded amax/448. Both use an amax floor of + 1e-4. A block never crosses a head boundary. + """ + (device, profile, spec, sections, byte_end, _, output_quant) = ( + _resolve_sparse_mla_plan( + device, + batch_size, + num_heads, + max_topk=max_topk, + max_extra_topk=max_extra_topk, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + q_data_type=q_data_type, + kv_layout=kv_layout, + has_sinks=has_sinks, + return_lse=return_lse, + assume_valid_prefix=assume_valid_prefix, + output_quant_block_size=output_quant_block_size, + out_scale_format=out_scale_format, + num_output_groups=num_output_groups, + ) + ) + kernel = spec.kernel + family = profile.family + splits = profile.split_kv + max_rows = batch_size * max_seq_len_q + workspace = self._workspace_buffer + if workspace is None: + workspace = torch.empty(byte_end, device=device, dtype=torch.uint8) + _validate_workspace_buffer(workspace, device=device, required_bytes=byte_end) + buffers = { + name: _workspace_section_view(workspace, section) + for name, section in sections.items() + } + # Finalize the joint distribution in the main kernel or its existing + # split-KV reducer. No per-source outputs or finishing launch are needed. + kernel = copy.copy(kernel) + kernel.output_quant = output_quant + kernel.assume_valid_prefix = assume_valid_prefix + if family == "2cta": + kernel.fuse_sparse_epilogue = splits == 1 + kernel.fuse_sparse_reduction = splits > 1 + else: + kernel.finalize_output = True + kernel.direct_sparse = True + kernel.direct_sparse_capacities = (max_topk, max_extra_topk) + spec = replace( + spec, kernel=kernel, kernel_signature=_mla_kernel_compile_signature(kernel) + ) + compiled = _get_compiled_mla_decode(spec) + static_kernel = copy.copy(kernel) + static_kernel.direct_static_scales = True + compiled_static = _get_compiled_mla_decode( + replace( + spec, + kernel=static_kernel, + kernel_signature=_mla_kernel_compile_signature(static_kernel), + ) + ) + self._prepared.clear() + self._state = dict( + device=device, + batch=batch_size, + heads=num_heads, + max_q=max_seq_len_q, + max_rows=max_rows, + packed=packed_query, + dtype=q_data_type, + ks=max_topk, + kc=max_extra_topk, + kv_layout=kv_layout, + has_sinks=has_sinks, + return_lse=return_lse, + output_quant=output_quant, + buffers=buffers, + workspace=workspace, + compiled=compiled, + compiled_static=compiled_static, + assume_valid_prefix=assume_valid_prefix, + splits=splits, + default_cl=torch.full( + (max_rows,), max_extra_topk, device=device, dtype=torch.int32 + ), + dummy_indices=torch.zeros((max_rows, 1), device=device, dtype=torch.int32), + dummy_cache=torch.zeros((1, 1, 512), device=device, dtype=q_data_type), + default_sinks=torch.full( + (num_heads,), -torch.inf, device=device, dtype=torch.float32 + ), + ) + + self.workspace_size_bytes = byte_end + self._scalar_cache.clear() + + def _scalar(self, value, name, validate): + state = self._state + if isinstance(value, torch.Tensor): + if ( + value.dtype != torch.float32 + or value.device != state["device"] + or value.numel() != 1 + or not value.is_contiguous() + ): + raise ValueError(f"{name} must be a scalar CUDA FP32 tensor") + if validate and ( + not torch.isfinite(value).all().item() or (value <= 0).any().item() + ): + raise ValueError(f"{name} must be positive and finite") + return value.reshape(1) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be positive and finite") + try: + key = struct.unpack("f", struct.pack("f", float(value)))[0] + except OverflowError as error: + raise ValueError(f"{name} must be representable in float32") from error + if not math.isfinite(key) or key <= 0: + raise ValueError(f"{name} must be positive finite float32") + if key not in self._scalar_cache: + self._scalar_cache[key] = torch.full( + (1,), key, device=state["device"], dtype=torch.float32 + ) + return self._scalar_cache[key] + + def _pool(self, cache, name): + state = self._state + if not isinstance(cache, torch.Tensor): + raise TypeError(f"{name} must be a tensor") + if cache.ndim == 4: + axis = 1 if state["kv_layout"] == "HND" else 2 + if cache.shape[axis] != 1: + raise ValueError(f"{name} requires one KV head") + cache = cache.squeeze(axis) + if ( + cache.ndim != 3 + or cache.shape[1] <= 0 + or cache.shape[2] != 512 + or cache.shape[0] <= 0 + ): + raise ValueError( + f"{name} must have shape [pages,page_size,512] with nonempty pages" + ) + page_size = cache.shape[1] + if cache.device != state["device"] or cache.dtype != state["dtype"]: + raise ValueError(f"{name} must match the planned device and dtype") + if ( + cache.stride(-1) != 1 + or cache.stride(1) != 512 + or cache.stride(0) < page_size * 512 + or cache.stride(0) % 512 + or cache.data_ptr() % 16 + ): + raise ValueError( + f"{name} requires aligned compact rows and a page stride divisible by 512 elements" + ) + row_stride = cache.stride(0) // 512 + rows = (cache.shape[0] - 1) * row_stride + page_size + if rows >= 0x7FFFFFFF: + raise ValueError(f"{name} exceeds signed gather coordinate range") + return cache.as_strided((rows, 1, 512), (512, 512, 1)) + + def _source_metadata( + self, indices, lengths, capacity, tokens, name, *, rows, validate + ): + state = self._state + if capacity == 0: + if indices is not None or lengths is not None: + raise ValueError(f"{name} metadata is not expected by the plan") + return state["dummy_indices"][:rows], state["default_cl"][:rows] + if not isinstance(indices, torch.Tensor) or ( + indices.shape != (rows, capacity) + or indices.dtype != torch.int32 + or indices.device != state["device"] + or indices.stride(-1) != 1 + ): + raise ValueError(f"invalid {name} indices shape/dtype/device/stride") + if not isinstance(lengths, torch.Tensor) or ( + lengths.shape != (rows,) + or lengths.dtype != torch.int32 + or lengths.device != state["device"] + or not lengths.is_contiguous() + ): + raise ValueError(f"invalid {name} lengths") + if validate: + if ((lengths < 0) | (lengths > capacity)).any().item(): + raise ValueError(f"{name} length exceeds capacity") + active = ( + torch.arange(capacity, device=state["device"])[None, :] + < lengths[:, None] + ) + if (active & ((indices < -1) | (indices >= tokens))).any().item(): + raise ValueError(f"{name} active index outside pool") + if state["assume_valid_prefix"] and (active & (indices < 0)).any().item(): + raise ValueError( + f"{name} active prefix contains a hole with assume_valid_prefix=True" + ) + return indices, lengths + + @flashinfer_experimental_api + def run( + self, + query, + kv_cache, + indices, + lengths, + extra_kv_cache=None, + *, + extra_indices=None, + extra_lengths=None, + qo_indptr=None, + softmax_scale=512**-0.5, + output_scale=1.0, + sinks=None, + out=None, + lse=None, + token_positions=None, + cos_sin_cache=None, + out_scale=None, + validate=True, + ): + """Attend to caller-selected storage rows in one or two native KV pools. + + ``indices`` and ``extra_indices`` are CUDA INT32 [T,K] tensors, with + T equal to the query token count and K equal to the corresponding + planned max_topk/max_extra_topk. The last dimension is contiguous; + row strides may include padding. Each index addresses its own pool's + flattened storage, including physical page strides. The caller performs + logical-token/page-table conversion and applies causality upstream. + ``lengths`` and ``extra_lengths`` are contiguous CUDA INT32 [T]. Only + entries before each length are considered; -1 entries are masked holes. + Omit both extra tensors when no extra pool was planned. Input widths + and live lengths need no 128-entry padding: the kernel aligns source + execution spans internally and bounds all source reads. + + ``softmax_scale`` and ``output_scale`` are positive finite Python + scalars or scalar CUDA FP32 tensors. They are already combined: + softmax_scale = attention_scale * q_descale * kv_descale; + output_scale = output_multiplier * kv_descale. Both pools share that + KV descale. These multiply QK scores and the weighted stored V values, + respectively; sink logits are not multiplied by softmax_scale. The + output multiplier applies before fused RoPE/quantization. ``out_scale`` + instead holds the generated per-block quantization scales. + Combine dynamic scales in upstream work before each launch/replay and + keep their tensor addresses stable for CUDA Graphs. + + Caller-owned inputs and plan workspace must outlive graph replay, + including tensor scales and sinks, and must not alias writable output + or workspace storage. Q and native KV are used directly; flattening + padded pages only creates a tensor view. Required attention reductions + remain included. Caller-owned LSE must be contiguous FP32 and 16-byte + aligned, including with validate=False. + + With output quantization enabled, ``token_positions`` is contiguous + INT32 with the query's batch/token prefix; ``cos_sin_cache`` is aligned + contiguous FP32[P,64] containing forward [cos(32), sin(32)] rows. Use + absolute query positions and the same layer-specific table as Q RoPE. + Positions must be in [0,P); validate=False leaves this to the caller. + + Return (out, out_scale), or (out, out_scale, lse) with return_lse=True. + Out is the [T,G,C] view (usual strides (C,T*C,1)) of a 16-byte aligned + contiguous group-major E4M3[G,T,C] tensor, where T is total query + tokens and C=(H/G)*512, in natural head/channel order. Equivalent + singleton-axis strides are accepted. A live prefix may retain a + larger [G,T_capacity,C] allocation; only its live T rows are written. Pass + ``torch.empty((G,T,C), dtype=torch.float8_e4m3fn).transpose(0,1)`` as + ``out`` to preallocate. For K=output_quant_block_size + and N=C/K, FP32 scales have logical shape [T,G,N]; UE8M0 scales have + shape [T,G,N/4], INT32 packing four block exponent bytes, first in the + least significant byte (scale=2**(byte-127)). Allocate scale backing as + 16-byte aligned [G,columns,pad4(T)] and use backing.permute(2,0,1)[:T], + with strides (1,columns*pad4(T),pad4(T)). Padding is not written. The scalar + ``output_scale`` multiplier remains applied before quantization. + + Preallocate out/out_scale/lse and warm up before graph capture. With + fusion disabled, the existing BF16 output and return contract apply. + """ + state = self._state + if state is None: + raise RuntimeError("plan() must be called before run()") + lse_arg = lse + prepared_key = ( + kv_cache.data_ptr(), + kv_cache.shape, + kv_cache.stride(), + kv_cache.dtype, + None + if extra_kv_cache is None + else ( + extra_kv_cache.data_ptr(), + extra_kv_cache.shape, + extra_kv_cache.stride(), + extra_kv_cache.dtype, + ), + None if sinks is None else sinks.data_ptr(), + softmax_scale, + output_scale, + ) + if ( + not validate + and out is not None + and lse is None + and qo_indptr is None + and token_positions is None + and cos_sin_cache is None + and out_scale is None + ): + prepared = self._prepared.get(prepared_key) + if prepared is not None: + return self._run_prepared( + prepared, query, indices, lengths, extra_indices, extra_lengths, out + ) + _validate_query( + query, + packed_query=state["packed"], + device=state["device"], + batch_size=state["batch"], + num_heads=state["heads"], + max_seq_len_q=state["max_q"], + q_dtype=state["dtype"], + query_dim=512, + ) + prefix = tuple(query.shape[:-2]) + rows = math.prod(prefix) + if state["packed"]: + _validate_qo_indptr( + qo_indptr, device=state["device"], batch_size=state["batch"] + ) + if validate: + maximum, total, _ = _derive_max_seq_len_q( + qo_indptr, batch_size=state["batch"] + ) + if total != rows or maximum > state["max_q"]: + raise ValueError("invalid packed query offsets") + elif qo_indptr is not None: + raise ValueError("fixed queries do not accept qo_indptr") + + primary = self._pool(kv_cache, "kv_cache") + if extra_kv_cache is None: + if state["kc"] or extra_indices is not None or extra_lengths is not None: + raise ValueError("extra pool and indices are required by the plan") + extra = state["dummy_cache"] + else: + extra = self._pool(extra_kv_cache, "extra_kv_cache") + + si, sl = self._source_metadata( + indices, + lengths, + state["ks"], + primary.shape[0], + "primary", + rows=rows, + validate=validate, + ) + ci, cl = self._source_metadata( + extra_indices, + extra_lengths, + state["kc"], + extra.shape[0], + "extra", + rows=rows, + validate=validate, + ) + if state["has_sinks"] != (sinks is not None): + raise ValueError("sinks presence must match the plan") + if sinks is None: + sinks = state["default_sinks"] + if ( + sinks.shape != (state["heads"],) + or sinks.dtype != torch.float32 + or sinks.device != state["device"] + or not sinks.is_contiguous() + ): + raise ValueError("sinks must be contiguous CUDA FP32[H]") + if validate and torch.isnan(sinks).any().item(): + raise ValueError("sinks must not contain NaN") + scale_tensors = [ + self._scalar(v, n, validate) + for v, n in ( + (softmax_scale, "softmax_scale"), + (output_scale, "output_scale"), + ) + ] + quant = state["output_quant"] + output_quant_args = None + if quant is not None: + out, out_scale, output_quant_args = _prepare_sparse_output_quant( + quant, query, token_positions, cos_sin_cache, out, out_scale, validate + ) + elif any(v is not None for v in (token_positions, cos_sin_cache, out_scale)): + raise ValueError( + "plan output quantization before passing RoPE/scale tensors" + ) + elif out is None: + out = torch.empty( + (*prefix, state["heads"], 512), + device=state["device"], + dtype=torch.bfloat16, + ) + if quant is None: + _validate_out( + out, + device=state["device"], + batch_size=state["batch"], + num_heads=state["heads"], + max_seq_len_q=state["max_q"], + packed_query=state["packed"], + total_q=rows if state["packed"] else None, + output_dtype=torch.bfloat16, + ) + if lse is None: + # A returned result must survive the next call on this wrapper. + # Use plan-owned scratch only when the LSE is not exposed. + lse = ( + torch.empty( + (*prefix, state["heads"]), + device=state["device"], + dtype=torch.float32, + ) + if state["return_lse"] + else state["buffers"]["public_lse"][:rows].view(*prefix, state["heads"]) + ) + if ( + lse.shape != (*prefix, state["heads"]) + or lse.dtype != torch.float32 + or lse.device != state["device"] + or not lse.is_contiguous() + or lse.data_ptr() % 16 + ): + raise ValueError( + "lse must be contiguous, 16-byte-aligned FP32 with one value per query/head" + ) + static_scales = all( + not isinstance(v, torch.Tensor) for v in (softmax_scale, output_scale) + ) + # Both paths consume the same effective scales; no Q/K/V scale + # composition or auxiliary preparation is performed here. + bmm1_scale = float(softmax_scale) if static_scales else 1.0 + bmm2_scale = float(output_scale) if static_scales else 1.0 + sparse_inputs = ( + si, + ci, + sl, + cl, + scale_tensors[0], + scale_tensors[1], + sinks, + ) + if rows: + buffers = state["buffers"] + # Quantized stores use output_quant_args, including group/bucket + # strides. The core O view supplies compact launch geometry only. + core_out = ( + out.view(rows, 1, state["heads"], 512) + if quant is None + else output_quant_args[0][: rows * state["heads"] * 512].view( + rows, 1, state["heads"], 512 + ) + ) + _launch_mla_decode( + _MLARuntime( + query.view(rows, 1, state["heads"], 512), + primary, + core_out, + primary.shape[0], + bmm1_scale, + bmm2_scale, + extra_cache=extra, + sparse_inputs=sparse_inputs, + output_quant_args=output_quant_args, + ), + block_tables=si, + seq_lens=sl, + qo_indptr=None, + packed_query=False, + kv_lora_rank=512, + split_kv=state["splits"], + workspace=_MLAWorkspaceViews( + buffers["core"] if buffers["core"].numel() else None, + lse.view(rows, 1, state["heads"]), + ), + compiled=state["compiled_static"] + if static_scales + else state["compiled"], + ) + if ( + rows + and quant is None + and static_scales + and not state["packed"] + and not state["return_lse"] + and state["max_q"] == 1 + and lse_arg is None + ): + # Static for this (plan, pools, scales): pool views, scalars, sinks, split, + # workspace. The next validate=False launch with the same pools re-forms + # only the query / index / output views (see _run_prepared). + if len(self._prepared) >= _PREPARED_LAUNCH_LIMIT: + self._prepared.clear() + self._prepared[prepared_key] = dict( + primary=primary, + extra=extra, + primary_rows=primary.shape[0], + c_latent=primary.permute(1, 2, 0), + c_rope=extra.permute(1, 2, 0), + scalars=(scale_tensors[0], scale_tensors[1], sinks), + bmm=(bmm1_scale, bmm2_scale), + ks=state["ks"], + kc=state["kc"], + dummy_indices=state["dummy_indices"], + default_cl=state["default_cl"], + ) + if quant is not None: + return (out, out_scale, lse) if state["return_lse"] else (out, out_scale) + return (out, lse) if state["return_lse"] else out + + def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_lengths, out): + """Launch with the static arguments recorded by the last full run(). + + Only the per-call tensors are checked and viewed: query and output + geometry, index/length shapes, dtypes and strides. Pool geometry, + scales and sinks were validated by the full path that prepared this + launch, and the caller passed validate=False. Byte-identical to the + full path for the same arguments. + """ + state = self._state + rows, heads, device = state["batch"], state["heads"], state["device"] + if ( + query.shape != (rows, 1, heads, 512) + or query.dtype != state["dtype"] + or query.device != device + or not query.is_contiguous() + or query.data_ptr() % 16 + ): + raise ValueError( + "prepared launch: query must be a contiguous, 16-byte aligned " + "[batch, 1, heads, 512] tensor of the planned dtype on the planned device" + ) + if ( + out.numel() != rows * heads * 512 + or out.dtype != torch.bfloat16 + or out.device != device + or not out.is_contiguous() + or out.data_ptr() % 16 + ): + raise ValueError( + "prepared launch: out must be a contiguous, 16-byte aligned BF16 tensor " + "holding one 512-wide row per query head" + ) + if ( + indices.shape != (rows, prepared["ks"]) + or indices.dtype != torch.int32 + or indices.device != device + or indices.stride(-1) != 1 + or lengths.shape != (rows,) + or lengths.dtype != torch.int32 + or lengths.device != device + or not lengths.is_contiguous() + ): + raise ValueError("prepared launch: invalid primary indices/lengths") + if prepared["kc"]: + if ( + not isinstance(extra_indices, torch.Tensor) + or not isinstance(extra_lengths, torch.Tensor) + or extra_indices.shape != (rows, prepared["kc"]) + or extra_indices.dtype != torch.int32 + or extra_indices.device != device + or extra_indices.stride(-1) != 1 + or extra_lengths.shape != (rows,) + or extra_lengths.dtype != torch.int32 + or extra_lengths.device != device + or not extra_lengths.is_contiguous() + ): + raise ValueError("prepared launch: invalid extra indices/lengths") + ci, cl = extra_indices, extra_lengths + else: + if extra_indices is not None or extra_lengths is not None: + raise ValueError("prepared launch: extra indices without a planned extra pool") + ci, cl = prepared["dummy_indices"][:rows], prepared["default_cl"][:rows] + buffers = state["buffers"] + s0, s1, sinks = prepared["scalars"] + q_kernel = query.permute(2, 3, 1, 0) + lse_kernel = buffers["public_lse"][:rows].view(rows, 1, heads).permute(2, 1, 0) + bmm1, bmm2 = prepared["bmm"] + state["compiled_static"]( + q_kernel, + q_kernel, + prepared["c_latent"], + prepared["c_rope"], + indices.transpose(0, 1), + out.view(rows, 1, heads, 512).permute(2, 3, 1, 0), + lse_kernel, + buffers["core"] if buffers["core"].numel() else None, + state["splits"], + lengths, + None, + (indices, ci, lengths, cl, s0, s1, sinks), + bmm1, + bmm2, + None, + ) + return out + + +def get_prims_ts_sparse_mla_decode_workspace_size(*plan_args, **plan_kwargs): + """Return workspace bytes using the same arguments as wrapper.plan(). + + Resolves the default kernel geometry without allocating GPU scratch or + compiling a kernel. Bind the resulting byte buffer to the wrapper constructor + and call plan() before graph capture; run() is the prepared standalone launch. + """ + _, _, _, _, size_bytes, _, _ = _resolve_sparse_mla_plan(*plan_args, **plan_kwargs) + return size_bytes + + +@flashinfer_experimental_api +def batch_sparse_mla_decode_with_paged_kv_cache( + query, + kv_cache, + indices, + lengths, + extra_kv_cache=None, + *, + extra_indices=None, + extra_lengths=None, + qo_indptr=None, + max_seq_len_q=None, + kv_layout="NHD", + softmax_scale=512**-0.5, + output_scale=1.0, + sinks=None, + out=None, + lse=None, + return_lse=False, + workspace_buffer=None, + assume_valid_prefix=False, + output_quant_block_size=None, + out_scale_format="ue8m0", + num_output_groups=None, + token_positions=None, + cos_sin_cache=None, + out_scale=None, +): + """Eager plan-and-run helper using separate source indices and lengths. + + Use a planned wrapper for CUDA Graph replay. Index conversion and scale + composition remain external. softmax_scale includes attention/Q/K descales; + output_scale includes the V descale and optional output multiplier. + Source contracts, output quantization arguments and return layouts follow + :class:`BatchSparseMLADecodePagedTSWrapper`. + """ + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "plan BatchSparseMLADecodePagedTSWrapper before CUDA Graph capture" + ) + packed = query.ndim == 3 + if packed: + if qo_indptr is None: + raise ValueError("packed queries require qo_indptr") + batch = qo_indptr.numel() - 1 + if max_seq_len_q is None: + max_seq_len_q = max(1, int((qo_indptr[1:] - qo_indptr[:-1]).max().item())) + else: + batch = query.shape[0] + max_seq_len_q = query.shape[1] if max_seq_len_q is None else max_seq_len_q + wrapper = BatchSparseMLADecodePagedTSWrapper(workspace_buffer) + wrapper.plan( + query.device, + batch, + query.shape[-2], + max_topk=indices.shape[-1], + max_extra_topk=0 if extra_indices is None else extra_indices.shape[-1], + max_seq_len_q=max_seq_len_q, + packed_query=packed, + q_data_type=query.dtype, + kv_layout=kv_layout, + has_sinks=sinks is not None, + return_lse=return_lse, + assume_valid_prefix=assume_valid_prefix, + output_quant_block_size=output_quant_block_size, + out_scale_format=out_scale_format, + num_output_groups=num_output_groups, + ) + return wrapper.run( + query, + kv_cache, + indices, + lengths, + extra_kv_cache, + extra_indices=extra_indices, + extra_lengths=extra_lengths, + qo_indptr=qo_indptr, + softmax_scale=softmax_scale, + output_scale=output_scale, + sinks=sinks, + out=out, + lse=lse, + token_positions=token_positions, + cos_sin_cache=cos_sin_cache, + out_scale=out_scale, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py index 0c5cf40dbd4e..dff16046cf7c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py @@ -25,11 +25,12 @@ Patterns can be shared across KV heads or supplied independently for each head. Construction uses one CUDA C++ CTA per route. Q1 maps selected and tail blocks -directly. Q2--Q8 radix-sort at most ``G * (topk + 1)`` selected/tail IDs, -segmented-OR equal-key memberships, and emit only unique pages; work and -temporary storage are independent of the model context length. On SM90 and -newer, the metadata grid releases the prepared attention grid through -programmatic dependent launch (PDL). +directly. Q2--Q8 union at most ``G * (topk + 1)`` selected/tail IDs using a +shared byte/bit map when its model-context footprint fits without adding CTA +waves, or a bounded radix-sort fallback. No global-cache-sized scratch is +required. On SM90 and newer, the metadata grid releases the attention prologue +at entry through programmatic dependent launch (PDL); the consumer wait still +covers metadata completion. """ from __future__ import annotations @@ -189,6 +190,7 @@ class _PrimsTSQTokenKvBlockSparseWorkspaceLayout: attention_workspace_byte_offset: int attention_scratch_bytes: int uses_split_kv: bool + query_major_memberships: bool max_seq_len: int total_bytes: int @@ -246,6 +248,7 @@ class _PrimsTSQTokenKvBlockSparseMetadataPlan: storage_page_size: int max_seq_len_kv: int release_attention_pdl: bool + locator_head_multiplier: int = 1 def run( self, @@ -271,6 +274,7 @@ def run( self.sparse_block_size, self.max_seq_len_kv, self.release_attention_pdl, + self.locator_head_multiplier, ) else: self.metadata_run( @@ -286,6 +290,7 @@ def run( self.sparse_block_size, self.max_seq_len_kv, self.release_attention_pdl, + self.locator_head_multiplier, ) @@ -403,7 +408,7 @@ class _QTokenKvBlockSparsePlanConfig: """Capacity and dtype contract published by the public wrapper.""" batch_size: int - seq_len_q: int + group_size: int num_qo_heads: int num_kv_heads: int head_dim: int @@ -468,7 +473,6 @@ def __init__(self) -> None: def plan( self, batch_size: int, - seq_len_q: int, num_qo_heads: int, num_kv_heads: int, head_dim: int, @@ -479,6 +483,7 @@ def plan( *, device: torch.device | str | int, workspace_buffer: torch.Tensor, + group_size: int = 1, use_packed_q: bool = False, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -491,8 +496,9 @@ def plan( ``batch_size`` is the number of query routes: ``len(qo_indptr)-1`` for packed Q, or ``B * Nq`` for fixed ``[B, Nq, G, Hq, D]`` Q. The planned - ``seq_len_q`` is the maximum packed route length or fixed group size - ``G``. ``kv_block_size`` is the indexer's sparse K/V atom and currently + ``group_size`` is the maximum packed route length or fixed group size + ``G``, defaulting to one (no grouping), not the request sequence length. + ``kv_block_size`` is the indexer's sparse K/V atom and currently supports 4/8/16/32/64/128 tokens; ``page_size`` is the physical cache page. Blocks may cross cache pages. The loader uses fragments of size gcd(kv_block_size, page_size), resolving every fragment independently. @@ -516,7 +522,7 @@ def plan( ) batch_size = _validate_positive_int(batch_size, "batch_size") - seq_len_q = _validate_positive_int(seq_len_q, "seq_len_q") + group_size = _validate_positive_int(group_size, "group_size") block_topk = _validate_positive_int(block_topk, "block_topk") max_seq_len_kv = _validate_positive_int(max_seq_len_kv, "max_seq_len_kv") page_size = _validate_positive_int(page_size, "page_size") @@ -524,7 +530,7 @@ def plan( _validate_head_geometry(num_qo_heads, num_kv_heads) kv_block_size = _validate_sparse_block_size(kv_block_size) _validate_storage_page_size(page_size) - _validate_shape_parameters(batch_size * seq_len_q, block_topk, seq_len_q) + _validate_shape_parameters(batch_size * group_size, block_topk, group_size) if mask_type != "causal": raise ValueError( "QToken-KvBlock-Sparse-Attention requires mask_type='causal'" @@ -551,11 +557,11 @@ def plan( max_num_storage_pages = (max_seq_len_kv + page_size - 1) // page_size layout = _get_prims_ts_q_token_kv_block_sparse_workspace_layout( - batch_size * seq_len_q, + batch_size * group_size, block_topk, max_num_storage_pages, page_size, - seq_len_q, + group_size, max_seq_len_kv=max_seq_len_kv, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -581,7 +587,7 @@ def plan( ) candidate = _QTokenKvBlockSparsePlanConfig( batch_size=batch_size, - seq_len_q=seq_len_q, + group_size=group_size, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, @@ -676,7 +682,7 @@ def run( qo_indptr : torch.Tensor, optional CUDA Int32 packed-route offsets. Required for packed Q and omitted for fixed Q; each route must stay within one request and contain - at most the planned ``seq_len_q`` tokens. + at most the planned ``group_size`` tokens. sm_scale : float, optional Softmax scale, defaulting to ``head_dim**-0.5``. Include Q/K dequantization scales here when using FP8 inputs. @@ -707,7 +713,7 @@ def run( ) if qo_indptr.numel() != config.batch_size + 1: raise ValueError("qo_indptr length must match the planned route count") - if q.shape[0] <= 0 or q.shape[0] > config.batch_size * config.seq_len_q: + if q.shape[0] <= 0 or q.shape[0] > config.batch_size * config.group_size: raise ValueError("packed q exceeds the planned query-token capacity") else: if qo_indptr is not None: @@ -715,7 +721,7 @@ def run( if ( q.ndim != 5 or q.shape[0] * q.shape[1] != config.batch_size - or q.shape[2] != config.seq_len_q + or q.shape[2] != config.group_size ): raise ValueError( "fixed q must be [B,Nq,G,Hq,D] with B*Nq and G matching plan" @@ -725,7 +731,7 @@ def run( num_query_tokens = ( int(q.shape[0]) if config.use_packed_q - else (config.batch_size * config.seq_len_q) + else (config.batch_size * config.group_size) ) expected_ids_shape = ( (num_query_tokens, config.block_topk) @@ -818,9 +824,7 @@ def run( max_seq_len_kv=config.max_seq_len_kv, bmm1_scale=sm_scale, qo_indptr=qo_indptr, - max_seq_len_q=( - config.seq_len_q if config.use_packed_q else None - ), + group_size=config.group_size, sparse_block_size=config.kv_block_size, split_kv=config.split_kv, share_pattern_across_kv_heads=config.share_pattern_across_kv_heads, @@ -974,7 +978,10 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( ) from .decode import ( - _resolve_decode_workspace_layout, + _dtype_key, + _make_decode_workspace_layout, + _resolve_cuda_device, + _resolve_decode_launch_spec, _validate_prims_ts_q_token_kv_block_sparse_group_capacity, ) @@ -1001,6 +1008,50 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( (max_seq_len_kv + sparse_block_size - 1) // sparse_block_size, ) page_capacity = block_capacity * (sparse_block_size // fragment_size) + # Padding is an allocation property, not additional selected KV work. + max_seq_len = ( + block_topk * sparse_block_size + (sparse_block_size - 1) + if group_size == 1 + else page_capacity * fragment_size + ) + max_seq_len = min(max_seq_len, max_seq_len_kv) + if kv_dtype is None: + kv_dtype = q_dtype + if out_dtype is None: + out_dtype = q_dtype + _, device_index = _resolve_cuda_device(device) + spec = _resolve_decode_launch_spec( + device_index, + groups, + num_qo_heads, + num_kv_heads, + head_dim, + fragment_size, + max_seq_len, + group_size, + _dtype_key(q_dtype), + _dtype_key(kv_dtype), + _dtype_key(kv_dtype), + _dtype_key(out_dtype), + "HND", + "causal", + use_packed_q, + -1, + storage_page_size, + True, + True, + split_kv, + share_pattern_across_kv_heads, + ) + # Eligible FP8 G2..G8/page-4 profiles consume one bit-plane per query. + # Other profiles and the standalone metadata API retain page-byte masks. + query_major_memberships = ( + sparse_block_size == 4 + and spec.config.use_fp8_qkv + and spec.config.supports_query_major_memberships + ) + if query_major_memberships: + page_capacity = (page_capacity + 31) // 32 * 32 membership_words = ( 0 if group_size == 1 @@ -1024,38 +1075,11 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( ) seq_lens_numel = metadata_rows seq_lens_bytes = seq_lens_numel * 4 - max_seq_len = ( - block_topk * sparse_block_size + (sparse_block_size - 1) - if group_size == 1 - else page_capacity * fragment_size - ) - max_seq_len = min(max_seq_len, max_seq_len_kv) - if kv_dtype is None: - kv_dtype = q_dtype - if out_dtype is None: - out_dtype = q_dtype - attention_layout = _resolve_decode_workspace_layout( - groups, - num_qo_heads, - num_kv_heads, - head_dim, - fragment_size, - max_seq_len, - group_size, - q_dtype, - kv_dtype, - kv_dtype, + attention_layout = _make_decode_workspace_layout( + spec.scratch_shapes, out_dtype, - "HND", - "causal", - use_packed_q, - -1, - storage_page_size, - device, - use_q_token_kv_block_sparse_route=True, - use_pdl=True, - split_kv=split_kv, - share_pattern_across_kv_heads=share_pattern_across_kv_heads, + use_separate_reduction_kernel=spec.config.use_separate_reduction_kernel, + use_split_kv=spec.config.use_split_kv, ) attention_scratch_bytes = attention_layout.total_bytes attention_workspace_byte_offset = _align_up_q_token_kv_block_sparse_workspace( @@ -1072,6 +1096,7 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( attention_workspace_byte_offset=attention_workspace_byte_offset, attention_scratch_bytes=attention_scratch_bytes, uses_split_kv=attention_layout.uses_split_kv, + query_major_memberships=query_major_memberships, max_seq_len=max_seq_len, total_bytes=attention_workspace_byte_offset + attention_scratch_bytes, ) @@ -1087,7 +1112,7 @@ def get_q_token_kv_block_sparse_workspace_size( max_seq_len_kv: int, o_data_type: Optional[torch.dtype] = None, qo_indptr: Optional[torch.Tensor] = None, - seq_len_q: Optional[int] = None, + group_size: int = 1, kv_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -1096,7 +1121,7 @@ def get_q_token_kv_block_sparse_workspace_size( Packed queries use ``[num_query_tokens, num_qo_heads, head_dim]`` and Int32 ``qo_indptr`` partitions their rows into request-safe routes of - length at most ``seq_len_q``. Fixed queries use + length at most ``group_size``. Fixed queries use ``[batch, num_query_groups, group_size, num_qo_heads, head_dim]`` and do not require query-offset metadata. The first two fixed axes are flattened into the attention route axis without copying tensor storage. @@ -1135,9 +1160,9 @@ def get_q_token_kv_block_sparse_workspace_size( dtype. qo_indptr : torch.Tensor, optional Contiguous CUDA Int32 cumulative offsets for packed query routes. - seq_len_q : int, optional - Maximum packed-route length and QToken-KvBlock-Sparse-Attention group size. It is required when - ``qo_indptr`` is supplied. + group_size : int, optional + Maximum query tokens per group, default one (no grouping). For fixed + Q, this must match its ``G`` dimension; packed final groups may be shorter. kv_block_size : int Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. split_kv : bool @@ -1165,7 +1190,7 @@ def get_q_token_kv_block_sparse_workspace_size( max_seq_len_kv=max_seq_len_kv, out_dtype=o_data_type, qo_indptr=qo_indptr, - max_seq_len_q=seq_len_q, + group_size=group_size, sparse_block_size=kv_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, @@ -1181,7 +1206,7 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( max_seq_len_kv: int, out_dtype: Optional[torch.dtype], qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + group_size: int = 1, sparse_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -1196,25 +1221,20 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( if not isinstance(block_table, torch.Tensor): raise TypeError("block_table must be a torch.Tensor") use_packed_q = qo_indptr is not None - if not use_packed_q and max_seq_len_q is not None: - raise ValueError( - "max_seq_len_q is only valid with packed QToken-KvBlock-Sparse-Attention qo_indptr" - ) if use_packed_q: if query.ndim != 3: raise ValueError( "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" ) - if max_seq_len_q is None: - raise ValueError( - "max_seq_len_q is required with QToken-KvBlock-Sparse-Attention qo_indptr" - ) _validate_q_token_kv_block_sparse_qo_indptr_layout_tensor(qo_indptr) - group_size = int(max_seq_len_q) groups = int(qo_indptr.numel()) - 1 num_query_tokens, num_qo_heads, head_dim = query.shape elif query.ndim == 5: - batch_size, groups_per_request, group_size, num_qo_heads, head_dim = query.shape + batch_size, groups_per_request, query_group_size, num_qo_heads, head_dim = ( + query.shape + ) + if group_size != query_group_size: + raise ValueError("fixed group_size must match query.shape[2]") if batch_size <= 0 or groups_per_request <= 0: raise ValueError( "fixed QToken-KvBlock-Sparse-Attention batch and query-group counts must be positive" @@ -1297,11 +1317,12 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( The membership output has one Int32 word per four page slots, or zero columns for Q1. Q1 uses one CUDA C++ CTA per route to map its selected blocks and causal - tail directly. Q2--Q8 use one CUDA C++ CTA per route to radix-sort the - bounded ``group_size * (block_topk + 1)`` candidates, unique them while - OR-reducing membership bits, and map the resulting logical pages through - the dense block table. The terminal metadata grid releases a following - PDL-capable QToken-KvBlock-Sparse-Attention attention launch when requested by the combined API. + tail directly. Q2--Q8 union the bounded + ``group_size * (block_topk + 1)`` candidates with a shared map or radix + sort, OR membership bits, and map the sorted unique logical blocks through + the dense block table. A map is used only when it fits shared memory and + requires no extra CTA waves. The combined API releases attention at + metadata entry; the consumer wait still covers every output store. Packed ``qo_indptr`` must be an Int32 device copy of CPU-validated route offsets; this builder checks only its structural tensor contract and does not read route values back to the host. @@ -1385,6 +1406,7 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( sparse_block_size, max_seq_len_kv, release_attention_pdl, + 1, ) else: metadata_module.run_fixed( @@ -1400,6 +1422,7 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( sparse_block_size, max_seq_len_kv, release_attention_pdl, + 1, ) return outputs @@ -1429,7 +1452,7 @@ def _build_q_token_kv_block_sparse_metadata( physical_page * (storage_page_size/F) + subpage. No CSR conversion is needed. Q1 uses one CUDA C++ direct-mapping kernel. Q2--Q8 use one CUDA C++ - radix-sort and union kernel per route. Neither path requires + map-or-sort union kernel per route. Neither path requires caller-provided scratch. Advanced callers that capture this raw path must preallocate ``out``, warm the same metadata geometry once before capture, and retain every tensor at a stable address through replay. The three @@ -1515,8 +1538,18 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( max_seq_len_kv: int, sparse_block_size: int = 4, qo_indptr: Optional[torch.Tensor] = None, + query_major_memberships: bool = False, + locator_head_multiplier: int = 1, ) -> _PrimsTSQTokenKvBlockSparseMetadataPlan: - """Freeze validated metadata tensors and launch constants.""" + """Freeze validated metadata tensors and launch constants. + + Query-major memberships use a private logical [row, KV tile, query slot] + layout: each uint32 masks 32 four-token union fragments for one query. + The rank-two tensor flattens eight slots per tile, with unused slots and + tail bits zero. This encoding must match the attention plan. Standalone + metadata retains per-fragment membership bytes; Q/K/V layouts, the union, + and page indices are unchanged. + """ sparse_block_size = _validate_sparse_block_size(sparse_block_size) rows = block_indices.shape[0] @@ -1532,9 +1565,16 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( ) release_attention_pdl = device_support_pdl(block_indices.device) metadata_module = _get_prims_ts_q_token_kv_block_sparse_metadata_module() - metadata_run = ( - metadata_module.run_packed if use_packed_q else metadata_module.run_fixed - ) + if query_major_memberships: + metadata_run = ( + metadata_module.run_packed_query_major + if use_packed_q + else metadata_module.run_fixed_query_major + ) + else: + metadata_run = ( + metadata_module.run_packed if use_packed_q else metadata_module.run_fixed + ) return _PrimsTSQTokenKvBlockSparseMetadataPlan( q_token_kv_block_sparse_page_indices=q_token_kv_block_sparse_page_indices, q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, @@ -1547,6 +1587,7 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( storage_page_size=storage_page_size, max_seq_len_kv=max_seq_len_kv, release_attention_pdl=release_attention_pdl, + locator_head_multiplier=locator_head_multiplier, ) @@ -1564,7 +1605,7 @@ def _prepare_q_token_kv_block_sparse_attention( bmm1_scale: Optional[float] = None, bmm2_scale: float = 1.0, qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + group_size: int = 1, sparse_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -1572,7 +1613,7 @@ def _prepare_q_token_kv_block_sparse_attention( """Prepare one graph-stable metadata-plus-attention QToken-KvBlock-Sparse-Attention launch. Packed Q/output use ``[num_query_tokens, Hq, D]`` and describe Q1--Q8 - routes with ``qo_indptr`` and ``max_seq_len_q``. A request's final route + routes with ``qo_indptr`` and ``group_size``. A request's final route may be shorter than the maximum group size. Fixed Q/output use ``[B, Nq, G, Hq, D]`` and omit ``qo_indptr``; ``B * Nq`` becomes the internal attention route count through a zero-copy view. @@ -1634,9 +1675,9 @@ def _prepare_q_token_kv_block_sparse_attention( qo_indptr : torch.Tensor, optional Contiguous CUDA Int32 cumulative route offsets selecting packed-query mode. - max_seq_len_q : int, optional - Maximum packed-route length and QToken-KvBlock-Sparse-Attention group size. It is required when - ``qo_indptr`` is supplied. + group_size : int, optional + Maximum query tokens per group, default one (no grouping). For fixed + Q, this must match its ``G`` dimension; packed final groups may be shorter. sparse_block_size : int Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. @@ -1666,11 +1707,6 @@ def _prepare_q_token_kv_block_sparse_attention( raise ValueError( "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" ) - if max_seq_len_q is None: - raise ValueError( - "max_seq_len_q is required with QToken-KvBlock-Sparse-Attention qo_indptr" - ) - group_size = int(max_seq_len_q) num_query_tokens = int(query.shape[0]) _validate_q_token_kv_block_sparse_qo_indptr_tensor( qo_indptr, @@ -1687,7 +1723,8 @@ def _prepare_q_token_kv_block_sparse_attention( raise ValueError( "fixed QToken-KvBlock-Sparse-Attention query and output must be contiguous" ) - group_size = int(query.shape[2]) + if group_size != query.shape[2]: + raise ValueError("fixed group_size must match query.shape[2]") num_query_tokens = int(query.shape[0] * query.shape[1]) * group_size attention_query = _flatten_fixed_q_token_kv_block_sparse_groups( query, group_size @@ -1724,26 +1761,12 @@ def _prepare_q_token_kv_block_sparse_attention( max_seq_len_kv=max_seq_len_kv, out_dtype=out.dtype, qo_indptr=qo_indptr, - max_seq_len_q=max_seq_len_q, + group_size=group_size, sparse_block_size=sparse_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, ) views = layout.bind(workspace_buffer) - metadata_plan = _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( - block_indices, - block_table, - token_to_request, - query_positions, - views.q_token_kv_block_sparse_page_indices, - views.q_token_kv_block_sparse_page_memberships, - views.seq_lens, - group_size=group_size, - storage_page_size=int(k_cache.shape[2]), - max_seq_len_kv=max_seq_len_kv, - sparse_block_size=sparse_block_size, - qo_indptr=qo_indptr, - ) from .decode import _prepare_prims_ts_batch_decode_plan, _validate_scale scale_qk = _validate_scale( @@ -1751,6 +1774,8 @@ def _prepare_q_token_kv_block_sparse_attention( "bmm1_scale", ) scale_v = _validate_scale(bmm2_scale, "bmm2_scale") + # Generic FMHA sees one query group as a route, so its route-local + # sequence-length parameters receive G, not the full request length. attention_plan, _ = _prepare_prims_ts_batch_decode_plan( attention_query, paged_kv_cache, @@ -1768,6 +1793,7 @@ def _prepare_q_token_kv_block_sparse_attention( kv_layout="HND", page_size=math.gcd(sparse_block_size, int(k_cache.shape[2])), direct_q1_sparse_block_size=sparse_block_size, + query_major_memberships=layout.query_major_memberships, use_q_token_kv_block_sparse_route=True, use_pdl=True, split_kv=split_kv, @@ -1781,6 +1807,25 @@ def _prepare_q_token_kv_block_sparse_attention( else () ), direct_q1_max_model_len=max_seq_len_kv if group_size == 1 else None, + allow_head_aware_locators=group_size > 1, + ) + # Preparation only binds shapes and storage; attention does not read these + # outputs until metadata runs. Use its resolved format for the producer. + metadata_plan = _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( + block_indices, + block_table, + token_to_request, + query_positions, + views.q_token_kv_block_sparse_page_indices, + views.q_token_kv_block_sparse_page_memberships, + views.seq_lens, + group_size=group_size, + storage_page_size=int(k_cache.shape[2]), + max_seq_len_kv=max_seq_len_kv, + sparse_block_size=sparse_block_size, + qo_indptr=qo_indptr, + query_major_memberships=layout.query_major_memberships, + locator_head_multiplier=attention_plan._flat_locator_heads, ) return _PrimsTSQTokenKvBlockSparsePlan( _metadata_plan=metadata_plan, @@ -1808,7 +1853,7 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( workspace_buffer: torch.Tensor, *, max_seq_len_kv: int, - seq_len_q: Optional[int] = None, + group_size: int = 1, kv_block_size: int = 4, mask_type: Literal["causal"] = "causal", sm_scale: Optional[float] = None, @@ -1827,9 +1872,10 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( wrapper, call :meth:`QTokenKvBlockSparsePagedTSWrapper.plan` outside capture, run once eagerly, and capture only :meth:`~QTokenKvBlockSparsePagedTSWrapper.run`. - Packed Q is ``[total_q, Hq, D]`` with ``qo_indptr``; ``seq_len_q`` + Packed Q is ``[total_q, Hq, D]`` with ``qo_indptr``; ``group_size`` is its maximum route length. Fixed Q is ``[B, Nq, G, Hq, D]`` without - ``qo_indptr``, and ``G`` determines ``seq_len_q``. The indexer + ``qo_indptr``, and ``group_size`` must match ``G``. The default is one + (no grouping) for both layouts. The indexer output is named ``indexer_block_ids`` to distinguish it from BSR ``block_indices``. ``kv_block_size`` is the semantic indexer atom (4/8/16/32/64/128), while the physical cache extent is inferred from @@ -1862,8 +1908,9 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( :func:`get_q_token_kv_block_sparse_workspace_size`. max_seq_len_kv : int Static maximum logical K/V length in tokens. - seq_len_q : int, optional - Maximum packed route length. Fixed Q derives it from ``G``. + group_size : int, optional + Maximum query tokens per group, default one (no grouping). Must match + ``G`` for fixed Q; packed final groups may contain fewer tokens. kv_block_size : int Semantic sparse K/V block size: 4/8/16/32/64/128, default four. mask_type : {"causal"} @@ -1918,24 +1965,19 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( raise ValueError( "fixed q must be [B,Nq,G,Hq,D], or supply qo_indptr for packed q" ) - fixed_seq_len_q = int(q.shape[2]) - if seq_len_q is not None and seq_len_q != fixed_seq_len_q: - raise ValueError("fixed seq_len_q must match q.shape[2]") - seq_len_q = fixed_seq_len_q + if group_size != q.shape[2]: + raise ValueError("fixed group_size must match q.shape[2]") batch_size = int(q.shape[0] * q.shape[1]) use_packed_q = False else: if q.ndim != 3: raise ValueError("packed q must have shape [total_q,Hq,D]") - if seq_len_q is None: - raise ValueError("seq_len_q is required with packed qo_indptr") batch_size = int(qo_indptr.numel()) - 1 use_packed_q = True wrapper = QTokenKvBlockSparsePagedTSWrapper() wrapper.plan( batch_size, - seq_len_q, int(q.shape[-2]), int(k_cache.shape[1]), int(q.shape[-1]), @@ -1945,6 +1987,7 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( max_seq_len_kv, device=q.device, workspace_buffer=workspace_buffer, + group_size=group_size, use_packed_q=use_packed_q, mask_type=mask_type, split_kv=split_kv, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 67bde65e7227..4ed9a04e968b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -241,6 +241,18 @@ def __post_init__(self): capture_graph=capture_graph, ) + # Absolute token positions feed the PrimTS sparse MLA FMHA (window + # lengths and the fused output RoPE). The DSA base allocates them only + # for the fp8_ds_mla layout, so provide them for every DSv4 cache mode. + if getattr(self, "token_positions_cuda", None) is None: + self.token_positions_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="token_positions_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # Compute max_compressed_indices for CUDA graph compatibility. # For ratio=4, the indexer selects index_topk compressed tokens. # For ratio=128, use max_seq_len / 128 rounded up to next power of 2 @@ -309,6 +321,28 @@ def __post_init__(self): ) for compress_ratio in self.compress_ratio_set } + # Live per-source lengths for the dual-pool FMHA libraries, derived once + # per step next to sparse_mla_topk_lens: SWA rows live per token + # (min(pos + 1, window)) and compressed rows live per compress ratio + # (sparse_mla_topk_lens[ratio] - window). Shape: [max_num_tokens]. + self.sparse_mla_swa_lens = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="sparse_mla_swa_lens", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.sparse_mla_extra_lens = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name=f"sparse_mla_extra_lens_{compress_ratio}", + dtype=torch.int32, + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + if compress_ratio > 1 + } # cached_token_lens_cuda: number of tokens already cached per request self.cached_token_lens_cuda = self.get_empty( @@ -504,7 +538,7 @@ def prepare_for_deepseek_v4_indices(self, token_positions=None): offsets = token_idx - cu_seq_lens[req_idx].to(torch.int32) token_positions = cached_tokens[req_idx].to(torch.int32) + offsets - if self.use_fp8_ds_mla: + if self.token_positions_cuda is not None: self.token_positions_cuda[: token_positions.shape[0]] = token_positions self._prepare_deepseek_v4_indices_compiled( @@ -516,6 +550,8 @@ def prepare_for_deepseek_v4_indices(self, token_positions=None): self.compressed_local_indices_cuda, self.sparse_mla_topk_lens, self._compress_ratios_sorted, + self.sparse_mla_swa_lens, + self.sparse_mla_extra_lens, ) @staticmethod @@ -529,8 +565,10 @@ def _prepare_deepseek_v4_indices_compiled( compressed_local_indices_buf: torch.Tensor, sparse_mla_topk_lens_bufs: Dict[int, torch.Tensor], compress_ratios: list, + swa_lens_buf: torch.Tensor, + extra_lens_bufs: Dict[int, torch.Tensor], ): - """Build SWA indices, compressed indices, and topk_lens in one fused graph.""" + """Build SWA indices, compressed indices, topk_lens and live lengths in one fused graph.""" device = token_positions.device num_tokens = token_positions.shape[0] @@ -568,6 +606,7 @@ def _prepare_deepseek_v4_indices_compiled( # For ratio>1, the SWA region always occupies exactly window_size (128) # slots. Invalid SWA positions are padded with -1 in the index buffer. kv_lens = token_positions + 1 + swa_lens_buf[:num_tokens] = kv_lens.clamp(max=window_size).to(torch.int32) for compress_ratio in compress_ratios: if compress_ratio == 1: total_count = kv_lens.clamp(max=window_size) @@ -580,6 +619,8 @@ def _prepare_deepseek_v4_indices_compiled( else: raise ValueError(f"Unsupported compress_ratio: {compress_ratio}") sparse_mla_topk_lens_bufs[compress_ratio][:num_tokens] = total_count.to(torch.int32) + if compress_ratio > 1: + extra_lens_bufs[compress_ratio][:num_tokens] = compressed_count.to(torch.int32) def _build_cache_buffer_data_pointers( self, manager: "DeepseekV4CacheManager", compress_ratios_by_layer: list[int] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 94cea2e5c098..368fc84c3750 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -107,7 +107,7 @@ l0_b200: # by absolute index, and they gate on >= 2 GPUs plus Blackwell. They also # rebind mpi4py's serializer at import time, which applies process-wide on # every stage that collects them, including where the tests themselves skip. - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.py TIMEOUT (90) + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.py --ignore=unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py TIMEOUT (90) - unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py # Keep cold-JIT MSA coverage out of the 60-minute attention aggregate. # Cold-JIT runs take about 25 minutes; allow a 45-minute budget. @@ -121,6 +121,7 @@ l0_b200: - unittest/disaggregated/test_deepseek_v4_kv_transfer.py - unittest/disaggregated/test_fp4_mla_kv_transfer.py - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py TIMEOUT (30) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index b04787ed8944..9d6cc8cf66b6 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -48,6 +48,12 @@ from tensorrt_llm.mapping import Mapping +@pytest.fixture(autouse=True) +def _strict_precompile_workers(monkeypatch): + """A PrimTS precompile worker failure is an error here, not a logged fallback.""" + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "1") + + @dataclass(kw_only=True, frozen=True) class Scenario: dtype: torch.dtype = torch.bfloat16 diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py new file mode 100644 index 000000000000..5beda56866d2 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py @@ -0,0 +1,1477 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the PrimTS native sparse MLA FMHA library (DeepSeek-V4, SM100).""" + +import gc +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.fmha import registry +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla import ( + PrimsTSSparseMlaFmha, + get_prims_ts_toolchain_unsupported_reason, +) +from tensorrt_llm._utils import get_sm_version + +from ..test_sparse_mla_forward import HAS_FLASH_MLA +from ..test_sparse_mla_forward import ( + test_forward_sparse_mla_unified as _run_unified_sparse_mla_forward, +) + +LIB_NAME = "prims_ts_sparse_mla" +LATENT = 512 +ROPE = 64 +BLOCK = 128 + + +@pytest.fixture(autouse=True) +def _strict_precompile_workers(monkeypatch): + """A precompile worker failure is an error here, not a logged fallback.""" + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "1") + + +pytestmark = pytest.mark.skipif( + get_sm_version() not in (100, 103), reason="PrimTS sparse MLA requires SM100 or SM103" +) + + +def _skip_unless_toolchain() -> None: + """Skip the tests that need the library to actually serve a batch.""" + if get_sm_version() != 100: + pytest.skip("PrimTS sparse MLA serves SM100 only (SM103 is not yet qualified)") + reason = get_prims_ts_toolchain_unsupported_reason() + if reason is not None: + pytest.skip(reason) + + +def test_registry_entry_is_default_before_fallback(monkeypatch) -> None: + """The library ships enabled and ahead of the TRTLLM-Gen fallback. + + Its ``is_available`` restricts it to DeepSeek-V4 on SM100, so being + in the default set changes nothing for other models or architectures. + """ + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + assert registry.FMHA_LIBS[LIB_NAME] is PrimsTSSparseMlaFmha + assert LIB_NAME in registry.DEFAULT_FMHA_LIBS + names = tuple(registry.FMHA_LIBS) + assert names.index(LIB_NAME) < names.index("flashinfer_trtllm_gen") + assert names.index(LIB_NAME) < names.index("fallback") + assert PrimsTSSparseMlaFmha in registry.get_enabled_fmha_lib_classes() + monkeypatch.setenv("TLLM_FMHA_LIBS", f"-{LIB_NAME}") + assert PrimsTSSparseMlaFmha not in registry.get_enabled_fmha_lib_classes() + + +def test_sm103_stays_on_the_fallback_until_qualified(monkeypatch) -> None: + """B300/GB300 report an SM reason; SM100 proceeds to the shape and toolchain checks. + + The vendored kernels admit SM103 through their architecture guard, but they + are signed off on SM100 only, so the adapter keeps SM103 on TRTLLM-Gen. + """ + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + + attn = SimpleNamespace( + use_fp8_ds_mla=False, + kv_cache_dtype="auto", + _uses_nvfp4_compress=False, + kv_lora_rank=448, + qk_rope_head_dim=64, + head_dim=512, + num_heads=16, + ) + monkeypatch.setattr(mod, "get_sm_version", lambda: 103) + assert "SM103" in PrimsTSSparseMlaFmha._unavailable_reason(attn) + monkeypatch.setattr(mod, "get_sm_version", lambda: 90) + assert "requires SM100" in PrimsTSSparseMlaFmha._unavailable_reason(attn) + monkeypatch.setattr(mod, "get_sm_version", lambda: 100) + reason = PrimsTSSparseMlaFmha._unavailable_reason(attn) + assert reason is None or reason == get_prims_ts_toolchain_unsupported_reason() + + +def test_softmax_scale_follows_qk_head_dim() -> None: + """The scale uses ``qk_nope + qk_rope`` like attentionOp.cpp, not the 512 latent. + + DeepSeek-V4 checkpoints have 448 + 64 = 512, so only a narrower nope width + (the 128 + 64 scenario of ``test_deepseek_v4_sparse_mla``) tells the two apart. + """ + attn = SimpleNamespace(qk_nope_head_dim=128, qk_rope_head_dim=64, q_scaling=0.5336) + expected = 1.0 / (0.5336 * math.sqrt(192.0)) + assert PrimsTSSparseMlaFmha._softmax_scale(attn) == pytest.approx(expected, rel=1e-12) + attn = SimpleNamespace(qk_nope_head_dim=448, qk_rope_head_dim=64, q_scaling=1.0) + assert PrimsTSSparseMlaFmha._softmax_scale(attn) == pytest.approx(LATENT**-0.5, rel=1e-12) + + +def test_pool_view_cache_follows_the_owning_manager() -> None: + """A new cache manager gets fresh pool views, even at the old manager's address. + + The views wrap the raw pool base pointer, so a view cached for a freed + manager would alias released memory. The cache is owned through a weak + reference and dropped whenever the owner changes. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType, + ) + + class _Owner: # weak-referenceable stand-ins for TrtllmAttention / the manager + pass + + def make_manager(rows: int) -> _Owner: + manager = _Owner() + manager.pool = torch.empty(rows, 1, LATENT, dtype=torch.bfloat16, device="cuda") + manager.swa_pool_ptr = manager.pool.data_ptr() + manager.compress_pool_ptrs = {} + manager.pp_layers = [0] + manager._compress_ratios = {0: 4} + manager.dtype = torch.bfloat16 + manager.get_buffers = lambda layer, attn_type: manager.pool + return manager + + attn = _Owner() + fmha = PrimsTSSparseMlaFmha(attn) + first = make_manager(8) + view = fmha._pool_view(first, DeepseekV4AttentionType.SWA, 1) + assert view.data_ptr() == first.swa_pool_ptr + assert tuple(view.shape) == (8, 1, LATENT) + assert fmha._pool_view(first, DeepseekV4AttentionType.SWA, 1) is view + + second = make_manager(16) + other = fmha._pool_view(second, DeepseekV4AttentionType.SWA, 1) + assert other is not view + assert other.data_ptr() == second.swa_pool_ptr + assert tuple(other.shape) == (16, 1, LATENT) + assert fmha._pool_views_owner() is second + + del second + gc.collect() + assert fmha._pool_views_owner() is None + third = make_manager(4) + assert tuple(fmha._pool_view(third, DeepseekV4AttentionType.SWA, 1).shape) == (4, 1, LATENT) + + +def test_plans_are_shared_through_the_metadata_cache(monkeypatch) -> None: + """Layers with one static profile share one plan via ``metadata.fmha_plan_caches``. + + The plan owns the kernel workspace, so caching it per layer would keep one + workspace per layer per CUDA-graph geometry; the metadata cache (the + convention of the block-sparse and FlashInfer wrappers) keeps one per + profile, and layers sharing a metadata object run serially. + """ + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + planned = [] + + class _FakeWrapper: + def __init__(self, workspace_buffer=None): + self.workspace_buffer = workspace_buffer + + def plan(self, device, rows, heads, **kwargs): + planned.append((device, rows, heads, kwargs, self.workspace_buffer)) + + def fake_resolve(device, rows, heads, **kwargs): + # The wrapper's resolver returns the workspace byte extent fifth. + return (device, None, None, {}, 1024 * (1 + kwargs["max_extra_topk"] // 256), None, None) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr(mla_decode, "_resolve_sparse_mla_plan", fake_resolve) + + class _Attn: # weak-referenceable stand-in for TrtllmAttention + num_heads = 16 + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=False, enable_heuristic_topk=False + ) + + attn_a, attn_b = _Attn(), _Attn() + layer_a, layer_b = PrimsTSSparseMlaFmha(attn_a), PrimsTSSparseMlaFmha(attn_b) + metadata = SimpleNamespace(fmha_plan_caches={}) + profile = dict( + metadata=metadata, + device=torch.device("cuda"), + rows=4, + extra=256, + q_dtype=torch.bfloat16, + has_sinks=False, + fused_groups=None, + ) + plan_a = layer_a._get_plan(**profile) + plan_b = layer_b._get_plan(**profile) + assert plan_a is plan_b + assert len(planned) == 1 + assert planned[0][3]["max_topk"] == 128 and planned[0][3]["max_extra_topk"] == 256 + # DeepSeek-V4's index producers emit valid prefixes, so the plan promises + # them unless KV-cache offload (which masks non-resident pages to -1 in + # place) or heuristic top-k is on; see test_plan_keeps_the_hole_checking_path. + assert planned[0][3]["assume_valid_prefix"] is True + cache = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.PLAN_CACHE_KEY] + assert list(cache.values()) == [plan_a] + workspace = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.WORKSPACE_CACHE_KEY]["buffer"] + assert planned[0][4] is workspace and workspace.numel() == 2048 + + # A smaller profile reuses the shared workspace; a larger one grows it. + smaller = layer_a._get_plan(**{**profile, "extra": 0}) + assert smaller is not plan_a and planned[1][4] is workspace + larger = layer_b._get_plan(**{**profile, "extra": 512}) + grown = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.WORKSPACE_CACHE_KEY]["buffer"] + assert larger is not plan_a and planned[2][4] is grown and grown.numel() == 3072 + assert len(planned) == 3 and len(cache) == 3 + + +@pytest.mark.parametrize("offload, heuristic_topk", [(True, False), (False, True)]) +def test_plan_keeps_the_hole_checking_path(monkeypatch, offload, heuristic_topk) -> None: + """Offload (pages masked to -1 in place) and the heuristic GVR top-k (pads from its + own candidate count) can leave -1 inside a live prefix, so no prefix promise.""" + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + planned = [] + + class _FakeWrapper: + def __init__(self, workspace_buffer=None): + pass + + def plan(self, device, rows, heads, **kwargs): + planned.append(kwargs) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr( + mla_decode, "_resolve_sparse_mla_plan", lambda *a, **k: (None,) * 4 + (1024, None, None) + ) + + class _Attn: + num_heads = 16 + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=offload, enable_heuristic_topk=heuristic_topk + ) + + attn = _Attn() # the library only holds a weak reference to its attention module + PrimsTSSparseMlaFmha(attn)._get_plan( + metadata=SimpleNamespace(fmha_plan_caches={}), + device=torch.device("cuda"), + rows=4, + extra=256, + q_dtype=torch.bfloat16, + has_sinks=False, + fused_groups=None, + ) + assert planned[0]["assume_valid_prefix"] is False + + +def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: + """Forward [cos(32), sin(32)] rows, the layout both kernels consume.""" + inv_freq = 1.0 / ( + 10000.0 ** (torch.arange(0, ROPE, 2, dtype=torch.float32, device=device) / ROPE) + ) + angles = torch.outer(torch.arange(max_positions, dtype=torch.float32, device=device), inv_freq) + return torch.cat([angles.cos(), angles.sin()], dim=-1).contiguous() + + +def _random_indices(rows: int, capacity: int, pool_rows: int, device: torch.device): + lengths = torch.randint(1, capacity + 1, (rows,), dtype=torch.int32, device=device) + indices = torch.randint(0, pool_rows, (rows, capacity), dtype=torch.int32, device=device) + inactive = torch.arange(capacity, device=device)[None, :] >= lengths[:, None] + indices[inactive] = -1 + return indices, lengths + + +def _two_source_case(rows: int, heads: int, topk: int, extra_topk: int, device, dtype): + """Pools, query and two index sets for one two-source BF16/FP8 attention.""" + pool_rows, extra_rows = 2048, 1024 + pool = torch.randn(pool_rows, 1, LATENT, dtype=torch.bfloat16, device=device).to(dtype) + extra_pool = torch.randn(extra_rows, 1, LATENT, dtype=torch.bfloat16, device=device).to(dtype) + query = (torch.randn(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) * 0.5).to( + dtype + ) + sets = [] + for _ in range(2): + indices, lengths = _random_indices(rows, topk, pool_rows, device) + extra_indices, extra_lengths = _random_indices(rows, extra_topk, extra_rows, device) + sets.append((indices, lengths, extra_indices, extra_lengths)) + return pool, extra_pool, query, sets + + +def test_reloaded_kernel_accepts_the_same_launch(): + """A kernel reloaded from the persistent compile cache runs the adapter's launch. + + CI processes always compile cold, while every later process (and + production restarts) reload the exported TVM-FFI kernels, whose positional + entry has no Python defaults. Drop the in-process kernel memo after the + first run so the second plan takes the reload path, and require identical + output. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(1) + device = torch.device("cuda") + rows, heads, topk, extra_topk = 5, 64, 128, 256 + pool, extra_pool, query, sets = _two_source_case( + rows, heads, topk, extra_topk, device, torch.bfloat16 + ) + indices, lengths, extra_indices, extra_lengths = sets[0] + + def run_fresh_plan() -> torch.Tensor: + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk) + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + wrapper.run( + query, + pool, + indices, + lengths, + extra_pool, + extra_indices=extra_indices, + extra_lengths=extra_lengths, + softmax_scale=LATENT**-0.5, + output_scale=1.0, + out=out, + ) + torch.cuda.synchronize() + return out + + first = run_fresh_plan() + mla_decode._get_compiled_mla_decode.cache_clear() + second = run_fresh_plan() + torch.testing.assert_close(second, first, rtol=0, atol=0) + + +def test_graph_replay_follows_new_indices_and_lengths(): + """A captured run reads the live index and length buffers on replay. + + Warm up eagerly, capture one run with ``validate=False``, swap in a second + index set and compare the replayed output with an eager run on that set. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(2) + device = torch.device("cuda") + rows, heads, topk, extra_topk = 8, 64, 128, 256 + pool, extra_pool, query, sets = _two_source_case( + rows, heads, topk, extra_topk, device, torch.bfloat16 + ) + buffers = [tensor.clone() for tensor in sets[0]] # fixed-address launch inputs + indices, lengths, extra_indices, extra_lengths = buffers + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk) + + def launch(validate: bool) -> None: + wrapper.run( + query, + pool, + indices, + lengths, + extra_pool, + extra_indices=extra_indices, + extra_lengths=extra_lengths, + softmax_scale=LATENT**-0.5, + output_scale=1.0, + out=out, + validate=validate, + ) + + launch(validate=True) # eager warmup on set 0, also prepares the scalar cache + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(validate=False) + for buffer, source in zip(buffers, sets[1]): + buffer.copy_(source) + graph.replay() + torch.cuda.synchronize() + replayed = out.clone() + launch(validate=True) # eager reference on set 1 through the same plan + torch.cuda.synchronize() + torch.testing.assert_close(replayed, out, rtol=0, atol=0) + assert not torch.equal(replayed, torch.zeros_like(replayed)) + + +def test_prepared_launch_matches_the_full_path_per_layer(): + """A validate=False run with the pools and sinks of an earlier full run takes + the prepared launch, byte-identical to the full path, with one prepared + entry per (pools, sinks) key so layers sharing the plan do not evict each + other; a new pool or new sinks address goes back through the full path. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(3) + device = torch.device("cuda") + rows, heads, topk, extra_topk = 6, 64, 128, 256 + pool, extra_pool, query, sets = _two_source_case( + rows, heads, topk, extra_topk, device, torch.bfloat16 + ) + indices, lengths, extra_indices, extra_lengths = sets[0] + sinks_a = torch.randn(heads, dtype=torch.float32, device=device) + sinks_b = torch.randn(heads, dtype=torch.float32, device=device) + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, has_sinks=True) + prepared_launches = [] + original = wrapper._run_prepared + + def counting_run_prepared(*args, **kwargs): + prepared_launches.append(1) + return original(*args, **kwargs) + + wrapper._run_prepared = counting_run_prepared + + def launch(sinks, kv=pool, validate=False, idx=0): + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + wrapper.run( + query, + kv, + sets[idx][0], + sets[idx][1], + extra_pool, + extra_indices=sets[idx][2], + extra_lengths=sets[idx][3], + softmax_scale=LATENT**-0.5, + output_scale=1.0, + sinks=sinks, + out=out, + validate=validate, + ) + torch.cuda.synchronize() + return out + + ref_a = launch(sinks_a, validate=True) # full path records layer A + ref_b = launch(sinks_b, validate=True) # full path records layer B + assert not prepared_launches and len(wrapper._prepared) == 2 + got_a = launch(sinks_a) + got_b = launch(sinks_b) + assert len(prepared_launches) == 2, "both layers must hit their prepared launch" + torch.testing.assert_close(got_a, ref_a, rtol=0, atol=0) + torch.testing.assert_close(got_b, ref_b, rtol=0, atol=0) + assert not torch.equal(got_a, got_b) + # The prepared launch reads the live index set, not the recorded one. + got_a1 = launch(sinks_a, idx=1) + assert len(prepared_launches) == 3 + torch.testing.assert_close(got_a1, launch(sinks_a, validate=True, idx=1), rtol=0, atol=0) + # A different pool storage or sinks address is not prepared: full path. + other_pool = pool.clone() + launch(sinks_a, kv=other_pool) + launch(sinks_a.clone()) + assert len(prepared_launches) == 3 and len(wrapper._prepared) == 4 + # Re-planning drops every prepared launch. + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, has_sinks=True) + assert not wrapper._prepared + + +def test_request_check_runs_once_per_geometry(monkeypatch) -> None: + """The full request re-check runs on the first call and again only when the + request geometry changes; repeated steps of one shape skip it.""" + adapter = PrimsTSSparseMlaFmha.__new__(PrimsTSSparseMlaFmha) + adapter._checked_key = None + reasons = [] + + def fake_reason(self, q, k, v, metadata, fwd): + reasons.append(1) + return None + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "_unsupported_reason", fake_reason) + + def make(rows: int, width: int = 640, fusion: bool = False, ctx: int = 0): + metadata = SimpleNamespace(num_tokens=ctx + rows, num_ctx_tokens=ctx) + indices = torch.zeros(rows, width, dtype=torch.int32) + output = torch.empty(rows, 64 * LATENT, dtype=torch.bfloat16) + fwd = SimpleNamespace( + sparse_runtime_params=SimpleNamespace(sparse_attn_indices=indices), + output=output, + enable_dsv4_epilogue_fusion=fusion, + quant_q_buffer=None, + attention_sinks=None, + output_sf=None, + out_scale=None, + dsv4_inv_rope_cos_sin_cache=None, + ) + return torch.empty(rows, 64 * LATENT, dtype=torch.bfloat16), metadata, fwd + + q, metadata, fwd = make(8) + adapter._check_request(q, None, None, metadata, fwd) + adapter._check_request(q, None, None, metadata, fwd) + adapter._check_request(q, None, None, metadata, fwd) + assert len(reasons) == 1 + q2, _, fwd2 = make(16) # new row count on the same metadata object + adapter._check_request(q2, None, None, metadata, fwd2) + assert len(reasons) == 2 + adapter._check_request(q2, None, None, metadata, fwd2) + assert len(reasons) == 2 + _, _, fwd3 = make(16, fusion=True) + adapter._check_request(q2, None, None, metadata, fwd3) + assert len(reasons) == 3 + adapter._check_request(q2, None, None, metadata, fwd3) + assert len(reasons) == 3 + _, metadata4, fwd4 = make(16, fusion=True) # only the metadata object differs + adapter._check_request(q2, None, None, metadata4, fwd4) + assert len(reasons) == 4 + + def rejecting_reason(self, q, k, v, metadata, fwd): + return "nope" + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "_unsupported_reason", rejecting_reason) + adapter._check_request(q2, None, None, metadata4, fwd4) # same geometry: skipped + q5, metadata5, fwd5 = make(4) + with pytest.raises(RuntimeError, match="nope"): + adapter._check_request(q5, None, None, metadata5, fwd5) + + +def test_prepared_launch_single_pool_matches_the_full_path(): + """SWA-only layers plan no extra pool; their prepared launch must match the + full path and reject extra indices.""" + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(4) + device = torch.device("cuda") + rows, heads, topk = 6, 64, 128 + pool, _, query, sets = _two_source_case(rows, heads, topk, 256, device, torch.bfloat16) + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk) + prepared_launches = [] + original = wrapper._run_prepared + + def counting_run_prepared(*args, **kwargs): + prepared_launches.append(1) + return original(*args, **kwargs) + + wrapper._run_prepared = counting_run_prepared + + def launch(validate: bool, idx: int = 0) -> torch.Tensor: + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + wrapper.run( + query, + pool, + sets[idx][0], + sets[idx][1], + softmax_scale=LATENT**-0.5, + output_scale=1.0, + out=out, + validate=validate, + ) + torch.cuda.synchronize() + return out + + ref = launch(validate=True) + assert not prepared_launches and len(wrapper._prepared) == 1 + got = launch(validate=False) + assert len(prepared_launches) == 1 + torch.testing.assert_close(got, ref, rtol=0, atol=0) + ref1 = launch(validate=True, idx=1) + got1 = launch(validate=False, idx=1) + assert len(prepared_launches) == 2 + torch.testing.assert_close(got1, ref1, rtol=0, atol=0) + with pytest.raises(ValueError, match="extra indices"): + wrapper.run( + query, + pool, + sets[0][0], + sets[0][1], + extra_indices=sets[0][2], + extra_lengths=sets[0][3], + softmax_scale=LATENT**-0.5, + output_scale=1.0, + out=torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device), + validate=False, + ) + + +def test_sinks_pass_through_or_are_prepared_once_per_layer() -> None: + """FP32 per-head sinks are passed through at a stable address; other + dtypes are converted once per source tensor, so the kernel key keyed by the + sinks address stays stable for every layer served by one adapter.""" + from tensorrt_llm._torch.attention.backends.fmha import prims_ts_sparse_mla as mod + + adapter = PrimsTSSparseMlaFmha.__new__(PrimsTSSparseMlaFmha) + adapter._sinks_cache = {} + fp32 = torch.nn.Parameter(torch.randn(64, device="cuda"), requires_grad=False) + assert adapter._sinks(SimpleNamespace(attention_sinks=fp32)) is fp32 + assert not adapter._sinks_cache + bf16_a = torch.randn(64, device="cuda").to(torch.bfloat16) + bf16_b = torch.randn(64, device="cuda").to(torch.bfloat16) + prep_a = adapter._sinks(SimpleNamespace(attention_sinks=bf16_a)) + prep_b = adapter._sinks(SimpleNamespace(attention_sinks=bf16_b)) + assert prep_a.dtype == torch.float32 and prep_a is not prep_b + assert adapter._sinks(SimpleNamespace(attention_sinks=bf16_a)) is prep_a + assert adapter._sinks(SimpleNamespace(attention_sinks=bf16_b)) is prep_b + torch.testing.assert_close(prep_a, bf16_a.float()) + assert adapter._sinks(SimpleNamespace(attention_sinks=None)) is None + assert mod._SINKS_CACHE_LIMIT >= 61 + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) +@pytest.mark.parametrize( + "rows,heads,groups,topk,extra_topk", + [ + (1, 128, 16, 128, 0), + (5, 128, 16, 128, 512), + (64, 128, 16, 128, 512), + (16, 64, 8, 128, 256), + (3, 64, 8, 128, 0), + ], +) +def test_fused_output_matches_trtllm_inverse_rope_quant( + dtype, rows, heads, groups, topk, extra_topk +): + """The kernel's fused epilogue must reproduce TRT-LLM's inverse-RoPE + FP8 quant. + + Reference: the same attention with a BF16 output, followed by + ``trtllm.fused_inv_rope_fp8_quant_vllm_port`` (the op the unfused DSv4 + path runs). The fused kernel quantizes from FP32 accumulators instead of + the BF16 output, so compare dequantized values up to FP8 rounding. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( + BatchSparseMLADecodePagedTSWrapper, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + pool_rows, extra_rows, max_positions = 4096, 2048, 4096 + + def make_pool(n): + return (torch.randn(n, 1, LATENT, dtype=torch.bfloat16, device=device)).to(dtype) + + pool = make_pool(pool_rows) + extra_pool = make_pool(extra_rows) if extra_topk else None + indices, lengths = _random_indices(rows, topk, pool_rows, device) + extra_indices, extra_lengths = ( + (None, None) if extra_topk == 0 else _random_indices(rows, extra_topk, extra_rows, device) + ) + query = (torch.randn(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) * 0.5).to( + dtype + ) + positions = torch.randint(0, max_positions, (rows,), dtype=torch.int32, device=device) + cos_sin = _rope_table(max_positions, device) + + softmax_scale = LATENT**-0.5 + descale = 0.75 if dtype == torch.float8_e4m3fn else 1.0 + # The kernel takes the combined scales (attention x Q descale x KV descale + # on the softmax, the KV descale on the output), as the adapter passes them. + run_scales = dict(softmax_scale=softmax_scale * descale * descale, output_scale=descale) + sources = dict(extra_indices=extra_indices, extra_lengths=extra_lengths) if extra_topk else {} + plan_args = dict(max_topk=topk, max_extra_topk=extra_topk, max_seq_len_q=1, q_data_type=dtype) + + reference = BatchSparseMLADecodePagedTSWrapper() + reference.plan(device, rows, heads, **plan_args) + out_bf16 = reference.run(query, pool, indices, lengths, extra_pool, **sources, **run_scales) + heads_per_group = heads // groups + ref_fp8, ref_scale = torch.ops.trtllm.fused_inv_rope_fp8_quant_vllm_port( + out_bf16.view(rows, heads, LATENT), + positions, + cos_sin, + groups, + heads_per_group, + LATENT - ROPE, + ROPE, + BLOCK, + False, + ) + + fused = BatchSparseMLADecodePagedTSWrapper() + fused.plan( + device, + rows, + heads, + output_quant_block_size=BLOCK, + out_scale_format="fp32", + num_output_groups=groups, + **plan_args, + ) + channels = heads_per_group * LATENT + padded_rows = (rows + 3) // 4 * 4 + # The exact buffers TRT-LLM's fused DSv4 epilogue allocates. + out_buf = torch.empty((groups, rows, channels), dtype=torch.float8_e4m3fn, device=device) + scale_buf = torch.empty( + (groups, channels // BLOCK, padded_rows), dtype=torch.float32, device=device + ) + fused_out, fused_scale = fused.run( + query, + pool, + indices, + lengths, + extra_pool, + **sources, + out=out_buf.transpose(0, 1), + out_scale=scale_buf.permute(2, 0, 1)[:rows], + token_positions=positions.view(rows, 1), + cos_sin_cache=cos_sin, + **run_scales, + ) + assert fused_out.data_ptr() == out_buf.data_ptr() + assert fused_scale.data_ptr() == scale_buf.data_ptr() + assert ref_fp8.shape == out_buf.shape and ref_scale.shape == scale_buf.shape + + ref_scales = ref_scale[..., :rows].permute(0, 2, 1) # [G, T, blocks] + got_scales = scale_buf[..., :rows].permute(0, 2, 1) + torch.testing.assert_close(got_scales, ref_scales, rtol=2e-2, atol=1e-6) + + def dequant(values, scales): + blocks = values.float().view(groups, rows, channels // BLOCK, BLOCK) + return (blocks * scales[..., None]).flatten(-2) + + ref = dequant(ref_fp8, ref_scales) + got = dequant(out_buf, got_scales) + # One E4M3 ulp is 1/8 of the magnitude. The reference rotates the + # BF16-rounded output, so near-cancelling RoPE channels (D[448:512]) carry + # an absolute error of a few BF16 ulps of their O(1) inputs, which the + # kernel (rotating FP32 accumulators) does not have; that shows up as up to + # about one quantization step of the block. Allow two steps there and half + # a step elsewhere, on top of two relative ulps. + block_scales = ref_scales.repeat_interleave(BLOCK, dim=-1) + is_rope = torch.zeros(channels, dtype=torch.bool, device=device) + is_rope.view(-1, LATENT)[:, LATENT - ROPE :] = True + step_allowance = torch.where(is_rope, 2.0, 0.5) * block_scales + ulp = ref.abs() * 0.125 + block_scales * 2.0**-6 + diff = (got - ref).abs() + assert (diff <= 2.0 * ulp + step_allowance + 1e-6).all(), ( + f"max fused-vs-reference gap {diff.max().item()}" + ) + assert diff.mean().item() <= 0.02 * ref.abs().mean().item() + 1e-6 + + +@pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") +@pytest.mark.parametrize("library", ["prims_ts", "fallback"]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize( + "batch_name", + [ + "small_decode", + "medium_decode", + "single_decode", + "small_mixed", + "mixed_heavy_decode", + ], +) +def test_forward_sparse_mla_unified_prims_ts(monkeypatch, batch_name, kv_cache_dtype, library): + """Run the DeepSeek-V4 module test on the default library set and on the opt-out. + + ``prims_ts``: no ``TLLM_FMHA_LIBS`` override; the PrimTS library is in the + default set and must win the generation phase on its own. ``fallback``: + ``TLLM_FMHA_LIBS=-prims_ts_sparse_mla``, the documented opt-out, must hand + generation back to the TRTLLM-Gen fallback (``FallbackFmha``) with PrimTS + never entered, so a regression on that path (what SM103, ``fp8_ds_mla`` and + NVFP4 deployments run) still fails here. + """ + from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + + if library == "prims_ts": + _skip_unless_toolchain() + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + else: + monkeypatch.setenv("TLLM_FMHA_LIBS", f"-{LIB_NAME}") + prims_calls: list = [] + fallback_calls: list = [] + original_forward = PrimsTSSparseMlaFmha.forward + original_fallback = FallbackFmha.forward + + def counting_forward(self, *args, **kwargs): + prims_calls.append(1) + return original_forward(self, *args, **kwargs) + + def counting_fallback(self, q, k, v, metadata, forward_args, *args, **kwargs): + if ( + metadata.num_generations > 0 + and forward_args.attention_input_type != AttentionInputType.context_only + ): + fallback_calls.append(1) + return original_fallback(self, q, k, v, metadata, forward_args, *args, **kwargs) + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", counting_forward) + monkeypatch.setattr(FallbackFmha, "forward", counting_fallback) + _run_unified_sparse_mla_forward(batch_name, kv_cache_dtype, "deepseek_v4") + if library == "prims_ts": + assert prims_calls, "the PrimTS sparse MLA FMHA did not serve the generation phase" + else: + assert not prims_calls, "TLLM_FMHA_LIBS=-prims_ts_sparse_mla did not disable PrimTS" + assert fallback_calls, "the TRTLLM-Gen fallback did not serve the generation phase" + + +@pytest.mark.parametrize( + "generation_seq_len_q", + [2, 6, [1, 2, 4]], + ids=["mtp_next_n_2", "mtp_next_n_6", "ragged_1_2_4"], +) +def test_speculative_decoding_rows_run_on_prims_ts(monkeypatch, generation_seq_len_q): + """MTP generation (``max_draft_tokens > 0``) on the default library set. + + ``test_deepseek_v4_sparse_mla`` drives the DeepSeek-V4 attention module with + several query tokens per request (uniform ``next_n`` widths and a ragged + batch) and checks the output against its reference; here it is run with the + library default and the spy asserts that PrimTS served those steps with one + row per query token (``rows = sum(generation_lens)``, not one per request), + which is the per-token-row contract the kernel is planned for + (``max_seq_len_q = 1``). + """ + _skip_unless_toolchain() + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + from . import test_deepseek_v4_sparse_mla as module_test + + context_lengths = [14, 508, 3947] + generation_lens = ( + generation_seq_len_q + if isinstance(generation_seq_len_q, list) + else [generation_seq_len_q] * len(context_lengths) + ) + rows_seen: list[int] = [] + original_forward = PrimsTSSparseMlaFmha.forward + + def spy_forward(self, q, k, v, metadata, forward_args): + rows_seen.append(int(q.shape[0])) + return original_forward(self, q, k, v, metadata, forward_args) + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", spy_forward) + module_test.test_deepseek_v4_sparse_mla(context_lengths, generation_seq_len_q) + assert rows_seen, "PrimTS did not serve the speculative-decoding generation steps" + assert set(rows_seen) == {sum(generation_lens)}, rows_seen + assert sum(generation_lens) > len(context_lengths) # more rows than requests: draft tokens + + +def _dequant_fused(values: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Dequantize a fused DSv4 epilogue buffer. + + ``values`` is ``[G, T, C]`` E4M3 and ``scales`` is ``[G, C // 128, pad4(T)]`` + FP32, exactly as ``_create_dsv4_epilogue_buffers`` allocates them. + """ + groups, rows, channels = values.shape + sc = scales[..., :rows].permute(0, 2, 1) # [G, T, blocks] + blocks = values.float().view(groups, rows, channels // BLOCK, BLOCK) + return (blocks * sc[..., None]).flatten(-2), sc + + +def _assert_fused_close(got, got_sc, ref, ref_sc, channels, device): + """Same tolerance as the kernel-level fused test: two relative E4M3 ulps plus + half a quantization step, two steps on the RoPE channels D[448:512].""" + torch.testing.assert_close(got_sc, ref_sc, rtol=2e-2, atol=1e-6) + block_scales = ref_sc.repeat_interleave(BLOCK, dim=-1) + is_rope = torch.zeros(channels, dtype=torch.bool, device=device) + is_rope.view(-1, LATENT)[:, LATENT - ROPE :] = True + step_allowance = torch.where(is_rope, 2.0, 0.5) * block_scales + diff = (got - ref).abs() + ulp = ref.abs() / 8.0 + assert (diff <= 2.0 * ulp + step_allowance + 1e-6).all(), ( + f"max diff {diff.max().item():.4g} vs allowance " + f"{(2.0 * ulp + step_allowance).max().item():.4g}" + ) + assert diff.mean().item() <= 0.02 * ref.abs().mean().item() + 1e-6 + + +@pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("batch_name", ["small_decode", "medium_decode", "single_decode"]) +def test_forward_sparse_mla_unified_fused_epilogue(monkeypatch, batch_name, kv_cache_dtype): + """Module-level coverage of the fused DSv4 epilogue (production ADP path). + + ``prepare_sparse_attn_outputs`` only hands out the 3D fused buffer under + attention-DP with 128 heads, an FP8 cache and an FP8 ``o_a_proj``, which the + module test harness does not configure. The generation entry point + ``forward_generation_sparse_attn`` takes ``enable_dsv4_epilogue_fusion`` + directly, so wrap the harness's call to it and, on the same prepared cache, + metadata and inputs, additionally run the fused mode on the default library + set: PrimTS must serve it with ``forward_args.enable_dsv4_epilogue_fusion`` + set, and its fused output must match the unfused BF16 output pushed through + ``trtllm.fused_inv_rope_fp8_quant_vllm_port`` (the path the module takes when + fusion is off). The harness's own reference check still runs on the unfused + call. + + The TRTLLM-Gen fallback is not run in fused mode here: the module only + enables the fused epilogue for it at 128 heads, and at the harness's 64 + heads its C++ op aborts the process rather than raising. + """ + _skip_unless_toolchain() + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + + from .. import test_sparse_mla_forward as harness + + original_gen = harness.forward_generation_deepseek_v4_mla + original_forward = PrimsTSSparseMlaFmha.forward + seen_fused_flags: list[bool] = [] + results: dict[str, object] = {} + + def spy_forward(self, q, k, v, metadata, forward_args): + seen_fused_flags.append(bool(forward_args.enable_dsv4_epilogue_fusion)) + return original_forward(self, q, k, v, metadata, forward_args) + + _INPLACE_INPUTS = ("q", "compressed_kv", "k_pe", "latent_cache") + + def run_fused(mla, kwargs, pristine): + # The RoPE/KV-write prologue rotates ``q`` and ``latent_cache`` in + # place, so every call needs the inputs as they were before the first. + fused_kwargs = dict(kwargs) + fused_kwargs.update({k: v.clone() for k, v in pristine.items()}) + fused_kwargs["output"] = None + fused_kwargs["enable_dsv4_epilogue_fusion"] = True + return original_gen(mla, **fused_kwargs) + + def wrapped_gen(mla, **kwargs): + pristine = {k: kwargs[k].clone() for k in _INPLACE_INPUTS if kwargs.get(k) is not None} + # 1. The harness's unfused call, untouched (it validates against the + # Python reference afterwards). + out = original_gen(mla, **kwargs) + output = kwargs["output"] + rows = output.shape[0] + heads = mla.num_heads_tp + groups = mla.n_local_groups + positions = kwargs["position_ids"].view(-1).to(torch.int32) + cos_sin = mla.inverse_rotary_emb.rotary_cos_sin + torch.cuda.synchronize() + + # Path B: the standalone inverse-RoPE + FP8 quant op on the BF16 output. + ref_fp8, ref_scale = torch.ops.trtllm.fused_inv_rope_fp8_quant_vllm_port( + output.view(rows, heads, LATENT).clone(), + positions, + cos_sin, + groups, + heads // groups, + LATENT - ROPE, + ROPE, + BLOCK, + False, + ) + + # 2. Path A on the default library set: PrimTS serves the fused epilogue. + # (The module hands out the 3D buffer only under ADP + 128 heads + FP8, + # which is what production throughput configs run.) + before = len(seen_fused_flags) + prims_fp8, prims_scale = run_fused(mla, kwargs, pristine) + torch.cuda.synchronize() + assert seen_fused_flags[before:] == [True], ( + "the PrimTS sparse MLA FMHA did not serve the fused generation call" + ) + + results.update( + rows=rows, + heads=heads, + groups=groups, + ref=(ref_fp8.clone(), ref_scale.clone()), + prims=(prims_fp8.clone(), prims_scale.clone()), + ) + return out + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", spy_forward) + monkeypatch.setattr(harness, "forward_generation_deepseek_v4_mla", wrapped_gen) + _run_unified_sparse_mla_forward(batch_name, kv_cache_dtype, "deepseek_v4") + + assert results, "the generation phase never ran" + assert seen_fused_flags.count(False) >= 1, "the unfused call did not reach PrimTS" + device = results["prims"][0].device + channels = results["heads"] // results["groups"] * LATENT + + ref, ref_sc = _dequant_fused(*results["ref"]) + got, got_sc = _dequant_fused(*results["prims"]) + assert got.shape == ref.shape + _assert_fused_close(got, got_sc, ref, ref_sc, channels, device) + + +def test_kernel_variants_are_warmed_once_per_profile(monkeypatch) -> None: + """The first eager forward of a static profile plans one representative row count + per distinct kernel compile signature, largest first, on the eager metadata only; + the same profile is not warmed again in the process and a second extra capacity + (another layer group) warms its own variants.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + planned = [] + + class _FakeWrapper: + def __init__(self, workspace_buffer=None): + pass + + def plan(self, device, rows, heads, **kwargs): + planned.append((rows, kwargs["max_extra_topk"])) + + class _FakeKernel: + def __init__(self, rows): + self.rows = rows + + def compile_signature(self): + return (self.rows // 8,) # one variant per 8 row counts + + def fake_resolve(device, rows, heads, **kwargs): + return (None, None, SimpleNamespace(kernel=_FakeKernel(rows)), None, 1024, None, None) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr(mla_decode, "_resolve_sparse_mla_plan", fake_resolve) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "0") # in-process planning only + + class _Attn: + num_heads = 16 + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=False, enable_heuristic_topk=False + ) + + attn = _Attn() + lib = PrimsTSSparseMlaFmha(attn) + profile = dict( + device=torch.device("cuda"), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None + ) + eager = SimpleNamespace( + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=False, + ) + lib._warm_kernel_variants(metadata=eager, extra=512, **profile) + assert planned == [(20, 512), (15, 512), (7, 512)] # one per signature, largest first + lib._warm_kernel_variants(metadata=eager, extra=512, **profile) + assert len(planned) == 3 # same profile: already warmed in this process + graph = SimpleNamespace( + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=True, + ) + lib._warm_kernel_variants(metadata=graph, extra=32, **profile) + assert len(planned) == 3 # CUDA graph metadata never warms + lib._warm_kernel_variants(metadata=eager, extra=32, **profile) + assert planned[3:] == [(20, 32), (15, 32), (7, 32)] # another layer group warms its own + assert lib._get_plan(metadata=eager, rows=15, extra=512, **profile) is not None + assert len(planned) == 6 # the warmed plan is reused + + +def test_warm_set_covers_the_speculative_decoding_rows(monkeypatch) -> None: + """With MTP the generation step carries ``requests * (1 + draft)`` rows; the warm-up + plans the largest row count of every variant up to that ceiling (bounded by the + token budget), so the MTP row counts never compile during the first serving step.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + + planned: list = [] + _fake_sparse_plan(monkeypatch, planned, signature_div=8) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "0") + attn, profile, _ = _attn_and_profile() + lib = PrimsTSSparseMlaFmha(attn) + + mtp = SimpleNamespace( + fmha_plan_caches={}, + max_num_tokens=256, + max_num_requests=4, + max_total_draft_tokens=5, # DeepSeek-V4 MTP: next_n = 6 rows per request + mapping=None, + is_cuda_graph=False, + ) + lib._warm_kernel_variants(metadata=mtp, extra=512, **profile) + rows = [r for r, _ in planned] + assert rows[0] == 4 * (1 + 5), rows + assert rows == sorted(rows, reverse=True) + assert set(rows) == {24, 23, 15, 7} # the largest row count of each variant + + # The step's token budget bounds the rows even when the MTP ceiling is larger. + planned.clear() + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + budget = SimpleNamespace( + fmha_plan_caches={}, + max_num_tokens=10, + max_num_requests=4, + max_total_draft_tokens=5, + mapping=None, + is_cuda_graph=False, + ) + lib._warm_kernel_variants(metadata=budget, extra=512, **profile) + assert [r for r, _ in planned][0] == 10 + + # A metadata class that does not declare the fields fails loudly. + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + renamed = SimpleNamespace( + fmha_plan_caches={}, max_num_tokens=10, max_num_requests=4, is_cuda_graph=False + ) + with pytest.raises(AttributeError): + lib._warm_kernel_variants(metadata=renamed, extra=512, **profile) + + +def test_index_contract_is_validated_once_per_process(monkeypatch) -> None: + """One validated eager launch per contract (producers + pools); the same contract never + re-validates, because that synchronizes the host inside the context ramp.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + + monkeypatch.setattr(mod, "_VALIDATED_CONTRACTS", set()) + key = ("cuda:0", 128, 512, 4, False, "torch.bfloat16", 1, 4096, 2, 4096) + assert PrimsTSSparseMlaFmha._claim_contract_validation(key) is True + assert PrimsTSSparseMlaFmha._claim_contract_validation(key) is False + other_pools = key[:6] + (3, 4096, 4, 4096) + assert PrimsTSSparseMlaFmha._claim_contract_validation(other_pools) is True + + +def _fake_sparse_plan(monkeypatch, planned, *, signature_div=8): + """Replace the sparse wrapper and plan resolver with fakes that record plan() calls.""" + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + class _FakeWrapper: + def __init__(self, workspace_buffer=None): + pass + + def plan(self, device, rows, heads, **kwargs): + planned.append((rows, kwargs["max_extra_topk"])) + + class _FakeKernel: + def __init__(self, rows): + self.rows = rows + + def compile_signature(self): + return (self.rows // signature_div,) + + def fake_resolve(device, rows, heads, **kwargs): + return (None, None, SimpleNamespace(kernel=_FakeKernel(rows)), None, 1024, None, None) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr(mla_decode, "_resolve_sparse_mla_plan", fake_resolve) + return mla_decode + + +def _attn_and_profile(): + class _Attn: + num_heads = 16 + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=False, enable_heuristic_topk=False + ) + + attn = _Attn() # the library holds a weak reference to its owner + profile = dict( + device=torch.device("cuda", 0), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None + ) + eager = SimpleNamespace( + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=False, + ) + return attn, profile, eager + + +def test_warm_up_fans_out_only_when_a_variant_is_missing(monkeypatch) -> None: + """The warm-up probes the cache first: all cached -> no workers; a miss -> every + representative row count goes to the workers as a JSON-safe payload before the + in-process plan loop; a failing helper never breaks the warm-up; strict mode raises.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + + planned, fanned = [], [] + _fake_sparse_plan(monkeypatch, planned) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "2") + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + + def fake_workers(payloads, workers): + fanned.append((list(payloads), workers)) + return True + + monkeypatch.setattr(pp, "run_precompile_workers", fake_workers) + attn, profile, eager = _attn_and_profile() + lib = PrimsTSSparseMlaFmha(attn) + + # 1. Every variant cached: the probe plans (under probe mode) and no worker runs. + lib._warm_kernel_variants(metadata=eager, extra=512, **profile) + assert fanned == [] + # probe planned the three variants, then the real loop planned them again + assert planned == [(20, 512), (15, 512), (7, 512)] * 2 + + # 2. A cache miss during the probe -> workers get every representative row count. + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + real_plan = mla_decode.BatchSparseMLADecodePagedTSWrapper.plan + + def probing_plan(self, device, rows, heads, **kwargs): + if cc._PRECOMPILE_MODE.get() == "probe": + raise cc.KernelNotCached("missing") + return real_plan(self, device, rows, heads, **kwargs) + + monkeypatch.setattr(mla_decode.BatchSparseMLADecodePagedTSWrapper, "plan", probing_plan) + planned.clear() + lib._warm_kernel_variants(metadata=eager, extra=32, **profile) + assert planned == [(20, 32), (15, 32), (7, 32)] + assert len(fanned) == 1 + payloads, workers = fanned[0] + assert workers == 2 # capped by the env value, below the three variants + assert [p["rows"] for p in payloads] == [20, 15, 7] + assert payloads[0]["plan_kwargs"]["q_data_type"] == "bfloat16" # JSON-safe dtype + assert payloads[0]["plan_kwargs"]["max_extra_topk"] == 32 + assert payloads[0]["device_index"] == 0 + import json + + json.dumps(payloads) # every payload field crosses the subprocess boundary + + # 3. A failing helper is logged and the in-process loop still plans everything... + def broken_workers(payloads, workers): + raise OSError("no subprocesses here") + + monkeypatch.setattr(pp, "run_precompile_workers", broken_workers) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "0") # the module fixture sets strict mode + planned.clear() + fresh = SimpleNamespace( # a new metadata object: the earlier plans are cached on `eager` + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=False, + ) + lib._warm_kernel_variants(metadata=fresh, extra=32, **profile) + assert planned == [(20, 32), (15, 32), (7, 32)] + # ...unless strict mode asks for the failure. + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "1") + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + with pytest.raises(OSError): + lib._warm_kernel_variants(metadata=eager, extra=32, **profile) + + +def test_precompile_worker_count_follows_the_env(monkeypatch) -> None: + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + + monkeypatch.delenv("TLLM_PRIMTS_WARM_WORKERS", raising=False) + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + assert pp.worker_count(12) == 4 # default for a single rank + assert pp.worker_count(3) == 3 # never more workers than variants + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "0") + assert pp.worker_count(12) == 0 + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "8") + assert pp.worker_count(12) == 8 # explicit values are never capped + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "many") + assert pp.worker_count(12) == 4 # unparsable: default with a warning + monkeypatch.setenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "1") + assert pp.worker_count(12) == 0 # nothing to persist into + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE") + monkeypatch.setenv("FLASHINFER_DISABLE_JIT", "1") + assert pp.worker_count(12) == 0 # compiling is forbidden + + +def test_default_worker_count_is_capped_by_the_ranks_on_the_host(monkeypatch) -> None: + """Under the default launcher every rank has a private FlashInfer workspace and fans + out on its own, so the per-rank default shrinks with the ranks on the host: the + host never starts more than 16 compile subprocesses at the default, and an + explicit ``TLLM_PRIMTS_WARM_WORKERS`` still wins.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + + monkeypatch.delenv("TLLM_PRIMTS_WARM_WORKERS", raising=False) + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + expected = {1: 4, 2: 4, 4: 4, 8: 2, 16: 1, 32: 1} + for ranks, workers in expected.items(): + assert pp.default_worker_count(ranks) == workers, ranks + assert pp.worker_count(12, local_ranks=ranks) == workers, ranks + assert ranks * workers <= 16 or workers == 1 + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "4") + assert pp.worker_count(12, local_ranks=8) == 4 # explicit values are never capped + + # The rank count comes from the model's Mapping (no MPI call inside the + # forward), else from the launcher's environment, else 1. + for name in pp._LOCAL_RANK_ENV: + monkeypatch.delenv(name, raising=False) + assert pp.local_rank_count(SimpleNamespace(world_size=8, gpus_per_node=8)) == 8 + assert pp.local_rank_count(SimpleNamespace(world_size=16, gpus_per_node=8)) == 8 + assert pp.local_rank_count(SimpleNamespace(world_size=4, gpus_per_node=8)) == 4 + assert pp.local_rank_count(None) == 1 + monkeypatch.setenv("OMPI_COMM_WORLD_LOCAL_SIZE", "8") + assert pp.local_rank_count(None) == 8 + monkeypatch.setenv("OMPI_COMM_WORLD_LOCAL_SIZE", "many") + monkeypatch.setenv("RAY_LOCAL_WORLD_SIZE", "4") + assert pp.local_rank_count(None) == 4 + + +def test_precompile_worker_entry_compiles_its_shard_first_then_the_rest(monkeypatch) -> None: + """The worker entry point runs the real plan path under precompile_only(): own + shard first, the rest afterwards, busy kernels deferred to a blocking pass.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + + planned, modes = [], [] + mla_decode = _fake_sparse_plan(monkeypatch, planned) + real_plan = mla_decode.BatchSparseMLADecodePagedTSWrapper.plan + busy_once = {"rows": 15, "armed": True} + + def recording_plan(self, device, rows, heads, **kwargs): + modes.append((rows, cc._PRECOMPILE_MODE.get(), cc._SKIP_BUSY.get())) + if rows == busy_once["rows"] and busy_once["armed"] and cc._SKIP_BUSY.get(): + busy_once["armed"] = False + raise cc.KernelBusy("another worker owns it") + return real_plan(self, device, rows, heads, **kwargs) + + monkeypatch.setattr(mla_decode.BatchSparseMLADecodePagedTSWrapper, "plan", recording_plan) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(cc, "expected_cache_meta", lambda: {"probe": "meta"}) + payloads = [ + pp.encode_payload( + 0, rows, 16, dict(max_topk=128, max_extra_topk=512, q_data_type=torch.bfloat16) + ) + for rows in (20, 15, 7, 3) + ] + import json + + assert pp._main(["worker", json.dumps(payloads), "1", "2"]) == 0 + # slot 1 of 2 owns rows 15 and 3; 15 was busy once and is retried blocking at the end + assert [m[0] for m in modes] == [15, 3, 20, 7, 15] + assert all(m[1] == "persist" for m in modes) + assert [m[2] for m in modes] == [True, True, True, True, False] + assert [r for r, _ in planned] == [3, 20, 7, 15] + assert pp.shard_order([0, 1, 2, 3, 4], 1, 2) == [1, 3, 0, 2, 4] + assert pp.shard_order([0, 1], 5, 8) == [0, 1] # a slot past the list owns nothing + + +def test_precompile_workers_get_a_clean_environment_and_bounded_wait(monkeypatch, tmp_path) -> None: + """The runner strips the MPI launcher environment, forwards sys.path and the + worker count 0, assigns one slot per worker, kills workers on timeout, and + reports a worker whose cache key differs from the parent's.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + + spawned = [] + + class _FakeProc: + def __init__(self, argv, env, stdout, returncode=0, text="", hang=False): + self.argv, self.env, self.returncode, self._hang = argv, env, returncode, hang + stdout.write(text) + self._done = False + + def wait(self, timeout=None): + if self._hang and not getattr(self, "killed", False): + raise subprocess.TimeoutExpired(self.argv, timeout) + self._done = True # like Popen.wait() after kill() + return self.returncode + + def poll(self): + return None if (self._hang and not self._done) else self.returncode + + def kill(self): + self._done = True + self.killed = True + + import subprocess + + def make_popen(**behaviour): + def fake_popen(argv, env, stdout, stderr, text): + proc = _FakeProc(argv, env, stdout, **behaviour) + spawned.append(proc) + return proc + + return fake_popen + + monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "2") + monkeypatch.setenv("SLURM_PROCID", "2") + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "4") + monkeypatch.setattr(pp, "_expected_cache_meta", lambda: {"k": 1}) + meta_ok = "[prims-ts precompile] meta={'k': 1}\n" + monkeypatch.setattr(subprocess, "Popen", make_popen(text=meta_ok)) + payloads = [{"rows": r} for r in (20, 15, 7)] + assert pp.run_precompile_workers(payloads, 2) is True + assert len(spawned) == 2 + env = spawned[0].env + assert "OMPI_COMM_WORLD_RANK" not in env and "SLURM_PROCID" not in env + assert env["TLLM_PRIMTS_WARM_WORKERS"] == "0" + import json + import sys + + assert json.loads(env["TLLM_PRIMTS_PRECOMPILE_SYS_PATH"]) == sys.path + assert [p.argv[-2:] for p in spawned] == [["0", "2"], ["1", "2"]] + assert spawned[0].argv[0] == sys.executable and spawned[0].argv[1] == "-c" + assert json.loads(spawned[0].argv[3]) == payloads + + spawned.clear() + monkeypatch.setattr( + subprocess, "Popen", make_popen(text="[prims-ts precompile] meta={'k': 2}\n") + ) + assert pp.run_precompile_workers(payloads, 1 + 1) is False # cache key mismatch is a failure + + spawned.clear() + monkeypatch.setattr(subprocess, "Popen", make_popen(returncode=3, text="boom")) + assert pp.run_precompile_workers(payloads, 2) is False + + spawned.clear() + monkeypatch.setenv("TLLM_PRIMTS_WARM_TIMEOUT", "0.01") + monkeypatch.setattr(subprocess, "Popen", make_popen(hang=True)) + assert pp.run_precompile_workers(payloads, 2) is False + assert all(getattr(p, "killed", False) for p in spawned) # nothing left running + assert pp.run_precompile_workers(payloads, 1) is False # one worker: nothing to fan out + + +def test_precompile_only_persists_without_loading(monkeypatch) -> None: + """Under precompile_only(), build_and_load_attention_kernel compiles and persists + through the unlocked path and hands back a placeholder that refuses to launch; + under probe_cached() a missing kernel raises; with the cache disabled both refuse.""" + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + + persisted, probed = [], [] + monkeypatch.setattr( + cc, + "_precompile_kernel", + lambda module, kernel, fn, files: persisted.append((module, kernel)), + ) + monkeypatch.setattr( + cc, "_probe_kernel", lambda module, kernel, fn, files: probed.append(kernel) + ) + monkeypatch.setattr(cc, "_kernel_source_files", lambda: ("a.py",)) + monkeypatch.setattr(cc, "_compile_cache_key", lambda spec, options, extra: "deadbeef") + monkeypatch.setattr( + "flashinfer.jit.cute_dsl_core.cute_dsl_cache_disabled", lambda: False, raising=False + ) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + compiled = [] + + def compile_fn(): + compiled.append(True) + return "kernel" + + with cc.precompile_only(): + placeholder = cc.build_and_load_attention_kernel( + "mla_decode", "main", object(), "", compile_fn + ) + assert persisted and persisted[0][0] == "prims_ts_mla_decode" + assert persisted[0][1].startswith("main_sm") and persisted[0][1].endswith("_deadbeef") + assert not compiled # the helper owns the compile; nothing was loaded in-process + with pytest.raises(RuntimeError, match="precompile-only"): + placeholder() + assert cc._PRECOMPILE_MODE.get() is None # the context manager restores the mode + with cc.probe_cached(): + cc.build_and_load_attention_kernel("mla_decode", "main", object(), "", compile_fn) + assert probed == [persisted[0][1]] + monkeypatch.setenv("FLASHINFER_DISABLE_JIT", "1") + with cc.precompile_only(), pytest.raises(RuntimeError, match="FLASHINFER_DISABLE_JIT"): + cc.build_and_load_attention_kernel("mla_decode", "main", object(), "", compile_fn) + # Outside the precompile modes FLASHINFER_DISABLE_JIT keeps the disk-cache path: + # the kernel goes through FlashInfer's builder, never a compile_fn() here. + loaded = [] + monkeypatch.setattr( + "flashinfer.jit.cute_dsl_core.build_and_load_cute_dsl_kernel", + lambda module, kernel, fn, extra_key_files=(): loaded.append((module, kernel)) or "loaded", + ) + assert ( + cc.build_and_load_attention_kernel("mla_decode", "main", object(), "", compile_fn) + == "loaded" + ) + assert loaded == [("prims_ts_mla_decode", persisted[0][1])] and not compiled + monkeypatch.delenv("FLASHINFER_DISABLE_JIT") + + +def test_precompile_placeholders_are_never_memoized(monkeypatch) -> None: + """A probe or persist build must not poison the compiled-kernel memo of the + serving process: the next real request compiles / loads the kernel.""" + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + calls = [] + + def fake_compile(spec): + calls.append(cc._PRECOMPILE_MODE.get()) + return cc._unloaded_kernel if cc._PRECOMPILE_MODE.get() else "real kernel" + + monkeypatch.setattr(mla_decode, "_compile_mla_decode", fake_compile) + mla_decode._get_compiled_mla_decode.cache_clear() + spec = ("fake", "compile", "spec") + with cc.probe_cached(): + assert cc.is_placeholder(mla_decode._get_compiled_mla_decode(spec)) + with cc.precompile_only(): + assert cc.is_placeholder(mla_decode._get_compiled_mla_decode(spec)) + assert ( + mla_decode._get_compiled_mla_decode(spec) == "real kernel" + ) # not the memoized placeholder + assert mla_decode._get_compiled_mla_decode(spec) == "real kernel" + assert calls == ["probe", "persist", None] # the real result is memoized + mla_decode._get_compiled_mla_decode.cache_clear() diff --git a/tests/unittest/_torch/attention/test_fmha_registry.py b/tests/unittest/_torch/attention/test_fmha_registry.py index bd6a65809afe..171feed20cb3 100644 --- a/tests/unittest/_torch/attention/test_fmha_registry.py +++ b/tests/unittest/_torch/attention/test_fmha_registry.py @@ -20,6 +20,7 @@ PRIMS_TS = "prims_ts" PRIMS_TS_BLOCK_SPARSE = "prims_ts_block_sparse" +PRIMS_TS_SPARSE_MLA = "prims_ts_sparse_mla" def _canonical_names() -> tuple[str, ...]: @@ -35,12 +36,15 @@ def _enabled_names() -> tuple[str, ...]: def test_prims_ts_precedes_trtllm_gen_in_canonical_order() -> None: names = _canonical_names() assert names.index(PRIMS_TS) < names.index("flashinfer_trtllm_gen") + assert names.index(PRIMS_TS_SPARSE_MLA) < names.index("flashinfer_trtllm_gen") + assert names.index(PRIMS_TS_SPARSE_MLA) < names.index("fallback") def test_default_fmha_libs_exclude_prims_ts(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) assert PRIMS_TS not in registry.DEFAULT_FMHA_LIBS + assert PRIMS_TS_SPARSE_MLA in registry.DEFAULT_FMHA_LIBS assert PRIMS_TS_BLOCK_SPARSE in registry.DEFAULT_FMHA_LIBS assert set(registry.DEFAULT_FMHA_LIBS) <= set(registry.FMHA_LIBS) assert _enabled_names() == registry.DEFAULT_FMHA_LIBS diff --git a/tests/unittest/_torch/attention/test_prims_ts_vendored_imports.py b/tests/unittest/_torch/attention/test_prims_ts_vendored_imports.py new file mode 100644 index 000000000000..78641108b537 --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_vendored_imports.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The vendored PrimTS tree must stay self-contained. + +TRT-LLM pins a flashinfer-python release whose own ``flashinfer.attention`` +tree is older than the vendored one and whose ``flashinfer_api`` has no trace +templates. A vendored module that imports either would silently bind to the +installed copy, so the compatibility patch rewrites those imports; this test +keeps a future re-vendor from reintroducing them. +""" + +import pathlib +import re + +import tensorrt_llm._torch.attention.backends.prims_ts as vendored + +_ESCAPING_IMPORT = re.compile( + r"^\s*(?:from\s+flashinfer\.(?:attention|trace)\b|import\s+flashinfer\.(?:attention|trace)\b)", + re.MULTILINE, +) + + +def test_vendored_prims_ts_does_not_import_flashinfer_attention_or_trace() -> None: + root = pathlib.Path(vendored.__file__).resolve().parent + offenders = sorted( + str(path.relative_to(root)) + for path in root.rglob("*.py") + if _ESCAPING_IMPORT.search(path.read_text(encoding="utf-8")) + ) + assert not offenders, f"vendored modules import the installed flashinfer tree: {offenders}"