From 2c71ac16d04040265ca016ea50b34e92479a9fb9 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:39:48 -0700 Subject: [PATCH 01/21] [None][feat] add FP4 MLA attention backend on Rubin Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 1 + .../attention/backends/fmha/__init__.py | 2 + .../attention/backends/fmha/fallback.py | 3 + .../_torch/attention/backends/fmha/fp4_mla.py | 145 +--- .../attention/backends/fmha/registry.py | 2 + .../attention/backends/fp4_mla/__init__.py | 601 +------------ .../backends/fp4_mla/fp4_mla_context.py | 59 +- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 784 ++++++----------- ...p4_mla_cutedsl_mufu16_fused_v_transpose.py | 820 ++++++------------ .../fp4_mla/fp4_mla_cutedsl_v_repack.py | 119 ++- .../backends/fp4_mla/fp4_mla_kernels.py | 284 +----- .../backends/fp4_mla/fp4_mla_triton.py | 527 ----------- .../_torch/attention/backends/trtllm.py | 450 +++++++++- tensorrt_llm/_torch/attention/mla.py | 17 +- 14 files changed, 1135 insertions(+), 2679 deletions(-) diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index aca96dcd2d47..fdfe6e628710 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -427,6 +427,7 @@ The FMHA package is split by role: MLA uses `query_input` with `is_fused_qkv=False`. - `fmha/combined.py` composes different context and generation implementations for non-MLA mixed batches. +- `fmha/fp4_mla.py` implements FP4 MLA context and no-dequant decode. - `fmha/triton_custom_mask.py` implements the Triton custom-mask context phase. Custom-mask data applies to context requests; for mixed batches, `TrtllmAttention` can pair it with a later causal-generation provider through diff --git a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py index 04d92f69630a..0fc278ce8c1c 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py @@ -18,6 +18,7 @@ from .fallback import FallbackFmha from .flashinfer_sparse_mla import FlashInferSparseMlaFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha +from .fp4_mla import Fp4MlaFmha from .interface import Fmha, FmhaPhase from .msa_decode import MsaDecodeFmha from .msa_prefill import MsaPrefillFmha @@ -34,6 +35,7 @@ "FallbackFmha", "FlashInferSparseMlaFmha", "FlashInferTrtllmGenFmha", + "Fp4MlaFmha", "Fmha", "FmhaCls", "FmhaParams", diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 5631950e9e06..bee07b32a455 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -99,6 +99,9 @@ def _is_supported( return False if q is not None and q.dtype == torch.float8_e4m3fn: return False + attn = self.attn + if attn.is_mla_enable and attn.has_fp4_kv_cache: + return False if forward_args.attention_mask == CustomAttentionMask.CUSTOM: return False if not forward_args.update_kv_cache and not metadata.is_cross: diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index b2703f364708..2a9a3ea9d73a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -14,7 +14,7 @@ # limitations under the License. from dataclasses import replace -from typing import TYPE_CHECKING, Optional, Tuple +from typing import TYPE_CHECKING, Optional import torch @@ -26,7 +26,6 @@ _FP8_CONTEXT_ATTN_ATTR, _FP8_CONTEXT_SCRATCH_ATTR, _build_fp8_mla_context_attn, - _build_fp8_mla_context_metadata, _execute_fp8_context_with_cache_update, _Fp8MlaContextScratch, _get_fp8_mla_context_metadata, @@ -39,7 +38,6 @@ ) from tensorrt_llm.bindings import DataType -from .fallback import FallbackFmha from .phased import FmhaParams, PhasedFmha if TYPE_CHECKING: @@ -53,8 +51,8 @@ class Fp4MlaFmha(PhasedFmha): """TRTLLM FMHA library for FP4 MLA context and no-dequant decode.""" @classmethod - def _is_available(cls, attn: "TrtllmAttention") -> bool: - return attn.uses_fp4_mla_attention + def is_available(cls, attn: "TrtllmAttention") -> bool: + return attn.is_mla_enable and attn.has_fp4_kv_cache def forward( self, @@ -83,17 +81,20 @@ def _validate_request( if forward_args.attention_sinks is not None: raise NotImplementedError("FP4 MLA does not support attention sinks.") - sparse_runtime_params = forward_args.sparse_runtime_params + sparse_prediction = forward_args.sparse_prediction + sparse_params = self.attn.sparse_params + uses_spcompress = getattr(sparse_params, "uses_spcompress", None) if ( ( - sparse_runtime_params.sparse_kv_indices is not None - and sparse_runtime_params.sparse_kv_indices.numel() > 0 + sparse_prediction.sparse_kv_indices is not None + and sparse_prediction.sparse_kv_indices.numel() > 0 ) or ( - sparse_runtime_params.sparse_attn_indices is not None - and sparse_runtime_params.sparse_attn_indices.numel() > 0 + sparse_prediction.sparse_attn_indices is not None + and sparse_prediction.sparse_attn_indices.numel() > 0 ) or metadata.num_sparse_topk > 0 + or uses_spcompress ): raise NotImplementedError("FP4 MLA does not support sparse attention.") @@ -127,8 +128,8 @@ def run_mla_context(self, params: FmhaParams) -> None: metadata = params.meta forward_args = params.fwd q = params.qkv_input - k = params.k_input - v = params.v_input + k = params.key_input + v = params.value_input output = params.context_buf if q is None or k is None or v is None: raise RuntimeError("FP4 MLA context requires expanded Q, K, and V tensors.") @@ -155,7 +156,7 @@ def run_mla_context(self, params: FmhaParams) -> None: num_tokens = q.shape[0] output = output.view(num_tokens, -1) - local_layer = attn.get_fp4_mla_local_layer_idx(metadata) + local_layer = attn.get_local_layer_idx(metadata) kv_lora_rank = attn.kv_lora_rank or 0 qk_rope_head_dim = attn.qk_rope_head_dim or 0 @@ -196,7 +197,6 @@ def update_fp4_cache() -> None: fp8_attention = getattr(attn, _FP8_CONTEXT_ATTN_ATTR, None) if fp8_attention is None: fp8_attention = _build_fp8_mla_context_attn(attn) - fp8_attention.fmha_libs = [FallbackFmha(fp8_attention)] setattr(attn, _FP8_CONTEXT_ATTN_ATTR, fp8_attention) fp8_attention.rotary_inv_freq = attn.rotary_inv_freq fp8_attention.rotary_cos_sin = attn.rotary_cos_sin @@ -220,121 +220,6 @@ def run_fp8_context() -> None: scratch.cache_done_event, ) - def forward_context_partition( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - metadata: "TrtllmAttentionMetadata", - forward_args: AttentionForwardArgs, - *, - kv_lens_cuda: torch.Tensor, - kv_lens_cpu: torch.Tensor, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """Run one explicit-KV partition through the FP8 MLA context kernel. - - Returns whatever ``TrtllmAttention.forward`` returns for the FP8 - context copy: the attention output plus its optional output scale. - """ - if forward_args.output is None: - raise RuntimeError("FP4 MLA context partition requires an output buffer.") - if forward_args.latent_cache is not None: - raise RuntimeError("FP4 MLA context partition expects explicit K/V tensors.") - if forward_args.output_sf is not None: - raise NotImplementedError( - "FP4 MLA context partition does not support quantized attention output." - ) - if forward_args.attention_mask not in ( - PredefinedAttentionMask.CAUSAL, - PredefinedAttentionMask.FULL, - ): - raise NotImplementedError( - "FP4 MLA context partition requires a causal or full attention mask." - ) - if forward_args.attention_mask_data is not None: - raise NotImplementedError( - "FP4 MLA context partition does not support custom attention masks." - ) - if forward_args.attention_sinks is not None: - raise NotImplementedError("FP4 MLA context partition does not support attention sinks.") - if metadata.is_cuda_graph: - raise NotImplementedError("FP4 MLA chunked prefill does not support CUDA graphs.") - if metadata.num_contexts <= 0 or q.shape[0] != metadata.num_ctx_tokens: - raise RuntimeError( - "FP4 MLA context partition query token count must match context metadata." - ) - if k.shape[0] != v.shape[0]: - raise RuntimeError("FP4 MLA context partition K/V token counts do not match.") - - sparse_runtime_params = forward_args.sparse_runtime_params - if ( - ( - sparse_runtime_params.sparse_kv_indices is not None - and sparse_runtime_params.sparse_kv_indices.numel() > 0 - ) - or ( - sparse_runtime_params.sparse_attn_indices is not None - and sparse_runtime_params.sparse_attn_indices.numel() > 0 - ) - or metadata.num_sparse_topk > 0 - ): - raise NotImplementedError("FP4 MLA chunked prefill does not support sparse attention.") - - require_fp4_mla_fp8_context_support() - kv_cache_manager = metadata.kv_cache_manager - if kv_cache_manager is None: - raise RuntimeError("FP8 MLA context scratch requires a KV cache manager.") - kv_lora_rank = self.attn.kv_lora_rank or 0 - qk_rope_head_dim = self.attn.qk_rope_head_dim or 0 - scratch_head_dim = kv_lora_rank + qk_rope_head_dim - scratch = getattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, None) - if not isinstance(scratch, _Fp8MlaContextScratch) or not scratch.matches( - metadata, - device=q.device, - head_dim=scratch_head_dim, - ): - scratch = _Fp8MlaContextScratch.create( - metadata, - device=q.device, - head_dim=scratch_head_dim, - ) - setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch) - - attn = self.attn - fp8_attention = getattr(attn, _FP8_CONTEXT_ATTN_ATTR, None) - if fp8_attention is None: - fp8_attention = _build_fp8_mla_context_attn(attn) - fp8_attention.fmha_libs = [FallbackFmha(fp8_attention)] - setattr(attn, _FP8_CONTEXT_ATTN_ATTR, fp8_attention) - fp8_attention.rotary_inv_freq = attn.rotary_inv_freq - fp8_attention.rotary_cos_sin = attn.rotary_cos_sin - - expected_kv_tokens = int(kv_lens_cpu[: metadata.num_contexts].sum().item()) - if k.shape[0] != expected_kv_tokens: - raise RuntimeError( - "FP4 MLA context partition K/V token count does not match " - f"the KV lengths: got {k.shape[0]}, expected {expected_kv_tokens}." - ) - scratch.prepare( - metadata, - context_lengths_cuda=kv_lens_cuda, - context_lengths_cpu=kv_lens_cpu, - ) - fp8_metadata = _build_fp8_mla_context_metadata( - metadata, - scratch, - kv_lens_cuda=kv_lens_cuda[: metadata.num_contexts], - kv_lens_cpu=kv_lens_cpu[: metadata.num_contexts], - ) - fp8_forward_args = replace( - forward_args, - output_sf=None, - kv_scale_orig_quant=None, - kv_scale_quant_orig=None, - latent_cache=None, - ) - return fp8_attention.forward(q, k, v, fp8_metadata, fp8_forward_args) - def run_mla_generation(self, params: FmhaParams) -> None: attn = params.attn metadata = params.meta @@ -350,7 +235,7 @@ def run_mla_generation(self, params: FmhaParams) -> None: if metadata.num_generations <= 0: raise RuntimeError("FP4 MLA generation requires generation requests.") - local_layer = attn.get_fp4_mla_local_layer_idx(metadata) + local_layer = attn.get_local_layer_idx(metadata) kv_lora_rank = attn.kv_lora_rank or 0 qk_rope_head_dim = attn.qk_rope_head_dim or 0 fused_head_dim = kv_lora_rank + qk_rope_head_dim diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index 8c4fa5bdf016..b7b60a0de6d1 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -19,6 +19,7 @@ from .cute_dsl_mla import CuteDslMlaFmha from .fallback import FallbackFmha from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha +from .fp4_mla import Fp4MlaFmha from .interface import Fmha from .prims_ts import PrimsTSFmha from .triton_custom_mask import TritonCustomMaskFmha @@ -39,6 +40,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: from .prims_ts_block_sparse import PrimsTSBlockSparseFmha return { + "fp4_mla": Fp4MlaFmha, "triton_custom_mask": TritonCustomMaskFmha, "cute_dsl_mla": CuteDslMlaFmha, # A pair, not a priority: msa_decode serves a MiniMax-M3 MSA layer's diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py index 7a310f0e9593..c60cd9512ff7 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -18,14 +18,11 @@ import triton import triton.language as tl -from tensorrt_llm._utils import get_sm_version, prefer_pinned -from tensorrt_llm.bindings import DataType +from tensorrt_llm._utils import get_sm_version from .fp4_mla_kernels import ( - _fp4_mla_chunked_cache_gather_kernel, _fp4_mla_context_cache_update_kernel, _fp4_mla_generation_fused_qk_rope_cache_update_kernel, - _fp4_mla_rebuild_v_scale_from_k_scale_kernel, ) HP_BLOCK_SIZE: int = 16 @@ -49,8 +46,6 @@ FP4_MLA_Q_PACKED_DIM: int = FP4_MLA_Q_LOGICAL_DIM // 2 FP4_MLA_Q_SF_GROUPS: int = FP4_MLA_Q_LOGICAL_DIM // FP4_BLOCK_SIZE FP4_MLA_ATTENTION_BACKEND_ENV = "TRTLLM_FP4_MLA_ATTENTION_BACKEND" -# Retained for compatibility; mufu16 is now the only non-fused CuTeDSL variant. -FP4_MLA_CUTEDSL_MUFU16_ENV = "TRTLLM_FP4_MLA_CUTEDSL_MUFU16" FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV = "TRTLLM_FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE" _FP4_MLA_CUTEDSL_BACKEND = "cutedsl" _FP4_MLA_K_RESIDUAL_BACKENDS = ("triton", _FP4_MLA_CUTEDSL_BACKEND) @@ -463,7 +458,7 @@ def populate_fp4_mla_generation_lengths( def _fp4_mla_page_table_spec(kv_cache_manager: Any) -> Any: get_spec = getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None) if not callable(get_spec): - raise RuntimeError("FP4 MLA requires V2 cache-layout page metadata.") + raise RuntimeError("FP4 MLA requires Fp4MlaKVCacheManagerV2 page metadata.") spec = get_spec() for field_name in ( "cache_pool_id", @@ -573,7 +568,7 @@ def configure_fp4_mla_device_page_table( ) -> bool: """Configure the fixed-stride, device-materialized page table. - Eager context and generation batches receive the full block-offset + Context, generation, and fresh mixed batches receive the full block-offset table on the GPU. The materialization kernel decodes V2 page indices and refreshes rows from the final device KV lengths before cache update. """ @@ -600,7 +595,18 @@ def configure_fp4_mla_device_page_table( tensors = (block_offsets, page_ids, paged_kv_indptr, paged_kv_indptr_decode) is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) generation_only = num_contexts == 0 - eager_context = not is_cuda_graph and num_contexts > 0 + fresh_mixed = ( + not is_cuda_graph + and num_contexts > 0 + and num_generation_sequences > 0 + and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 + ) + fresh_context_only = ( + not is_cuda_graph + and num_contexts > 0 + and num_generation_sequences == 0 + and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 + ) has_valid_generation = num_generation_sequences == 0 or ( num_generation_tokens >= num_generation_sequences and num_generation_tokens % num_generation_sequences == 0 @@ -608,12 +614,13 @@ def configure_fp4_mla_device_page_table( # NVFP4 exposes one data pool plus its paired block-scale pool. The # materializer reads encoded data offsets from pool 0. supported = ( - (generation_only or eager_context) + (generation_only or fresh_mixed or fresh_context_only) and kv_cache_manager is not None and has_valid_generation and int(getattr(metadata, "beam_width", 1)) == 1 and not bool(getattr(metadata, "is_spec_dec_tree", False)) and not bool(getattr(metadata, "locality_domain_enabled", False)) + and not bool(getattr(metadata, "enable_helix", False)) and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK and max_page_capacity > 0 and page_index_scale > 0 @@ -645,7 +652,7 @@ def configure_fp4_mla_device_page_table( and kv_lens.ndim == 1 and kv_lens.numel() >= num_sequences ) - if eager_context and num_generation_sequences > 0 and not host_kv_lens_available: + if fresh_mixed and not host_kv_lens_available: return False if not is_cuda_graph and host_kv_lens_available: generation_tokens_per_sequence = ( @@ -1146,251 +1153,6 @@ def get_fp4_mla_v_scale_pool_view( return torch.as_strided(pool, size=shape, stride=strides) -def _rebuild_fp4_mla_v_scales_from_k_scales( - sf_cache: torch.Tensor, - v_scale_pool: torch.Tensor, - page_ids: torch.Tensor, - page_valid_tokens: torch.Tensor, - *, - local_layer: int, - v_head_dim: int, - page_size: int, -) -> None: - """Bit-exactly rebuild imported MLA V scales from transferred K scales.""" - if page_ids.numel() == 0: - return - if page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - "FP4 MLA imported V-scale rebuild requires " - f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." - ) - for name, tensor in ( - ("sf_cache", sf_cache), - ("v_scale_pool", v_scale_pool), - ("page_ids", page_ids), - ("page_valid_tokens", page_valid_tokens), - ): - if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: - raise ValueError(f"{name} must be a CUDA tensor.") - if page_ids.dtype != torch.int32 or page_valid_tokens.dtype != torch.int32: - raise TypeError("FP4 MLA imported page IDs and valid-token counts must use int32.") - if page_ids.ndim != 1 or page_valid_tokens.ndim != 1: - raise ValueError("FP4 MLA imported page metadata must be one-dimensional.") - if page_ids.numel() != page_valid_tokens.numel(): - raise ValueError("FP4 MLA imported page IDs and valid-token counts must have equal length.") - if not page_ids.is_contiguous() or not page_valid_tokens.is_contiguous(): - raise ValueError("FP4 MLA imported page metadata must be contiguous.") - if not (sf_cache.device == v_scale_pool.device == page_ids.device == page_valid_tokens.device): - raise ValueError("FP4 MLA imported cache tensors must be on the same device.") - - k_sf_bytes = sf_cache.view(torch.uint8) - v_sf_bytes = v_scale_pool.view(torch.uint8) - if k_sf_bytes.ndim < 2 or v_sf_bytes.ndim < 3: - raise ValueError( - "FP4 MLA imported scale pools require per-page K storage and " - "per-layer/per-page V storage." - ) - num_layers = int(v_sf_bytes.shape[0]) - num_pages = int(v_sf_bytes.shape[1]) - if not 0 <= local_layer < num_layers: - raise IndexError( - f"local_layer={local_layer} is outside the V-scale pool with {num_layers} layers." - ) - if int(k_sf_bytes.shape[0]) != num_pages: - raise ValueError( - "FP4 MLA K/V scale pools disagree on their physical page count: " - f"{int(k_sf_bytes.shape[0])} != {num_pages}." - ) - sf_per_token = int(k_sf_bytes.shape[-1]) - required_sf_per_token = _ceil_div(v_head_dim, FP4_BLOCK_SIZE) - if sf_per_token < required_sf_per_token: - raise ValueError( - "FP4 MLA K-scale storage is too narrow for the compressed V head: " - f"{sf_per_token} < {required_sf_per_token}." - ) - required_v_page_elems = get_fp4_mla_v_scale_pool_size(v_head_dim, page_size) - if int(v_sf_bytes.shape[-1]) < required_v_page_elems: - raise ValueError( - "FP4 MLA V-scale page storage is too small for import rebuild: " - f"{int(v_sf_bytes.shape[-1])} < {required_v_page_elems}." - ) - - token_groups = page_size // HP_BLOCK_SIZE - _fp4_mla_rebuild_v_scale_from_k_scale_kernel[ - (page_ids.numel(), triton.cdiv(v_head_dim, FP4_BLOCK_SIZE)) - ]( - k_sf_bytes, - v_sf_bytes, - page_ids, - page_valid_tokens, - page_ids.numel(), - num_pages, - num_layers, - local_layer, - page_size, - k_sf_bytes.stride(0), - v_sf_bytes.stride(0), - v_sf_bytes.stride(1), - V_HEAD_D=v_head_dim, - HP_BLOCK=HP_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=token_groups, - BLOCK_TOKEN_GROUPS=triton.next_power_of_2(token_groups), - num_warps=4, - ) - - -def _stage_fp4_mla_import_page_metadata( - prompt_block_ids: list[int], - *, - prompt_len: int, - page_size: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - """Stage imported page metadata without synchronizing the CUDA stream.""" - if device.type != "cuda": - raise ValueError("FP4 MLA disaggregated import requires a CUDA device.") - num_prompt_pages = len(prompt_block_ids) - pin_memory = prefer_pinned() - page_ids_host = torch.tensor( - prompt_block_ids, - dtype=torch.int32, - device="cpu", - pin_memory=pin_memory, - ) - page_valid_tokens_host = torch.full( - (num_prompt_pages,), - page_size, - dtype=torch.int32, - device="cpu", - pin_memory=pin_memory, - ) - page_valid_tokens_host[-1] = prompt_len - (num_prompt_pages - 1) * page_size - # Constructing a CUDA tensor directly from a Python list synchronizes the - # current stream. During overlap scheduling that stream already waits for - # the previous forward, exposing the import rebuild as an inter-step gap. - # Explicit non-blocking copies keep the CPU free to enqueue every rebuild - # and the next forward while preserving their existing stream order. - page_ids = torch.empty_like(page_ids_host, device=device) - page_valid_tokens = torch.empty_like(page_valid_tokens_host, device=device) - page_ids.copy_(page_ids_host, non_blocking=True) - page_valid_tokens.copy_(page_valid_tokens_host, non_blocking=True) - return page_ids, page_valid_tokens - - -def rebuild_fp4_mla_disagg_imported_cache( - kv_cache_manager: Any, - request_id: int, - prompt_len: int, -) -> bool: - """Rebuild GEN-local FP4 MLA sidecars after a disaggregated KV import. - - The disaggregated payload carries the native V2 K, K-scale, and BF16 HP - roles. V scales and CuTeDSL's V-packed layout are deterministic - process-local views, so rebuilding them here avoids transfer bandwidth and - guarantees they are ready before a first decode step that may execute - through a pre-captured CUDA graph. - """ - if ( - kv_cache_manager is None - or getattr(kv_cache_manager, "dtype", None) != DataType.NVFP4 - or getattr(kv_cache_manager, "kv_factor", None) != 1 - or getattr(kv_cache_manager, "mla_v_scale_head_dim", None) is None - or not callable(getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None)) - ): - return False - if not isinstance(prompt_len, int) or prompt_len < 0: - raise ValueError( - f"FP4 MLA disaggregated import needs a nonnegative prompt_len, got {prompt_len}." - ) - # Helix assigns whole pages round-robin, so a rank may own no prompt pages. - # There are no process-local V sidecars to rebuild on that rank. - if prompt_len == 0: - return True - - page_size = int(kv_cache_manager.tokens_per_block) - if page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - "FP4 MLA disaggregated import requires " - f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." - ) - pp_layers = list(getattr(kv_cache_manager, "pp_layers", ())) - num_local_layers = int(getattr(kv_cache_manager, "num_local_layers", len(pp_layers))) - if len(pp_layers) != num_local_layers: - raise RuntimeError( - "FP4 MLA disaggregated import cannot map local to global layers: " - f"{len(pp_layers)} PP layers for {num_local_layers} local layers." - ) - fp4_local_layers = list( - getattr(kv_cache_manager, "_fp4_mla_compact_to_local", range(num_local_layers)) - ) - if not fp4_local_layers: - raise RuntimeError("FP4 MLA disaggregated import found no local MLA layers.") - if any(local_layer < 0 or local_layer >= num_local_layers for local_layer in fp4_local_layers): - raise RuntimeError( - "FP4 MLA disaggregated import has invalid compact-to-local layer mapping: " - f"{fp4_local_layers}." - ) - - num_prompt_pages = _ceil_div(prompt_len, page_size) - first_attention_layer = pp_layers[fp4_local_layers[0]] - block_ids_per_seq = kv_cache_manager.get_batch_cache_indices( - [int(request_id)], layer_idx=first_attention_layer - ) - if len(block_ids_per_seq) != 1 or len(block_ids_per_seq[0]) < num_prompt_pages: - available = len(block_ids_per_seq[0]) if block_ids_per_seq else 0 - raise RuntimeError( - "FP4 MLA disaggregated import is missing prompt pages for request " - f"{request_id}: need {num_prompt_pages}, have {available}." - ) - prompt_block_ids = [int(block_id) for block_id in block_ids_per_seq[0][:num_prompt_pages]] - - v_scale_pool = kv_cache_manager.get_mla_v_scale_pool() - if not isinstance(v_scale_pool, torch.Tensor): - raise RuntimeError("FP4 MLA disaggregated import requires the manager V-scale pool.") - page_ids, page_valid_tokens = _stage_fp4_mla_import_page_metadata( - prompt_block_ids, - prompt_len=prompt_len, - page_size=page_size, - device=v_scale_pool.device, - ) - - v_scale_head_dim = int(kv_cache_manager.mla_v_scale_head_dim) - cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND - for compact_layer, local_layer in enumerate(fp4_local_layers): - layer_idx = pp_layers[local_layer] - kv_cache, sf_cache = kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) - _rebuild_fp4_mla_v_scales_from_k_scales( - sf_cache, - v_scale_pool, - page_ids, - page_valid_tokens, - local_layer=compact_layer, - v_head_dim=v_scale_head_dim, - page_size=page_size, - ) - if cutedsl_backend and not _fp4_mla_cutedsl_fused_v_transpose_enabled(): - v_head_dim = getattr(kv_cache_manager, "mla_v_head_dim", None) - if v_head_dim is None: - raise RuntimeError( - "CuTeDSL FP4 MLA disaggregated import requires a persistent V head dimension." - ) - v_packed = kv_cache_manager.get_mla_v_packed_pool(compact_layer) - if not isinstance(v_packed, torch.Tensor): - raise RuntimeError( - "CuTeDSL FP4 MLA disaggregated import requires the persistent V-packed pool." - ) - _repack_cutedsl_v_packed_cache( - v_packed, - kv_cache, - page_ids, - v_head_dim=int(v_head_dim), - page_size=page_size, - block_v=FP4_MLA_SCALE_ROW_GROUP, - ) - return True - - # Python launch helpers @@ -1494,8 +1256,6 @@ def _scatter_fp4_mla_kv_cache_2d_context( v_sf: torch.Tensor, global_scale: torch.Tensor, rotary_cos_sin: Optional[torch.Tensor], - q_context: Optional[torch.Tensor], - q_nope_head_dim: Optional[int], *, token_offset: int, local_layer: int, @@ -1527,34 +1287,6 @@ def _scatter_fp4_mla_kv_cache_2d_context( ) rotary_cos_sin_ptr = rotary_cos_sin if rotary_cos_sin is not None else latent_cache - apply_q_rope = q_context is not None - block_q_heads = 16 - if apply_q_rope: - if rotary_cos_sin is None or q_nope_head_dim is None: - raise ValueError("FP4 MLA fused context Q-RoPE requires a rotary table and Q layout.") - q_head_dim = q_nope_head_dim + rope_dim - if ( - q_context.dtype != torch.bfloat16 - or q_context.device != latent_cache.device - or q_context.ndim != 2 - or q_context.shape[0] != num_tokens - or q_context.shape[1] <= 0 - or q_nope_head_dim <= 0 - or q_context.shape[1] % q_head_dim != 0 - or not q_context.is_contiguous() - ): - raise ValueError( - "FP4 MLA fused context Q-RoPE requires a contiguous same-device BF16 " - f"tensor shaped [tokens, heads * ({q_nope_head_dim} + {rope_dim})]." - ) - num_q_heads = q_context.shape[1] // q_head_dim - q_context_view = q_context.view(num_tokens, num_q_heads, q_head_dim) - q_head_blocks = triton.cdiv(num_q_heads, block_q_heads) - else: - num_q_heads = 0 - q_context_view = latent_cache - q_head_blocks = 0 - hp_pool = getattr(metadata, "high_precision_kv_pool", None) if not isinstance(hp_pool, torch.Tensor): raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") @@ -1590,7 +1322,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( _fp4_mla_context_cache_update_kernel[ ( num_tokens, - num_dim_blocks + q_head_blocks, + num_dim_blocks, ) ]( kv_cache, @@ -1598,7 +1330,6 @@ def _scatter_fp4_mla_kv_cache_2d_context( v_sf, v_packed_output, latent_cache, - q_context_view, global_scale, rotary_cos_sin_ptr, hp_pool, @@ -1625,9 +1356,6 @@ def _scatter_fp4_mla_kv_cache_2d_context( sf_cache.stride(0), latent_cache.stride(0), latent_cache.stride(1), - q_context_view.stride(0), - q_context_view.stride(1) if apply_q_rope else 0, - q_context_view.stride(2) if apply_q_rope else 0, v_sf.stride(0), v_sf.stride(1), v_packed_s0, @@ -1645,11 +1373,6 @@ def _scatter_fp4_mla_kv_cache_2d_context( STORE_K_RESIDUAL=(_fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS), ROPE_DIM=rope_dim, APPLY_K_ROPE=apply_k_rope, - APPLY_Q_ROPE=apply_q_rope, - NUM_DIM_BLOCKS=num_dim_blocks, - NUM_Q_HEADS=num_q_heads, - Q_NOPE_DIM=q_nope_head_dim if q_nope_head_dim is not None else 0, - BLOCK_Q_HEADS=block_q_heads, POOL_HEAD_D=pool_head_dim, STORE_HP_TAIL=store_hp_tail, WRITE_V_PACKED=write_v_packed, @@ -1781,8 +1504,6 @@ def _scatter_fp4_mla_kv_cache_2d_generation( q_sf_out: torch.Tensor, v_packed_base: Optional[torch.Tensor], v_page_offset: int, - helix_position_offsets: Optional[torch.Tensor], - helix_is_inactive_rank: Optional[torch.Tensor], ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: num_contexts = metadata.num_contexts num_seqs = metadata.num_seqs @@ -1825,59 +1546,6 @@ def _scatter_fp4_mla_kv_cache_2d_generation( num_hp_pages = pool.shape[0] max_gen_len = num_tokens // num_gen - use_helix = helix_position_offsets is not None or helix_is_inactive_rank is not None - use_helix_local_slots = use_helix and bool(getattr(metadata, "_helix_spec_tokens_valid", False)) - if use_helix: - if helix_position_offsets is None or helix_is_inactive_rank is None: - raise RuntimeError( - "FP4 MLA Helix requires both position-offset and inactive-rank metadata." - ) - if max_gen_len != 1 and not use_helix_local_slots: - raise NotImplementedError( - "FP4 MLA multi-token Helix requires speculative per-token metadata." - ) - if ( - helix_position_offsets.dtype != torch.int32 - or helix_position_offsets.device != latent_cache.device - or helix_position_offsets.ndim != 1 - or helix_position_offsets.numel() < num_tokens - or not helix_position_offsets.is_contiguous() - ): - raise ValueError( - "FP4 MLA Helix position offsets must be a contiguous same-device " - "int32 tensor covering every generation token." - ) - if ( - helix_is_inactive_rank.dtype != torch.bool - or helix_is_inactive_rank.device != latent_cache.device - or helix_is_inactive_rank.ndim != 1 - or helix_is_inactive_rank.numel() < num_gen - or not helix_is_inactive_rank.is_contiguous() - ): - raise ValueError( - "FP4 MLA Helix inactive-rank metadata must be a contiguous " - "same-device bool tensor covering every generation sequence." - ) - if use_helix_local_slots: - helix_local_slots = getattr(metadata, "helix_local_slots", None) - if ( - not isinstance(helix_local_slots, torch.Tensor) - or helix_local_slots.dtype != torch.int32 - or helix_local_slots.device != latent_cache.device - or helix_local_slots.ndim != 1 - or helix_local_slots.numel() < num_tokens - or not helix_local_slots.is_contiguous() - ): - raise ValueError( - "FP4 MLA speculative Helix local slots must be a contiguous " - "same-device int32 tensor covering every generation token." - ) - else: - helix_local_slots = helix_position_offsets - else: - helix_position_offsets = kv_lens_gen - helix_local_slots = kv_lens_gen - helix_is_inactive_rank = gen_lens_gen _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) max_rewind_len = hp_pool_size - HP_BLOCK_SIZE page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) @@ -1991,9 +1659,6 @@ def launch_generation_update( q_sf_output, kv_lens_gen, gen_lens_gen, - helix_position_offsets, - helix_local_slots, - helix_is_inactive_rank, page_ids, hp_page_ids, metadata.paged_kv_indptr_decode, @@ -2032,8 +1697,6 @@ def launch_generation_update( K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, STORE_K_RESIDUAL=store_k_residual, FUSE_ROPE_CACHE_STORE=True, - USE_HELIX=use_helix, - USE_HELIX_LOCAL_SLOTS=use_helix_local_slots, WRITE_V_PACKED=write_v_packed, MAX_GEN_TILES=max_gen_tiles_variant, ROPE_DIM=rope_dim, @@ -2088,8 +1751,6 @@ def launch_generation_update( metadata.page_size, write_v_packed, store_k_residual, - use_helix, - use_helix_local_slots, tuple(q1_variants), multi_token_tiles, tuple(kv_cache.stride()), @@ -2125,7 +1786,6 @@ def launch_generation_update( v_sf, v_packed_output, latent_cache, - latent_cache, global_scale, rotary_table, pool, @@ -2152,9 +1812,6 @@ def launch_generation_update( sf_cache.stride(0), latent_cache.stride(0), latent_cache.stride(1), - latent_cache.stride(0), - 0, - 0, v_sf.stride(0), v_sf.stride(1), v_packed_s0, @@ -2172,11 +1829,6 @@ def launch_generation_update( STORE_K_RESIDUAL=store_k_residual, ROPE_DIM=rope_dim, APPLY_K_ROPE=True, - APPLY_Q_ROPE=False, - NUM_DIM_BLOCKS=num_dim_blocks, - NUM_Q_HEADS=0, - Q_NOPE_DIM=0, - BLOCK_Q_HEADS=16, POOL_HEAD_D=hp_head_dim, STORE_HP_TAIL=True, WRITE_V_PACKED=write_v_packed, @@ -2224,140 +1876,6 @@ def launch_generation_update( # Public cache update and decode entry points -def load_fp4_mla_chunked_kv_cache( - metadata: Any, - layer_idx: int, - *, - num_ctx_cached_tokens: int, - cu_chunked_seq_len: torch.Tensor, - chunked_global_offset: torch.Tensor, - chunked_max_seq_len: int, - out_dtype: torch.dtype, - kv_lora_rank: int, - qk_rope_head_dim: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Gather one cached-prefix partition from dense FP4 MLA V2 storage.""" - if out_dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError(f"FP4 MLA chunk gather does not support output dtype {out_dtype}.") - if num_ctx_cached_tokens < 0: - raise ValueError("FP4 MLA chunk gather token count must be non-negative.") - if chunked_max_seq_len < 0: - raise ValueError("FP4 MLA chunk gather max sequence length must be non-negative.") - if num_ctx_cached_tokens > 0 and chunked_max_seq_len == 0: - raise ValueError("A non-empty FP4 MLA chunk gather requires a positive max length.") - if ( - kv_lora_rank <= 0 - or kv_lora_rank % FP4_BLOCK_SIZE != 0 - or qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM - ): - raise ValueError( - "FP4 MLA chunk gather requires a positive kv_lora_rank and " - f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got " - f"{kv_lora_rank} and {qk_rope_head_dim}." - ) - - num_contexts = int(metadata.num_contexts) - if num_ctx_cached_tokens > 0 and num_contexts <= 0: - raise ValueError("A non-empty FP4 MLA chunk gather requires context requests.") - tensors = (cu_chunked_seq_len, chunked_global_offset) - if any(not isinstance(tensor, torch.Tensor) or not tensor.is_cuda for tensor in tensors): - raise ValueError("FP4 MLA chunk gather metadata must use CUDA tensors.") - if ( - cu_chunked_seq_len.dtype != torch.int64 - or cu_chunked_seq_len.ndim != 1 - or cu_chunked_seq_len.numel() < num_contexts + 1 - or not cu_chunked_seq_len.is_contiguous() - ): - raise ValueError( - "FP4 MLA chunk gather requires a contiguous int64 cumulative-length tensor." - ) - if ( - chunked_global_offset.dtype != torch.int64 - or chunked_global_offset.ndim != 1 - or chunked_global_offset.numel() < num_contexts - or not chunked_global_offset.is_contiguous() - ): - raise ValueError("FP4 MLA chunk gather requires contiguous int64 global offsets.") - if cu_chunked_seq_len.device != chunked_global_offset.device: - raise ValueError("FP4 MLA chunk gather metadata tensors must share one device.") - - compressed_kv = torch.empty( - (num_ctx_cached_tokens, kv_lora_rank), - dtype=out_dtype, - device=cu_chunked_seq_len.device, - ) - k_pe = torch.empty( - (num_ctx_cached_tokens, qk_rope_head_dim), - dtype=out_dtype, - device=cu_chunked_seq_len.device, - ) - if num_ctx_cached_tokens == 0: - return compressed_kv, k_pe - - if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): - raise RuntimeError("FP4 MLA chunk gather requires fixed-stride device page metadata.") - _materialize_fp4_mla_device_page_table_for_forward(metadata) - page_table_stride = int(metadata.fp4_mla_page_table_stride) - page_ids = metadata._paged_kv_indices - if ( - page_table_stride <= 0 - or not isinstance(page_ids, torch.Tensor) - or page_ids.dtype != torch.int32 - or not page_ids.is_cuda - or page_ids.device != cu_chunked_seq_len.device - or page_ids.numel() < num_contexts * page_table_stride - ): - raise RuntimeError("FP4 MLA chunk gather received invalid device page metadata.") - - kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) - head_dim = kv_lora_rank + qk_rope_head_dim - storage_head_dim = _validate_fp4_mla_kv_storage_shape( - kv_cache, - sf_cache, - head_dim=head_dim, - backend=_fp4_mla_attention_backend(), - ) - if storage_head_dim != head_dim + FP4_MLA_K_RESIDUAL_DIM: - raise RuntimeError( - "FP4 MLA chunk gather requires K residual storage for BF16 reconstruction." - ) - sf_cache = sf_cache.view(torch.float8_e4m3fn) - global_scale = _get_fp4_mla_global_scale(metadata, kv_cache.device) - token_block = 16 - grid = ( - triton.cdiv(chunked_max_seq_len, token_block), - num_contexts, - triton.cdiv(head_dim, FP4_BLOCK_SIZE), - ) - _fp4_mla_chunked_cache_gather_kernel[grid]( - compressed_kv, - k_pe, - kv_cache, - sf_cache, - page_ids, - cu_chunked_seq_len, - chunked_global_offset, - global_scale, - chunked_max_seq_len, - page_table_stride, - kv_cache.shape[0], - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - compressed_kv.stride(0), - k_pe.stride(0), - KV_LORA_RANK=kv_lora_rank, - QK_ROPE_HEAD_DIM=qk_rope_head_dim, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=storage_head_dim // FP4_BLOCK_SIZE, - TOKEN_BLOCK=token_block, - num_warps=4, - ) - return compressed_kv, k_pe - - def scatter_fp4_mla_kv_cache( metadata: Any, latent_cache: torch.Tensor, @@ -2371,10 +1889,6 @@ def scatter_fp4_mla_kv_cache( q_pe: Optional[torch.Tensor] = None, q_rope_out: Optional[torch.Tensor] = None, q_quant_input: Optional[torch.Tensor] = None, - helix_position_offsets: Optional[torch.Tensor] = None, - helix_is_inactive_rank: Optional[torch.Tensor] = None, - q_context: Optional[torch.Tensor] = None, - q_nope_head_dim: Optional[int] = None, ) -> bool: """Quantize MLA latent tokens and scatter them into the paged FP4 cache. @@ -2394,10 +1908,8 @@ def scatter_fp4_mla_kv_cache( layouts. Tail K-only dimensions use K's per-token 1D scales. For exclusively owned CuTeDSL pages, context scatter also writes the persistent packed-V sidecar. - Context scatter can rotate Q in place and rotate the K tail directly from - the unassembled latent tensor. When context Q is supplied for chunked - prefill, it also writes rotated K back for current-chunk attention. - Generation scatter rewrites each touched 16-token tile by reading + Context scatter can rotate the K tail directly from the unassembled latent + tensor. Generation scatter rewrites each touched 16-token tile by reading old tokens from the HP pool and new tokens from ``latent_cache``. The static-scale generation specialization can also rotate Q and new K tails while updating the HP pool. @@ -2412,11 +1924,6 @@ def scatter_fp4_mla_kv_cache( metadata._fp4_mla_q_batch_capacity = None if latent_cache.numel() == 0: raise ValueError("FP4 MLA cache scatter requires at least one latent token.") - if q_context is not None and not latent_cache.is_contiguous(): - raise ValueError( - "FP4 MLA fused context Q/K RoPE requires contiguous latent_cache " - "storage for in-place current-K update." - ) latent_cache = latent_cache.reshape(latent_cache.shape[0], -1).contiguous() num_tokens = latent_cache.shape[0] @@ -2484,17 +1991,8 @@ def scatter_fp4_mla_kv_cache( if phase == "context": if any(arg is not None for arg in (q_pe, q_rope_out, q_quant_input)): raise ValueError("FP4 MLA context cache update does not accept generation Q tensors.") - if (q_context is None) != (q_nope_head_dim is None): - raise ValueError("FP4 MLA context Q and q_nope_head_dim must be provided together.") hp_pool_updated = False else: - if q_context is not None or q_nope_head_dim is not None: - raise ValueError("FP4 MLA generation cache update does not accept context Q tensors.") - if (helix_position_offsets is None) != (helix_is_inactive_rank is None): - raise ValueError( - "FP4 MLA Helix position-offset and inactive-rank metadata " - "must be provided together." - ) if not all(arg is not None for arg in generation_inputs): raise ValueError( "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " @@ -2595,8 +2093,6 @@ def scatter_fp4_mla_kv_cache( v_sf, global_scale, rotary_cos_sin, - q_context, - q_nope_head_dim, token_offset=token_offset, local_layer=local_layer, v_head_dim=v_head_dim, @@ -2633,8 +2129,6 @@ def scatter_fp4_mla_kv_cache( q_sf_out=q_sf_out, v_packed_base=v_packed_base, v_page_offset=v_page_offset, - helix_position_offsets=helix_position_offsets, - helix_is_inactive_rank=helix_is_inactive_rank, ) v_pack_page_ids = _fp4_mla_generation_page_ids( metadata, metadata.num_seqs - metadata.num_contexts @@ -3956,7 +3450,6 @@ def run_fp4_mla_attention_decode( prequantized_q: torch.Tensor, prequantized_q_sf: torch.Tensor, q_batch_capacity: int, - softmax_stats_tensor: Optional[torch.Tensor] = None, ) -> None: """Run MLA decode with FP4 QK and FP4 PV tensor-core matmuls. @@ -4004,43 +3497,6 @@ def run_fp4_mla_attention_decode( raise ValueError("FP4 MLA attention output batch dimensions do not match.") backend = _fp4_mla_attention_backend() - helix_spec_tokens_valid = bool(getattr(metadata, "_helix_spec_tokens_valid", False)) - helix_kv_bounds = None - if softmax_stats_tensor is not None: - if backend != _FP4_MLA_CUTEDSL_BACKEND: - raise NotImplementedError( - "FP4 MLA Helix softmax stats require the cutedsl attention backend." - ) - if query_len_per_seq != 1 and not helix_spec_tokens_valid: - raise NotImplementedError( - "FP4 MLA multi-token Helix requires speculative per-token metadata." - ) - expected_stats_shape = (num_queries, num_heads, 2) - if ( - softmax_stats_tensor.shape != expected_stats_shape - or softmax_stats_tensor.dtype != torch.float32 - or softmax_stats_tensor.device != q.device - or not softmax_stats_tensor.is_contiguous() - ): - raise ValueError( - "FP4 MLA Helix requires contiguous same-device float32 softmax " - f"stats with shape {expected_stats_shape}." - ) - if helix_spec_tokens_valid: - helix_kv_bounds = getattr(metadata, "helix_kv_bounds", None) - if ( - not isinstance(helix_kv_bounds, torch.Tensor) - or helix_kv_bounds.dtype != torch.int32 - or helix_kv_bounds.device != q.device - or helix_kv_bounds.ndim != 1 - or helix_kv_bounds.numel() < num_queries - or not helix_kv_bounds.is_contiguous() - ): - raise ValueError( - "FP4 MLA speculative Helix KV bounds must be a contiguous " - "same-device int32 tensor covering every query token." - ) - helix_kv_bounds = helix_kv_bounds[:num_queries] if getattr(metadata, "fp4_mla_v_scale_pool", None) is None: raise RuntimeError( "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." @@ -4243,15 +3699,6 @@ def run_fp4_mla_attention_decode( ) kernel_output = output - kernel_softmax_stats = None - if softmax_stats_tensor is not None: - kernel_softmax_stats = _ensure_workspace_tensor( - metadata, - "_fp4_mla_cutedsl_softmax_stats_buf", - (2, num_queries, physical_heads), - dtype=torch.float32, - device=output.device, - ) if num_heads < physical_heads: kernel_output = _ensure_workspace_tensor( metadata, @@ -4282,15 +3729,9 @@ def run_fp4_mla_attention_decode( v_page_offset=v_page_offset, q_batch_capacity=q_batch_capacity, partition_runtime_valid_k=bool(getattr(metadata, "is_cuda_graph", False)), - softmax_row_max=(None if kernel_softmax_stats is None else kernel_softmax_stats[0]), - softmax_row_sum=(None if kernel_softmax_stats is None else kernel_softmax_stats[1]), - helix_kv_bounds=helix_kv_bounds, ) if kernel_output is not output: output.copy_(kernel_output[:, :num_heads]) - if kernel_softmax_stats is not None: - softmax_stats_tensor[..., 0].copy_(kernel_softmax_stats[0, :, :num_heads]) - softmax_stats_tensor[..., 1].copy_(kernel_softmax_stats[1, :, :num_heads]) return total_p_rows = num_queries * max_pages * num_heads diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py index b0c05229c288..6759b9026996 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py @@ -97,7 +97,6 @@ class _Fp8MlaContextScratch: max_num_sequences: int max_blocks_per_seq: int capacity_blocks: int - max_num_tokens: int page_size: int head_dim: int cache_stream: torch.cuda.Stream @@ -122,12 +121,6 @@ def create( max_num_sequences = int(meta.max_num_sequences or meta.max_num_requests) max_blocks_per_seq = int(kv_cache_manager.max_blocks_per_seq) max_num_tokens = int(meta.max_num_tokens) - runtime_features = meta.runtime_features - if runtime_features.chunked_prefill: - max_num_tokens *= max( - 1, - int(runtime_features.chunked_prefill_buffer_batch_size), - ) max_nonempty_sequences = min(max_num_sequences, max_num_tokens) capacity_blocks = max( 1, @@ -188,7 +181,6 @@ def create( max_num_sequences=max_num_sequences, max_blocks_per_seq=max_blocks_per_seq, capacity_blocks=capacity_blocks, - max_num_tokens=max_num_tokens, page_size=page_size, head_dim=head_dim, cache_stream=torch.cuda.Stream(device=device), @@ -205,43 +197,19 @@ def matches( head_dim: int, ) -> bool: kv_cache_manager = meta.kv_cache_manager - required_max_num_tokens = int(meta.max_num_tokens) - if meta.runtime_features.chunked_prefill: - required_max_num_tokens *= max( - 1, - int(meta.runtime_features.chunked_prefill_buffer_batch_size), - ) return ( kv_cache_manager is not None and self.pool.device == device and self.head_dim == head_dim - and self.max_num_tokens >= required_max_num_tokens and self.page_size == meta.tokens_per_block and self.cache_manager_view.max_seq_len == int(kv_cache_manager.max_seq_len) and self.max_num_sequences >= int(meta.max_num_sequences or meta.max_num_requests) and self.max_blocks_per_seq >= int(kv_cache_manager.max_blocks_per_seq) ) - def prepare( - self, - meta: "TrtllmAttentionMetadata", - *, - context_lengths_cuda: Optional[torch.Tensor] = None, - context_lengths_cpu: Optional[torch.Tensor] = None, - ) -> None: - if context_lengths_cuda is None: - context_lengths_cuda = meta.prompt_lens_cuda_runtime - if context_lengths_cpu is None: - context_lengths_cpu = meta.prompt_lens_cpu_runtime - if ( - context_lengths_cpu.device.type != "cpu" - or context_lengths_cpu.dtype != torch.int32 - or context_lengths_cpu.ndim != 1 - or context_lengths_cpu.numel() < meta.num_contexts - ): - raise ValueError("FP8 MLA context metadata requires a CPU int32 context-length tensor.") + def prepare(self, meta: "TrtllmAttentionMetadata") -> None: context_lengths = tuple( - int(length) for length in context_lengths_cpu[: meta.num_contexts].tolist() + int(length) for length in meta.prompt_lens_cpu_runtime[: meta.num_contexts].tolist() ) self.host_total_kv_lens[0] = sum(context_lengths) self.host_total_kv_lens[1] = 0 @@ -270,6 +238,7 @@ def prepare( f"{self.capacity_blocks} were allocated." ) + context_lengths_cuda = meta.prompt_lens_cuda_runtime if ( context_lengths_cuda.dtype != torch.int32 or not context_lengths_cuda.is_cuda @@ -299,14 +268,19 @@ def prepare( def _build_fp8_mla_context_attn(attn: "TrtllmAttention") -> "TrtllmAttention": """Build a direct-attribute FP8 view without per-access Python forwarding.""" + from ..fmha.fallback import FallbackFmha + from ..fmha.manager import FmhaManager + fp8_attn = copy.copy(attn) fp8_attn.quant_mode = int(QuantMode(0).set_fp8_kv_cache()) fp8_attn.has_fp4_kv_cache = False fp8_attn.has_fp8_kv_cache = True # FMHA instances hold weak references to their owning attention object. - # Do not reuse instances copied from the FP4 attention; the caller binds - # this FP8 view explicitly to TRTLLM's regular FMHA implementation. - fp8_attn.fmha_libs = [] + # Do not reuse the manager copied from the FP4 attention; bind this FP8 + # view explicitly to TRTLLM's regular FMHA implementation. + fp8_manager = FmhaManager(fp8_attn) + fp8_manager.fmha_libs = [FallbackFmha(fp8_attn)] + fp8_attn._fmha_manager = fp8_manager fp8_attn.local_layer_idx = 0 # This branch resolves local cache layers through layer_idx. The # disposable cache has exactly one layer, so bind the copied view to it. @@ -317,15 +291,8 @@ def _build_fp8_mla_context_attn(attn: "TrtllmAttention") -> "TrtllmAttention": def _build_fp8_mla_context_metadata( meta: "TrtllmAttentionMetadata", scratch: _Fp8MlaContextScratch, - *, - kv_lens_cuda: Optional[torch.Tensor] = None, - kv_lens_cpu: Optional[torch.Tensor] = None, ) -> "TrtllmAttentionMetadata": """Route the mandatory FP8 cache write through a direct metadata view.""" - if kv_lens_cuda is None: - kv_lens_cuda = meta.prompt_lens_cuda_runtime[: meta.num_contexts] - if kv_lens_cpu is None: - kv_lens_cpu = meta.prompt_lens_cpu_runtime[: meta.num_contexts] fp8_meta = copy.copy(meta) fp8_meta._fp4_mla_fp8_context_state = None fp8_meta.kv_cache_manager = scratch.cache_manager_view @@ -337,8 +304,8 @@ def _build_fp8_mla_context_metadata( # The disposable cache represents only this context invocation. Expose # exact context-only lengths so cached prefixes, trailing generation # metadata, and stale totals cannot extend FP8 K/V quantization. - fp8_meta.kv_lens_cuda_runtime = kv_lens_cuda - fp8_meta.kv_lens_runtime = kv_lens_cpu + fp8_meta.kv_lens_cuda_runtime = meta.prompt_lens_cuda_runtime[: meta.num_contexts] + fp8_meta.kv_lens_runtime = meta.prompt_lens_cpu_runtime[: meta.num_contexts] fp8_meta.host_total_kv_lens = scratch.host_total_kv_lens # Scratch lengths intentionally start from zero. Preserve the actual # absolute positions for Q/K RoPE through the native kernel's explicit diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 7187a8601215..785f06656c68 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -18,6 +18,19 @@ import cutlass.cute as cute import cutlass.pipeline as pipeline import torch +from ctm.Operations import ptx +from ctm.Operations.ptx import ( + AtomicOpKind, + CvtaSpace, + MBarrierArriveScope, + MBarrierArriveSem, + MBarrierSpace, + MemScopeKind, + SharedSpace, + cvta_to, + mbarrier_arrive, +) +from ctm.Operations.ptx import cp_async as _cp_async from cutlass._mlir.dialects import llvm from cutlass.base_dsl.dsl import BaseDSL from cutlass.cute.arch.nvvm_wrappers import inline_ptx as cute_inline_ptx @@ -25,12 +38,11 @@ from cutlass.experimental import cuda as cuda_tma from cutlass.experimental import primitives as prims -nvvm_add_packed_f32x2 = partial(prims.add_packed_f32x2, rnd=prims.FPRoundingMode.RN) -nvvm_mul_packed_f32x2 = partial(prims.mul_packed_f32x2, rnd=prims.FPRoundingMode.RN) -nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_add_packed_f32x2 = partial(ptx.add_packed_f32x2, rnd="rn") +nvvm_mul_packed_f32x2 = partial(ptx.mul_packed_f32x2, rnd="rn") +nvvm_fma_packed_f32x2 = partial(ptx.fma_packed_f32x2, rnd="rn") PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 -INT32_MAX = (1 << 31) - 1 _CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" _PYIR_STDOUT_LINES = frozenset( { @@ -47,7 +59,7 @@ def _as_shared_cta(address, *, loc=None, ip=None): address_ir = address.ir_value() if hasattr(address, "ir_value") else address address_space = llvm.PointerType(address_ir.type).address_space if address_space == ctm.AddressSpace.dsmem: - return prims.cvta_to(address, prims.CvtaSpace.SHARED, loc=loc, ip=ip) + return cvta_to(address, CvtaSpace.SHARED, loc=loc, ip=ip) return address @@ -58,10 +70,12 @@ def _mapa_shared_cluster(address, rank, *, loc=None, ip=None): @ctm.dsl_user_op def _mbarrier_arrive_release_cta_shared_cluster(mbar, count=1, *, loc=None, ip=None) -> None: - prims.mbarrier_arrive( + mbarrier_arrive( mbar, - count=count, - scope=prims.MemScope.CTA, + count, + sem=MBarrierArriveSem.RELEASE, + scope=MBarrierArriveScope.CTA, + space=MBarrierSpace.SHARED_CLUSTER, loc=loc, ip=ip, ) @@ -293,14 +307,9 @@ def _initial_float_from_argv(option: str, default: float) -> float: # normalization below cancels this factor without changing the attention math. FP4_MLA_E4M3_MAX_FINITE = 448.0 FP4_MLA_P_GLOBAL_SCALE = FP4_MLA_E4M3_MAX_FINITE * 6.0 -SMEM_P4_PAGE_PLAN_PROFILE_KV = 160 * 1024 -SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT = 4 -SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT = SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT * TRTLLM_PAGE_SIZE -# Keep the device-side ceil division ``valid_k + KV_TILE - 1`` in Int32 range. -SMEM_P4_RUNTIME_MAX_KV = ( - (INT32_MAX - (KV_TILE - 1)) // SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT -) * SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT -SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_PAGE_PLAN_PROFILE_KV // TRTLLM_PAGE_SIZE +SMEM_P4_RUNTIME_MAX_KV = 160 * 1024 +SMEM_P4_RUNTIME_MAX_PAGES = SMEM_P4_RUNTIME_MAX_KV // TRTLLM_PAGE_SIZE +SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_RUNTIME_MAX_PAGES SMEM_P4_PAGE_ID_PLAN_BYTES = SMEM_P4_PAGE_ID_PLAN_INTS * ctm.Int32.bytes SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS = 2 SMEM_P4_RUNTIME_SCALE_PAIR_BYTES = SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS * ctm.Float32.bytes @@ -501,40 +510,13 @@ def _initial_float_from_argv(option: str, default: float) -> float: _EXPLICIT_TORCH_STREAM: torch.cuda.Stream | None = None -@contextlib.contextmanager -def _launch_stream(): - """Yield the CUDA driver stream these kernels should launch on. - - cuda2ctl capture cannot reliably query the default null stream, so SMART - wrappers can ask for an explicit one through DKG_MLA_EXPLICIT_STREAM. The - substitution must be bracketed rather than assigned: the caller's stream - already carries the RoPE, Q quantization and KV-cache writes these kernels - read, and it is also the stream their consumers read the output from. - Entering waits on the caller's work, leaving publishes ours back to it, and - ``torch.cuda.stream`` restores the thread's current stream on the way out. - A bare ``set_stream`` does none of the three. - - Switching the current stream is illegal while a CUDA graph is capturing, so - the knob is ignored under capture and the caller's stream is used as-is. - """ - entry_stream = torch.cuda.current_stream() - if os.environ.get("DKG_MLA_EXPLICIT_STREAM") != "1" or torch.cuda.is_current_stream_capturing(): - yield cuda.CUstream(entry_stream.cuda_stream) - return - +def _current_cu_stream() -> cuda.CUstream: global _EXPLICIT_TORCH_STREAM - if _EXPLICIT_TORCH_STREAM is None: - _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() - start_event = torch.cuda.Event() - done_event = torch.cuda.Event() - start_event.record(entry_stream) - with torch.cuda.stream(_EXPLICIT_TORCH_STREAM): - _EXPLICIT_TORCH_STREAM.wait_event(start_event) - try: - yield cuda.CUstream(_EXPLICIT_TORCH_STREAM.cuda_stream) - finally: - done_event.record(_EXPLICIT_TORCH_STREAM) - entry_stream.wait_event(done_event) + if os.environ.get("DKG_MLA_EXPLICIT_STREAM") == "1": + if _EXPLICIT_TORCH_STREAM is None: + _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() + torch.cuda.set_stream(_EXPLICIT_TORCH_STREAM) + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) @dataclass(frozen=True) @@ -995,26 +977,6 @@ def _ceil_div(a: int, b: int) -> int: return (a + b - 1) // b -def _select_runtime_kv_profile(max_kv_len: int) -> int: - if type(max_kv_len) is not int: - raise TypeError(f"max_kv_len must be an int, got {type(max_kv_len).__name__}") - if max_kv_len <= 0: - raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") - if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: - raise ValueError( - "max_kv_len exceeds the Int32-safe runtime-KV limit: " - f"{max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" - ) - if max_kv_len <= SMEM_P4_PAGE_PLAN_PROFILE_KV: - return SMEM_P4_PAGE_PLAN_PROFILE_KV - # Eager execution can report a different batch maximum on every step. - # Shifted power-of-two buckets bound the number of compiled variants and - # retain the established 1 Mi-token plus four-page-reserve profile. - profile_payload = max_kv_len - SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT - profile_kv = (1 << (profile_payload - 1).bit_length()) + (SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT) - return min(profile_kv, SMEM_P4_RUNTIME_MAX_KV) - - def _validate_output_dtype(output_dtype: torch.dtype) -> None: if output_dtype not in {torch.float16, torch.bfloat16}: raise TypeError(f"output dtype must be torch.float16 or torch.bfloat16, got {output_dtype}") @@ -1068,7 +1030,6 @@ def fused_fp4_mla_decode_ctm( sfb_ptr: cute.Pointer, page_table_ptr: cute.Pointer, valid_k_ptr: cute.Pointer, - helix_kv_bounds_ptr: cute.Pointer, c_ptr: cute.Pointer, accum_ptr: cute.Pointer, row_max_ptr: cute.Pointer, @@ -1092,11 +1053,8 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, - use_helix_kv_bounds: ctm.Constexpr = False, - use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, - write_softmax_stats: ctm.Constexpr = False, ) -> None: n, k = problem_size m = runtime_m @@ -1117,7 +1075,9 @@ def fused_fp4_mla_decode_ctm( ) page_table_tensor = cute.make_tensor( cute.recast_ptr(page_table_ptr, dtype=cutlass.Int32), - cute.make_layout(((k // page_size) * (batch_size // query_len_per_seq),), stride=(1,)), + cute.make_layout( + (SMEM_P4_RUNTIME_MAX_PAGES * (batch_size // query_len_per_seq),), stride=(1,) + ), ) page_indptr_tensor = cute.make_tensor( cute.recast_ptr(page_indptr_ptr, dtype=cutlass.Int32), @@ -1135,10 +1095,6 @@ def fused_fp4_mla_decode_ctm( cute.recast_ptr(valid_k_ptr, dtype=cutlass.Int32), cute.make_layout((batch_size // query_len_per_seq,), stride=(1,)), ) - helix_kv_bounds_tensor = cute.make_tensor( - cute.recast_ptr(helix_kv_bounds_ptr, dtype=cutlass.Int32), - cute.make_layout((batch_size,), stride=(1,)), - ) v_tma_tensor = cute.make_tensor( b_ptr, cute.make_layout( @@ -1395,7 +1351,6 @@ def fused_fp4_mla_decode_ctm( page_table_tensor, page_indptr_tensor, valid_k_tensor, - helix_kv_bounds_tensor, q_global_scale_tensor, kv_global_scale_tensor, tma_q_desc, @@ -1423,12 +1378,9 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, - use_helix_kv_bounds, - use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, - write_softmax_stats, ).launch( grid=( cute.ceil_div(c_tensor.shape[0], SMEM_P4_CTA_GROUP_M) * CLUSTER_SHAPE_MNK[0], @@ -1541,31 +1493,6 @@ def _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx: ctm.Int32) -> t return (physical_page0, physical_page1) -@cute.jit -def _load_global_page_native_tile_pair( - mPageTable_pl: cute.Tensor, - kv_tile_idx: ctm.Int32, - page_begin: ctm.Int32, - page_count: ctm.Int32, -) -> tuple: - tile_page_idx = ctm.Int32(kv_tile_idx * SMEM_P4_PAGES_PER_KV_TILE) - # Keep these as two scalar loads: a CSR row may start at an odd Int32, and - # the second logical page must clamp to the first for a one-page tail. - physical_page0 = _lookup_physical_page( - mPageTable_pl, tile_page_idx, ctm.Int32(0), page_begin, page_count - ) - physical_page1 = _lookup_physical_page( - mPageTable_pl, - tile_page_idx + ctm.Int32(1), - ctm.Int32(0), - page_begin, - page_count, - ) - physical_page0 = cute.arch.make_warp_uniform(physical_page0) - physical_page1 = cute.arch.make_warp_uniform(physical_page1) - return (physical_page0, physical_page1) - - @cute.jit def _tma_gather4_cluster( smem_dst, @@ -1584,7 +1511,7 @@ def _tma_gather4_cluster( leader_barrier = _as_shared_cta(_mapa_shared_cluster(barrier, ctm.Int32(0))) barrier_ptr = leader_barrier.data_ptr() multicast_mask_u16 = ctm.Uint16(multicast_mask) - cute_inline_ptx( + _cp_async._predicated_inline_ptx( "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" " [{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], " @@ -2105,7 +2032,6 @@ def _load_v_tile_stage( v_tma_phase: ctm.Int32, stage: ctm.Int32 = 0, manage_mbarrier: ctm.Constexpr = True, - use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, ) -> None: del tidx @@ -2119,15 +2045,15 @@ def _load_v_tile_stage( if should_wait_v_tma: if prims.elect_sync(): prims.mbarrier_arrive_expect_tx(v_tma_mbar, v_tma_bytes) - del bidz, tile_physical_page0, tile_physical_page1 - if cutlass.const_expr(use_smem_page_plan): - resolved_physical_page0, resolved_physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) - else: - resolved_physical_page0, resolved_physical_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, kv_tile_idx, page_begin, page_count - ) + # Resolve the page pair at the point of use so V-ring wrap cannot reuse a + # stale loop-carried ID. The plan was populated once before the producer + # warps start, so reloading it here avoids a second CSR/global lookup and + # its serial uniform-address chain on every tile. + del mPageTable_pl, bidz, page_begin, page_count + del tile_physical_page0, tile_physical_page1 + resolved_physical_page0, resolved_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) if cutlass.const_expr(use_consecutive_page_pair): resolved_physical_page1 = resolved_physical_page0 + ctm.Int32(1) if prims.elect_sync(): @@ -2368,11 +2294,7 @@ def _softmax_exp2_pair_packed_f16x2( (score0, score1), (softmax_scale_log2, softmax_scale_log2), (p_bias, p_bias) ) shifted_h2 = _pack_f32_pair_to_f16x2(shifted[0], shifted[1]) - return cute_inline_ptx( - "ex2.approx.f16x2 {$w0}, {$r0};", - write_only_types=[ctm.Int32], - read_only_args=[shifted_h2], - ) + return ptx.ex2_f16x2(shifted_h2) @cute.jit @@ -3669,41 +3591,28 @@ def _smem_p4_p4_materialize_group_col( @cute.jit def _float_to_ordered_u32_for_atomic_max(value: ctm.Float32) -> ctm.Uint32: - bits = prims.mov_b32(value, target_type=ctm.Int32) + bits = ptx.mov_b32(value, target_type=ctm.Int32) sign_mask = bits >> ctm.Int32(31) | ctm.Int32(2147483648) encoded = bits ^ sign_mask - return prims.mov_b32(encoded, target_type=ctm.Uint32) + return ptx.mov_b32(encoded, target_type=ctm.Uint32) @cute.jit def _ordered_u32_to_float_after_atomic_max(value: ctm.Uint32) -> ctm.Float32: - encoded = prims.mov_b32(value, target_type=ctm.Int32) + encoded = ptx.mov_b32(value, target_type=ctm.Int32) sign_mask = ~(encoded >> ctm.Int32(31)) | ctm.Int32(2147483648) bits = encoded ^ sign_mask - return prims.mov_b32(bits, target_type=ctm.Float32) + return ptx.mov_b32(bits, target_type=ctm.Float32) @cute.jit def _smem_atomic_max_ordered_u32(pointer, value: ctm.Uint32) -> None: - prims.atomicrmw( - prims.AtomicOp.MAX, + ptx.atom( + AtomicOpKind.MAX, pointer, value, - syncscope=prims.MemScope.CTA, - space=prims.SharedSpace.shared_cta, - ) - - -@cute.jit -def _tcgen05_ld_red_32x32b_x16_max_f32(tmem) -> tuple: - tmem_addr = tmem.toint(ctm.Int32) - return cute_inline_ptx( - "tcgen05.ld.red.sync.aligned.32x32b.x16.max.f32 " - "{{$w0}, {$w1}, {$w2}, {$w3}, {$w4}, {$w5}, {$w6}, {$w7}, " - "{$w8}, {$w9}, {$w10}, {$w11}, {$w12}, {$w13}, {$w14}, {$w15}}, " - "{$w16}, [{$r0}];", - write_only_types=[ctm.Int32] * 17, - read_only_args=[tmem_addr], + syncscope=MemScopeKind.CTA, + space=SharedSpace.shared_cta, ) @@ -3789,7 +3698,13 @@ def _load_p4_n256_score_half_from_tmem( pending_score_groups.append(group_scores) pending_group_stats.append(None) else: - regs = _tcgen05_ld_red_32x32b_x16_max_f32(tmem) + regs = ptx.tcgen05_ld_red( + ptx.Tcgen05LdStShape.SHAPE_32X32B, + tmem, + num=SF_VEC_SIZE, + red_op="max", + type_="f32", + ) pending_score_groups.append(regs) pending_group_stats.append(regs[SF_VEC_SIZE]) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) @@ -4638,7 +4553,6 @@ def _store_final_o_from_tmem( final_row_sum: ctm.Float32, final_stat_scale: ctm.Float32, output_normalizer: ctm.Float32, - write_softmax_stats: ctm.Constexpr = False, producer_warp_base: ctm.Constexpr = 0, ) -> None: gC_arr = ctm.make_array_view(mC_mnl) @@ -4658,10 +4572,6 @@ def _store_final_o_from_tmem( + local_row ) batch_offset = bidz * m * n - if cutlass.const_expr(write_softmax_stats): - if col_band == ctm.Int32(0) and bidy == ctm.Int32(0): - mRowMax_ml[row, bidz] = final_row_max * final_stat_scale - mRowSum_ml[row, bidz] = final_row_sum n_tile_total = n // ctm.Int32(OUT_DIM) subtile_cols: ctm.Constexpr = 32 subtiles_per_n_tile: ctm.Constexpr = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N // subtile_cols @@ -5522,13 +5432,14 @@ def _runtime_t336_producer_tile( ) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - if common_source_ready_pred: - prims.mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - count=1, - scope=prims.MemScope.CLUSTER, - relaxed=True, - ) + mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + 1, + sem=MBarrierArriveSem.RELAXED, + scope=MBarrierArriveScope.CLUSTER, + space=MBarrierSpace.SHARED_CLUSTER, + pred=common_source_ready_pred, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -5780,13 +5691,14 @@ def _runtime_t336_producer_tile_v23( warp_rebase = cute.arch.vote_any_sync(lane_rebase) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - if common_source_ready_pred: - prims.mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - count=1, - scope=prims.MemScope.CLUSTER, - relaxed=True, - ) + mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + 1, + sem=MBarrierArriveSem.RELAXED, + scope=MBarrierArriveScope.CLUSTER, + space=MBarrierSpace.SHARED_CLUSTER, + pred=common_source_ready_pred, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -5923,7 +5835,6 @@ def _run_mla_decode_body( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, - mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_ptr, @@ -5951,12 +5862,9 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, - use_helix_kv_bounds: ctm.Constexpr, - use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, - write_softmax_stats: ctm.Constexpr = False, ) -> None: pv_psf_rescale = ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) * pv_output_scale warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -5967,14 +5875,11 @@ def _run_mla_decode_body( page_batch = ctm.Int32(0) valid_k_for_l = ctm.Int32(0) page_batch = cute.arch.make_warp_uniform(bidz // ctm.Int32(query_len_per_seq)) - if cutlass.const_expr(use_helix_kv_bounds): - valid_k_for_l = mHelixKvBounds_l[bidz] - else: - query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) - valid_k_for_l = ctm.max( - mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), - ctm.Int32(0), - ) + query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) + valid_k_for_l = ctm.max( + mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), + ctm.Int32(0), + ) valid_k_for_l = cute.arch.make_warp_uniform(valid_k_for_l) csr_page_begin = ctm.Int32(0) csr_page_count = ctm.Int32(0) @@ -6252,10 +6157,9 @@ def _run_mla_decode_body( _stage_smem_p4_runtime_scale_pair( mQGlobalScale, mKvGlobalScale, sRuntimeScalePair, softmax_scale_log2 ) - if cutlass.const_expr(use_smem_page_plan): - _stage_smem_p4_page_id_plan( - mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count - ) + _stage_smem_p4_page_id_plan( + mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count + ) cute.arch.cluster_arrive() cute.arch.cluster_wait() runtime_scale_pair = sRuntimeScalePair.data_ptr().load(count=2, alignment=8) @@ -6305,14 +6209,7 @@ def _run_mla_decode_body( qk0_handle = qk_smem_producer.acquire_and_advance() qk0_tma_mbar = ctm.Array(qk0_handle.barrier, shape=1) if is_leader_cta: - if cutlass.const_expr(use_smem_page_plan): - qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(0) - ) - else: - qk0_page0, qk0_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, ctm.Int32(0), csr_page_begin, csr_page_count - ) + qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, ctm.Int32(0)) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6346,14 +6243,9 @@ def _run_mla_decode_body( ) qk_prefix_tma_mbar = ctm.Array(qk_prefix_handle.barrier, shape=1) if is_leader_cta: - if cutlass.const_expr(use_smem_page_plan): - qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) - else: - qk_prefix_page0, qk_prefix_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count - ) + qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6390,14 +6282,9 @@ def _run_mla_decode_body( ) qk_tma_mbar = ctm.Array(qk_handle.barrier, shape=1) if is_leader_cta: - if cutlass.const_expr(use_smem_page_plan): - qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) - else: - qk_steady_page0, qk_steady_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count - ) + qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6429,14 +6316,9 @@ def _run_mla_decode_body( ) qk15_tma_mbar = ctm.Array(qk15_handle.barrier, shape=1) if is_leader_cta: - if cutlass.const_expr(use_smem_page_plan): - qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(15) - ) - else: - qk15_page0, qk15_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, ctm.Int32(15), csr_page_begin, csr_page_count - ) + qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(15) + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6466,12 +6348,9 @@ def _run_mla_decode_body( cute.arch.cp_async_bulk_wait_group(0, read=True) q_smem_producer.tail() if warp_idx == TMA_V_WARP_ID: - v_physical_page0 = ctm.Int32(0) - v_physical_page1 = ctm.Int32(0) - if cutlass.const_expr(use_smem_page_plan): - v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(0) - ) + v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(0) + ) for kv_tile_idx in cutlass.range(0, kv_tiles, 1, unroll=1): v_handle = v_smem_producer.acquire_and_advance() v_tma_mbar = ctm.Array(v_handle.barrier, shape=1) @@ -6500,27 +6379,20 @@ def _run_mla_decode_body( v_tma_phase=v_stage, stage=v_stage, manage_mbarrier=False, - use_smem_page_plan=use_smem_page_plan, use_consecutive_page_pair=use_consecutive_page_pair, ) - if cutlass.const_expr(use_smem_page_plan): - if kv_tile_idx + ctm.Int32(1) < kv_tiles: - v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx + ctm.Int32(1) - ) + if kv_tile_idx + ctm.Int32(1) < kv_tiles: + v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ctm.Int32(1) + ) v_smem_producer.tail() if warp_idx == ROWMETA_WARP_ID: if is_leader_cta: qk_prefix_end = min(kv_tiles, ctm.Int32(3)) for qk_prefix_li_idx in cutlass.range(1, qk_prefix_end, 1, unroll=1): - if cutlass.const_expr(use_smem_page_plan): - rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) - else: - rowmeta_page0, rowmeta_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count - ) + rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) _issue_runtime_t336_qk_prefix_rank1( mPageTable_pl, tma_k_page_ptr, @@ -6925,19 +6797,8 @@ def _run_mla_decode_body( final_anchor_row_sum = sFinalAnchorRowSum[final_row_state_local_row] final_stat_scale = ctm.Float32(1.0) final_row_sum = final_anchor_row_sum - final_row_max = running_row_max - if cutlass.const_expr(write_softmax_stats): - final_stat_scale = softmax_scale_log2 / ctm.Float32(LOG2_E) - final_row_sum = final_anchor_row_sum * cute.exp2( - (running_row_anchor - running_row_max) * softmax_scale_log2, - fastmath=True, - ) - if valid_k_for_l == ctm.Int32(0): - final_stat_scale = ctm.Float32(1.0) - final_row_max = ctm.Float32(-ctm.Float32.inf) - final_row_sum = ctm.Float32(0.0) output_normalizer = ctm.Float32(0.0) - if valid_k_for_l != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): + if final_anchor_row_sum != ctm.Float32(0.0): output_normalizer = cute.arch.rcp_approx(final_anchor_row_sum) * pv_output_scale last_pv_li_idx = stream_li_total - ctm.Int32(1) last_pv_slot = last_pv_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -6963,11 +6824,10 @@ def _run_mla_decode_body( bidz, m, n, - final_row_max=final_row_max, + final_row_max=running_row_max, final_row_sum=final_row_sum, final_stat_scale=final_stat_scale, output_normalizer=output_normalizer, - write_softmax_stats=write_softmax_stats, producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, ) prims.barrier(barrier_id=O_STORE_BAR_ID, number_of_threads=O_STORE_BAR_THREADS) @@ -6986,7 +6846,6 @@ def kernel( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, - mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_desc: ctm.GridConstant[cuda_tma.TensorMap], @@ -7014,18 +6873,14 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, - use_helix_kv_bounds: ctm.Constexpr, - use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, - write_softmax_stats: ctm.Constexpr, ) -> None: _run_mla_decode_body( mPageTable_pl, mPageIndptr_s, mValidK_l, - mHelixKvBounds_l, mQGlobalScale, mKvGlobalScale, tma_q_desc.get_ptr(), @@ -7053,12 +6908,9 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, - use_helix_kv_bounds, - use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, - write_softmax_stats, ) @@ -7079,7 +6931,6 @@ def _make_fused_ptrs( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, - helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7099,7 +6950,6 @@ def _make_fused_ptrs( make_ptr(cutlass.Uint8, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=32), make_ptr(cutlass.Int32, page_table_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr(cutlass.Int32, valid_k_data_ptr, cute.AddressSpace.gmem, assumed_align=4), - make_ptr(cutlass.Int32, helix_kv_bounds_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr( _cutlass_output_dtype(output_dtype), c_data_ptr, @@ -7289,7 +7139,6 @@ def _compile_fused( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, - helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7311,36 +7160,20 @@ def _compile_fused( ksf_page_stride_bytes: int = 0, vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, - write_softmax_stats: bool = False, - use_helix_kv_bounds: bool = False, ) -> Callable: - if type(kv) is not int: - raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") - is_page_plan_profile = kv == SMEM_P4_PAGE_PLAN_PROFILE_KV - is_bucketed_runtime_profile = ( - SMEM_P4_PAGE_PLAN_PROFILE_KV < kv <= SMEM_P4_RUNTIME_MAX_KV - and _select_runtime_kv_profile(kv) == kv - ) - if not (is_page_plan_profile or is_bucketed_runtime_profile): - raise ValueError( - "runtime-KV compile requires the fixed page-plan profile " - f"{SMEM_P4_PAGE_PLAN_PROFILE_KV} or a larger geometric profile " - f"up to {SMEM_P4_RUNTIME_MAX_KV}, got {kv}" - ) + if kv != SMEM_P4_RUNTIME_MAX_KV: + raise ValueError(f"runtime-KV compile requires fixed K={SMEM_P4_RUNTIME_MAX_KV}, got {kv}") use_ksf_gather4 = ( not use_consecutive_page_pair and ksf_page_stride_bytes == page_size * TRTLLM_K_SF_GROUPS ) cache_key = ( n, - kv, page_size, use_mixed_imlp, output_dtype, query_len_per_seq, use_consecutive_page_pair, use_ksf_gather4, - write_softmax_stats, - use_helix_kv_bounds, ) cached = _FUSED_COMPILE_CACHE.get(cache_key) if cached is not None: @@ -7355,7 +7188,6 @@ def _compile_fused( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, - helix_kv_bounds_data_ptr, c_data_ptr, accum_data_ptr, row_max_data_ptr, @@ -7384,11 +7216,8 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, - use_helix_kv_bounds=use_helix_kv_bounds, - use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, - write_softmax_stats=write_softmax_stats, options="--opt-level 2 --ptxas-options '--uumn'", ) _FUSED_COMPILE_CACHE[cache_key] = compiled @@ -7420,9 +7249,6 @@ def run_trtllm_fp4_mla_decode_page_native( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, - softmax_row_max: torch.Tensor | None = None, - softmax_row_sum: torch.Tensor | None = None, - helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7437,7 +7263,12 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"page-native decode requires page_size={TRTLLM_PAGE_SIZE}, got {page_size}" ) - physical_k = _select_runtime_kv_profile(max_kv_len) + if max_kv_len <= 0: + raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") + if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: + raise ValueError( + f"max_kv_len exceeds the fixed runtime-KV profile: {max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" + ) for name, value in ( ("assume_valid_k_prefix_tiles", assume_valid_k_prefix_tiles), ( @@ -7463,30 +7294,13 @@ def run_trtllm_fp4_mla_decode_page_native( del assume_valid_k_prefix_tiles, partition_runtime_valid_k required_consecutive_tiles = _ceil_div(max_kv_len, KV_TILE) use_consecutive_page_pair = assume_consecutive_page_prefix_tiles >= required_consecutive_tiles + physical_k = SMEM_P4_RUNTIME_MAX_KV if q_internal.dtype != torch.uint8 or q_internal.dim() != 3: raise TypeError( "q_internal must be a uint8 [M, Q640/2, L] tensor, " f"got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" ) physical_m, q_bytes, l_batch = q_internal.shape - if (softmax_row_max is None) != (softmax_row_sum is None): - raise ValueError("softmax_row_max and softmax_row_sum must be provided together") - write_softmax_stats = softmax_row_max is not None - if write_softmax_stats: - expected_stats_shape = (l_batch, physical_m) - for name, tensor in ( - ("softmax_row_max", softmax_row_max), - ("softmax_row_sum", softmax_row_sum), - ): - if ( - tensor.dtype != torch.float32 - or tensor.shape != expected_stats_shape - or not tensor.is_contiguous() - ): - raise ValueError( - f"{name} must be contiguous float32 with shape {expected_stats_shape}" - ) - _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) if l_batch <= 0 or l_batch > CUDA_GRID_Z_MAX: raise ValueError(f"queries must be in [1, {CUDA_GRID_Z_MAX}], got {l_batch}") if q_batch_capacity is None: @@ -7540,15 +7354,10 @@ def run_trtllm_fp4_mla_decode_page_native( "src_page_ids must be a 1D int32 physical-page list, " f"got dtype={src_page_ids.dtype} shape={tuple(src_page_ids.shape)}" ) - src_page_id_count = src_page_ids.numel() - if src_page_id_count == 0 or src_page_ids.stride(0) != 1: + if src_page_ids.numel() == 0 or src_page_ids.stride(0) != 1: raise ValueError( "src_page_ids must be non-empty and contiguous, " - f"got numel={src_page_id_count} stride={src_page_ids.stride()}" - ) - if src_page_id_count > INT32_MAX: - raise ValueError( - f"src_page_ids exceeds the Int32 element limit: {src_page_id_count} > {INT32_MAX}" + f"got numel={src_page_ids.numel()} stride={src_page_ids.stride()}" ) if ( paged_kv_indptr_decode.dtype != torch.int32 @@ -7566,16 +7375,6 @@ def run_trtllm_fp4_mla_decode_page_native( f"got dtype={valid_k.dtype} shape={tuple(valid_k.shape)} " f"stride={valid_k.stride()}" ) - if helix_kv_bounds is not None and ( - helix_kv_bounds.dtype != torch.int32 - or helix_kv_bounds.shape != (l_batch,) - or helix_kv_bounds.stride(0) != 1 - ): - raise ValueError( - "helix_kv_bounds must be contiguous int32 with shape " - f"[{l_batch}], got dtype={helix_kv_bounds.dtype} " - f"shape={tuple(helix_kv_bounds.shape)} stride={helix_kv_bounds.stride()}" - ) cache_layout = _validate_v_packed_cache_args( v_packed, kv_cache, @@ -7596,14 +7395,9 @@ def run_trtllm_fp4_mla_decode_page_native( ) num_cache_pages = cache_layout.num_pages page_table_capacity = num_sequences * (physical_k // page_size) - if page_table_capacity > INT32_MAX: - raise ValueError( - "bucketed CSR capacity exceeds the Int32 layout limit: " - f"{page_table_capacity} > {INT32_MAX}" - ) - if src_page_id_count > page_table_capacity: + if src_page_ids.numel() > page_table_capacity: raise ValueError( - f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_id_count} > {page_table_capacity}" + f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_ids.numel()} > {page_table_capacity}" ) expected_v_rows = num_cache_pages * (TRTLLM_V_HEAD_DIM // v_pack_block) * v_pack_block if ( @@ -7626,11 +7420,12 @@ def run_trtllm_fp4_mla_decode_page_native( if not isinstance(v_page_offset, int): raise TypeError(f"v_page_offset must be an int, got {type(v_page_offset).__name__}") num_v_cache_pages = v_packed.shape[0] // v_rows_per_page - if v_page_offset < 0 or v_page_offset > INT32_MAX: - raise ValueError(f"v_page_offset must be in [0, {INT32_MAX}], got {v_page_offset}") - if num_v_cache_pages > INT32_MAX: + int32_max = torch.iinfo(torch.int32).max + if v_page_offset < 0 or v_page_offset > int32_max: + raise ValueError(f"v_page_offset must be in [0, {int32_max}], got {v_page_offset}") + if num_v_cache_pages > int32_max: raise ValueError( - f"v_packed exceeds the Int32 physical-page limit: {num_v_cache_pages} > {INT32_MAX}" + f"v_packed exceeds the Int32 physical-page limit: {num_v_cache_pages} > {int32_max}" ) if v_page_offset + num_cache_pages > num_v_cache_pages: raise ValueError( @@ -7692,187 +7487,170 @@ def run_trtllm_fp4_mla_decode_page_native( q_global_scale, kv_global_scale, ) - if write_softmax_stats: - tensors += (softmax_row_max, softmax_row_sum) - if helix_kv_bounds is not None: - tensors += (helix_kv_bounds,) if device.type != "cuda" or any((tensor.device != device for tensor in tensors)): raise ValueError("all page-native decode tensors must share one CUDA device") if q_global_scale.dtype != torch.float32 or q_global_scale.numel() != 1: raise TypeError("q_global_scale must be a scalar FP32 CUDA tensor") if kv_global_scale.dtype != torch.float32 or kv_global_scale.numel() != 1: raise TypeError("kv_global_scale must be a scalar FP32 CUDA tensor") - with _launch_stream() as stream: - device_index = torch.cuda.current_device() - context_result, current_context = cuda.cuCtxGetCurrent() - if context_result != cuda.CUresult.CUDA_SUCCESS: + stream = _current_cu_stream() + device_index = torch.cuda.current_device() + context_result, current_context = cuda.cuCtxGetCurrent() + if context_result != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError( + f"Failed to query the current CUDA context for FP4 MLA launch: {context_result}." + ) + q_data_ptr = q_internal.data_ptr() + k_data_ptr = kv_cache.data_ptr() + q_sf_data_ptr = q_sf_internal.data_ptr() + k_sf_data_ptr = sf_cache.data_ptr() + b_data_ptr = v_packed.data_ptr() + scratch_ptr = output.data_ptr() + sfb_data_ptr = v_sf.data_ptr() + page_table_data_ptr = src_page_ids.data_ptr() + valid_k_data_ptr = valid_k.data_ptr() + c_data_ptr = output.data_ptr() + page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() + q_global_scale_data_ptr = q_global_scale.data_ptr() + kv_global_scale_data_ptr = kv_global_scale.data_ptr() + kv_page_stride_bytes = cache_layout.stride_page + ksf_page_stride_bytes = int(sf_cache.stride(0)) + vsf_page_stride_bytes = int(v_sf.stride(0)) + softmax_scale_log2 = sm_scale * LOG2_E + fused = _compile_fused( + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + scratch_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + c_data_ptr, + scratch_ptr, + scratch_ptr, + scratch_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + physical_m, + TRTLLM_V_HEAD_DIM, + physical_k, + l_batch, + page_size, + use_mixed_imlp=enable_mxi_imlp, + output_dtype=output.dtype, + stream=stream, + num_cache_pages=num_cache_pages, + query_len_per_seq=query_len_per_seq, + kv_page_stride_bytes=kv_page_stride_bytes, + ksf_page_stride_bytes=ksf_page_stride_bytes, + vsf_page_stride_bytes=vsf_page_stride_bytes, + use_consecutive_page_pair=use_consecutive_page_pair, + ) + supports_prepared = _class_defines_callables( + fused, "to", "generate_execution_args", "run_compiled_program" + ) + prepared_key = ( + fused, + device_index, + int(current_context), + int(stream), + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + c_data_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + physical_m, + l_batch, + q_batch_capacity, + num_cache_pages, + num_v_cache_pages, + v_page_offset, + kv_page_stride_bytes, + ksf_page_stride_bytes, + vsf_page_stride_bytes, + softmax_scale_log2, + output.dtype, + enable_mxi_imlp, + use_consecutive_page_pair, + ) + if supports_prepared: + prepared = _get_prepared_fused_call(prepared_key) + if prepared is not None: + prepared.run() + return + ptrs = _make_fused_ptrs( + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + scratch_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + c_data_ptr, + scratch_ptr, + scratch_ptr, + scratch_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + output.dtype, + ) + runtime_args = ( + *ptrs, + (TRTLLM_V_HEAD_DIM, physical_k), + ctm.Int32(physical_m), + ctm.Int32(l_batch), + ctm.Int32(q_batch_capacity), + ctm.Int32(num_cache_pages), + ctm.Int32(num_v_cache_pages), + ctm.Int32(v_page_offset), + ctm.Int64(kv_page_stride_bytes), + ctm.Int64(ksf_page_stride_bytes), + ctm.Int64(vsf_page_stride_bytes), + ctm.Float32(softmax_scale_log2), + ctm.Float32(1.0), + stream, + ) + if not supports_prepared: + fused(*runtime_args) + return + executor_key = (fused, device_index, int(current_context)) + executor = _get_fused_executor(executor_key) + if executor is None: + if torch.cuda.is_current_stream_capturing(): raise RuntimeError( - f"Failed to query the current CUDA context for FP4 MLA launch: {context_result}." + "FP4 MLA CUDA Graph capture requires an eager warmup for the compiled kernel, device, and CUDA context." ) - q_data_ptr = q_internal.data_ptr() - k_data_ptr = kv_cache.data_ptr() - q_sf_data_ptr = q_sf_internal.data_ptr() - k_sf_data_ptr = sf_cache.data_ptr() - b_data_ptr = v_packed.data_ptr() - scratch_ptr = output.data_ptr() - row_max_data_ptr = softmax_row_max.data_ptr() if write_softmax_stats else scratch_ptr - row_sum_data_ptr = softmax_row_sum.data_ptr() if write_softmax_stats else scratch_ptr - sfb_data_ptr = v_sf.data_ptr() - page_table_data_ptr = src_page_ids.data_ptr() - valid_k_data_ptr = valid_k.data_ptr() - helix_kv_bounds_data_ptr = ( - helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else valid_k_data_ptr - ) - c_data_ptr = output.data_ptr() - page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() - q_global_scale_data_ptr = q_global_scale.data_ptr() - kv_global_scale_data_ptr = kv_global_scale.data_ptr() - kv_page_stride_bytes = cache_layout.stride_page - ksf_page_stride_bytes = int(sf_cache.stride(0)) - vsf_page_stride_bytes = int(v_sf.stride(0)) - softmax_scale_log2 = sm_scale * LOG2_E - fused = _compile_fused( - q_data_ptr, - k_data_ptr, - q_sf_data_ptr, - k_sf_data_ptr, - b_data_ptr, - scratch_ptr, - sfb_data_ptr, - page_table_data_ptr, - valid_k_data_ptr, - helix_kv_bounds_data_ptr, - c_data_ptr, - scratch_ptr, - row_max_data_ptr, - row_sum_data_ptr, - page_indptr_data_ptr, - q_global_scale_data_ptr, - kv_global_scale_data_ptr, - physical_m, - TRTLLM_V_HEAD_DIM, - physical_k, - l_batch, - page_size, - use_mixed_imlp=enable_mxi_imlp, - output_dtype=output.dtype, - stream=stream, - num_cache_pages=num_cache_pages, - query_len_per_seq=query_len_per_seq, - kv_page_stride_bytes=kv_page_stride_bytes, - ksf_page_stride_bytes=ksf_page_stride_bytes, - vsf_page_stride_bytes=vsf_page_stride_bytes, - use_consecutive_page_pair=use_consecutive_page_pair, - write_softmax_stats=write_softmax_stats, - use_helix_kv_bounds=helix_kv_bounds is not None, - ) - supports_prepared = _class_defines_callables( - fused, "to", "generate_execution_args", "run_compiled_program" - ) - prepared_key = ( - fused, - device_index, - int(current_context), - int(stream), - q_data_ptr, - k_data_ptr, - q_sf_data_ptr, - k_sf_data_ptr, - b_data_ptr, - sfb_data_ptr, - page_table_data_ptr, - valid_k_data_ptr, - helix_kv_bounds_data_ptr, - c_data_ptr, - row_max_data_ptr, - row_sum_data_ptr, - page_indptr_data_ptr, - q_global_scale_data_ptr, - kv_global_scale_data_ptr, - physical_m, - l_batch, - q_batch_capacity, - num_cache_pages, - num_v_cache_pages, - v_page_offset, - kv_page_stride_bytes, - ksf_page_stride_bytes, - vsf_page_stride_bytes, - softmax_scale_log2, - output.dtype, - enable_mxi_imlp, - use_consecutive_page_pair, - ) - if supports_prepared: - prepared = _get_prepared_fused_call(prepared_key) - if prepared is not None: - prepared.run() - return - ptrs = _make_fused_ptrs( - q_data_ptr, - k_data_ptr, - q_sf_data_ptr, - k_sf_data_ptr, - b_data_ptr, - scratch_ptr, - sfb_data_ptr, - page_table_data_ptr, - valid_k_data_ptr, - helix_kv_bounds_data_ptr, - c_data_ptr, - scratch_ptr, - row_max_data_ptr, - row_sum_data_ptr, - page_indptr_data_ptr, - q_global_scale_data_ptr, - kv_global_scale_data_ptr, - output.dtype, - ) - runtime_args = ( - *ptrs, - (TRTLLM_V_HEAD_DIM, physical_k), - ctm.Int32(physical_m), - ctm.Int32(l_batch), - ctm.Int32(q_batch_capacity), - ctm.Int32(num_cache_pages), - ctm.Int32(num_v_cache_pages), - ctm.Int32(v_page_offset), - ctm.Int64(kv_page_stride_bytes), - ctm.Int64(ksf_page_stride_bytes), - ctm.Int64(vsf_page_stride_bytes), - ctm.Float32(softmax_scale_log2), - ctm.Float32(1.0), - stream, - ) - if not supports_prepared: + candidate = fused.to(device_index) + if not _class_defines_callables( + candidate, "generate_execution_args", "run_compiled_program" + ): fused(*runtime_args) return - executor_key = (fused, device_index, int(current_context)) - executor = _get_fused_executor(executor_key) - if executor is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "FP4 MLA CUDA Graph capture requires an eager warmup for the " - "compiled kernel, device, and CUDA context." - ) - candidate = fused.to(device_index) - if not _class_defines_callables( - candidate, "generate_execution_args", "run_compiled_program" - ): - fused(*runtime_args) - return - executor = _cache_fused_executor(executor_key, candidate) - execution_args, adapted_args = executor.generate_execution_args(*runtime_args) - prepared = _cache_prepared_fused_call( - prepared_key, - _PreparedFusedCall( - executor=executor, - runtime_args=runtime_args, - execution_args=execution_args, - adapted_args=adapted_args, - ), - ) - prepared.run() + executor = _cache_fused_executor(executor_key, candidate) + execution_args, adapted_args = executor.generate_execution_args(*runtime_args) + prepared = _cache_prepared_fused_call( + prepared_key, + _PreparedFusedCall( + executor=executor, + runtime_args=runtime_args, + execution_args=execution_args, + adapted_args=adapted_args, + ), + ) + prepared.run() def run_trtllm_fp4_mla_decode_page_native_from_raw( @@ -7901,9 +7679,6 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, - softmax_row_max: torch.Tensor | None = None, - softmax_row_sum: torch.Tensor | None = None, - helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7955,7 +7730,4 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles=assume_consecutive_page_prefix_tiles, partition_runtime_valid_k=partition_runtime_valid_k, enable_mxi_imlp=enable_mxi_imlp, - softmax_row_max=softmax_row_max, - softmax_row_sum=softmax_row_sum, - helix_kv_bounds=helix_kv_bounds, ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py index 0c2f590462e3..df51457a43c2 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py @@ -19,6 +19,18 @@ import cutlass.cute as cute import cutlass.pipeline as pipeline import torch +from ctm.Operations import ptx +from ctm.Operations.ptx import ( + AtomicOpKind, + CvtaSpace, + MBarrierArriveScope, + MBarrierArriveSem, + MBarrierSpace, + MemScopeKind, + SharedSpace, + cvta_to, +) +from ctm.Operations.ptx import cp_async as _cp_async from cutlass._mlir.dialects import llvm from cutlass.base_dsl.array import EvictPriority from cutlass.base_dsl.dsl import BaseDSL @@ -27,12 +39,11 @@ from cutlass.experimental import cuda as cuda_tma from cutlass.experimental import primitives as prims -nvvm_add_packed_f32x2 = partial(prims.add_packed_f32x2, rnd=prims.FPRoundingMode.RN) -nvvm_mul_packed_f32x2 = partial(prims.mul_packed_f32x2, rnd=prims.FPRoundingMode.RN) -nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_add_packed_f32x2 = partial(ptx.add_packed_f32x2, rnd="rn") +nvvm_mul_packed_f32x2 = partial(ptx.mul_packed_f32x2, rnd="rn") +nvvm_fma_packed_f32x2 = partial(ptx.fma_packed_f32x2, rnd="rn") PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 -INT32_MAX = (1 << 31) - 1 _CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" _PYIR_STDOUT_LINES = frozenset( { @@ -48,7 +59,7 @@ def _as_tma_completion_mbar(mbar, *, loc=None, ip=None): mbar_ir = mbar.ir_value() if hasattr(mbar, "ir_value") else mbar if llvm.PointerType(mbar_ir.type).address_space == ctm.AddressSpace.dsmem: - return prims.cvta_to(mbar, prims.CvtaSpace.SHARED, loc=loc, ip=ip) + return cvta_to(mbar, CvtaSpace.SHARED, loc=loc, ip=ip) return mbar @@ -62,10 +73,12 @@ def _mapa_shared_cluster(mbar, rank, *, loc=None, ip=None): def _mbarrier_arrive_shared_cluster(mbar, count=1, *, loc=None, ip=None) -> None: # Keep the validated release.cta form. The cluster-scoped NVVM form fails # to lower on the target toolchain (CUDA error 715). - prims.mbarrier_arrive( + ptx.mbarrier_arrive( mbar, - count=count, - scope=prims.MemScope.CTA, + count, + sem=MBarrierArriveSem.RELEASE, + scope=MBarrierArriveScope.CTA, + space=MBarrierSpace.SHARED_CLUSTER, loc=loc, ip=ip, ) @@ -295,14 +308,9 @@ def _initial_float_from_argv(option: str, default: float) -> float: # normalization below cancels this factor without changing the attention math. FP4_MLA_E4M3_MAX_FINITE = 448.0 FP4_MLA_P_GLOBAL_SCALE = FP4_MLA_E4M3_MAX_FINITE * 6.0 -SMEM_P4_PAGE_PLAN_PROFILE_KV = 160 * 1024 -SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT = 4 -SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT = SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT * TRTLLM_PAGE_SIZE -# Keep the device-side ceil division ``valid_k + KV_TILE - 1`` in Int32 range. -SMEM_P4_RUNTIME_MAX_KV = ( - (INT32_MAX - (KV_TILE - 1)) // SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT -) * SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT -SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_PAGE_PLAN_PROFILE_KV // TRTLLM_PAGE_SIZE +SMEM_P4_RUNTIME_MAX_KV = 160 * 1024 +SMEM_P4_RUNTIME_MAX_PAGES = SMEM_P4_RUNTIME_MAX_KV // TRTLLM_PAGE_SIZE +SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_RUNTIME_MAX_PAGES SMEM_P4_PAGE_ID_PLAN_BYTES = SMEM_P4_PAGE_ID_PLAN_INTS * ctm.Int32.bytes SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS = 2 SMEM_P4_RUNTIME_SCALE_PAIR_BYTES = SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS * ctm.Float32.bytes @@ -503,40 +511,13 @@ def _initial_float_from_argv(option: str, default: float) -> float: _EXPLICIT_TORCH_STREAM: torch.cuda.Stream | None = None -@contextlib.contextmanager -def _launch_stream(): - """Yield the CUDA driver stream these kernels should launch on. - - cuda2ctl capture cannot reliably query the default null stream, so SMART - wrappers can ask for an explicit one through DKG_MLA_EXPLICIT_STREAM. The - substitution must be bracketed rather than assigned: the caller's stream - already carries the RoPE, Q quantization and KV-cache writes these kernels - read, and it is also the stream their consumers read the output from. - Entering waits on the caller's work, leaving publishes ours back to it, and - ``torch.cuda.stream`` restores the thread's current stream on the way out. - A bare ``set_stream`` does none of the three. - - Switching the current stream is illegal while a CUDA graph is capturing, so - the knob is ignored under capture and the caller's stream is used as-is. - """ - entry_stream = torch.cuda.current_stream() - if os.environ.get("DKG_MLA_EXPLICIT_STREAM") != "1" or torch.cuda.is_current_stream_capturing(): - yield cuda.CUstream(entry_stream.cuda_stream) - return - +def _current_cu_stream() -> cuda.CUstream: global _EXPLICIT_TORCH_STREAM - if _EXPLICIT_TORCH_STREAM is None: - _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() - start_event = torch.cuda.Event() - done_event = torch.cuda.Event() - start_event.record(entry_stream) - with torch.cuda.stream(_EXPLICIT_TORCH_STREAM): - _EXPLICIT_TORCH_STREAM.wait_event(start_event) - try: - yield cuda.CUstream(_EXPLICIT_TORCH_STREAM.cuda_stream) - finally: - done_event.record(_EXPLICIT_TORCH_STREAM) - entry_stream.wait_event(done_event) + if os.environ.get("DKG_MLA_EXPLICIT_STREAM") == "1": + if _EXPLICIT_TORCH_STREAM is None: + _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() + torch.cuda.set_stream(_EXPLICIT_TORCH_STREAM) + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) @dataclass(frozen=True) @@ -996,26 +977,6 @@ def _ceil_div(a: int, b: int) -> int: return (a + b - 1) // b -def _select_runtime_kv_profile(max_kv_len: int) -> int: - if type(max_kv_len) is not int: - raise TypeError(f"max_kv_len must be an int, got {type(max_kv_len).__name__}") - if max_kv_len <= 0: - raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") - if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: - raise ValueError( - "max_kv_len exceeds the Int32-safe runtime-KV limit: " - f"{max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" - ) - if max_kv_len <= SMEM_P4_PAGE_PLAN_PROFILE_KV: - return SMEM_P4_PAGE_PLAN_PROFILE_KV - # Eager execution can report a different batch maximum on every step. - # Shifted power-of-two buckets bound the number of compiled variants and - # retain the established 1 Mi-token plus four-page-reserve profile. - profile_payload = max_kv_len - SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT - profile_kv = (1 << (profile_payload - 1).bit_length()) + (SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT) - return min(profile_kv, SMEM_P4_RUNTIME_MAX_KV) - - def _validate_output_dtype(output_dtype: torch.dtype) -> None: if output_dtype not in {torch.float16, torch.bfloat16}: raise TypeError(f"output dtype must be torch.float16 or torch.bfloat16, got {output_dtype}") @@ -1062,7 +1023,6 @@ def fused_fp4_mla_decode_ctm( sfb_ptr: cute.Pointer, page_table_ptr: cute.Pointer, valid_k_ptr: cute.Pointer, - helix_kv_bounds_ptr: cute.Pointer, c_ptr: cute.Pointer, accum_ptr: cute.Pointer, row_max_ptr: cute.Pointer, @@ -1086,11 +1046,8 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, - use_helix_kv_bounds: ctm.Constexpr = False, - use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, - write_softmax_stats: ctm.Constexpr = False, ) -> None: n, k = problem_size m = runtime_m @@ -1117,9 +1074,10 @@ def fused_fp4_mla_decode_ctm( stride=(1, TRTLLM_K_STORAGE_DIM // 2, kv_page_stride_bytes), ), ) + page_count = cute.assume(k // page_size, 1) page_table_tensor = cute.make_tensor( cute.recast_ptr(page_table_ptr, dtype=cutlass.Int32), - cute.make_layout(((k // page_size) * (l // query_len_per_seq),), stride=(1,)), + cute.make_layout((SMEM_P4_RUNTIME_MAX_PAGES * (l // query_len_per_seq),), stride=(1,)), ) page_indptr_tensor = cute.make_tensor( cute.recast_ptr(page_indptr_ptr, dtype=cutlass.Int32), @@ -1137,10 +1095,6 @@ def fused_fp4_mla_decode_ctm( cute.recast_ptr(valid_k_ptr, dtype=cutlass.Int32), cute.make_layout((l // query_len_per_seq,), stride=(1,)), ) - helix_kv_bounds_tensor = cute.make_tensor( - cute.recast_ptr(helix_kv_bounds_ptr, dtype=cutlass.Int32), - cute.make_layout((l,), stride=(1,)), - ) v_tma_tensor = cute.make_tensor( b_ptr, cute.make_layout( @@ -1403,7 +1357,6 @@ def fused_fp4_mla_decode_ctm( page_table_tensor, page_indptr_tensor, valid_k_tensor, - helix_kv_bounds_tensor, q_global_scale_tensor, kv_global_scale_tensor, tma_q_desc, @@ -1432,12 +1385,9 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, - use_helix_kv_bounds, - use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, - write_softmax_stats, ).launch( grid=( cute.ceil_div(c_tensor.shape[0], SMEM_P4_CTA_GROUP_M) * CLUSTER_SHAPE_MNK[0], @@ -1519,7 +1469,7 @@ def _load_qk_qonly_kblock_stage( if cta_rank == ctm.Int32(0): while not prims.mbarrier_try_wait_parity(q_tma_mbar, q_tma_phase, time_limit=10000000): pass - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) @cute.jit @@ -1537,31 +1487,6 @@ def _lookup_physical_page( return ctm.Int32(mPageTable_pl[page_table_idx]) -@cute.jit -def _load_global_page_native_tile_pair( - mPageTable_pl: cute.Tensor, - kv_tile_idx: ctm.Int32, - page_begin: ctm.Int32, - page_count: ctm.Int32, -) -> tuple: - tile_page_idx = ctm.Int32(kv_tile_idx * SMEM_P4_PAGES_PER_KV_TILE) - # Keep these as two scalar loads: a CSR row may start at an odd Int32, and - # the second logical page must clamp to the first for a one-page tail. - physical_page0 = _lookup_physical_page( - mPageTable_pl, tile_page_idx, ctm.Int32(0), page_begin, page_count - ) - physical_page1 = _lookup_physical_page( - mPageTable_pl, - tile_page_idx + ctm.Int32(1), - ctm.Int32(0), - page_begin, - page_count, - ) - physical_page0 = cute.arch.make_warp_uniform(physical_page0) - physical_page1 = cute.arch.make_warp_uniform(physical_page1) - return (physical_page0, physical_page1) - - @cute.jit def _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx: ctm.Int32) -> tuple: physical_page0 = ctm.Int32(0) @@ -1590,13 +1515,10 @@ def _tma_gather4_cluster( """Gather four KSF atoms and multicast the compact image to both CTAs.""" smem_ptr = smem_dst.data_ptr() tma_ptr = tma_desc.data_ptr() if hasattr(tma_desc, "data_ptr") else tma_desc - leader_barrier = prims.cvta_to( - _mapa_shared_cluster(barrier, ctm.Int32(0)), - prims.CvtaSpace.SHARED, - ) + leader_barrier = cvta_to(_mapa_shared_cluster(barrier, ctm.Int32(0)), CvtaSpace.SHARED) barrier_ptr = leader_barrier.data_ptr() multicast_mask_u16 = ctm.Uint16(multicast_mask) - cute_inline_ptx( + _cp_async._predicated_inline_ptx( "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" " [{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], " @@ -2312,6 +2234,7 @@ def _issue_runtime_t336_qk_prefix_rank1( @cute.jit def _load_v_tile_stage( mPageTable_pl: cute.Tensor, + sPageIdPlan, tma_v_tile_ptr, tma_v_pair_ptr, tma_vsf_ptr, @@ -2346,11 +2269,15 @@ def _load_v_tile_stage( if should_wait_v_tma: if prims.elect_sync(): prims.mbarrier_arrive_expect_tx(v_tma_mbar, v_tma_bytes) - # Resolve the supplied pair at the point of use so V-ring wrap cannot - # reuse stale loop-carried IDs. + # Resolve the page pair at the point of use so V-ring wrap cannot reuse a + # stale loop-carried ID. The plan was populated once before the producer + # warps start, so reloading it here avoids a second CSR/global lookup and + # its serial uniform-address chain on every tile. del mPageTable_pl, bidz, page_begin, page_count - resolved_physical_page0 = tile_physical_page0 - resolved_physical_page1 = tile_physical_page1 + del tile_physical_page0, tile_physical_page1 + resolved_physical_page0, resolved_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) if cutlass.const_expr(use_consecutive_page_pair): resolved_physical_page1 = resolved_physical_page0 + ctm.Int32(1) if prims.elect_sync(): @@ -2438,7 +2365,7 @@ def _load_v_tile_stage( if should_wait_v_tma: while not prims.mbarrier_try_wait_parity(v_tma_mbar, v_tma_phase, time_limit=10000000): pass - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) @cute.jit @@ -2446,20 +2373,15 @@ def _issue_raw_k_pages_to_v_staging( tma_k_v_raw_ptr, raw_v_tma_mbar, sRawV, + sV, sPageIdPlan, + tidx: ctm.Int32, cta_rank: ctm.Int32, kv_tile_idx: ctm.Int32, - direct_physical_page0: ctm.Int32, - direct_physical_page1: ctm.Int32, - use_smem_page_plan: ctm.Constexpr, + v_stage: ctm.Int32, ) -> None: """Stage the V-bearing K prefix in a transpose-friendly layout.""" - physical_page0 = direct_physical_page0 - physical_page1 = direct_physical_page1 - if cutlass.const_expr(use_smem_page_plan): - physical_page0, physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) + physical_page0, physical_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx) raw_packed_dims_per_n_tile: ctm.Constexpr = SMEM_P4_V_N_PER_CTA // 2 raw_n_tile_bytes: ctm.Constexpr = TRTLLM_PAGE_SIZE * raw_packed_dims_per_n_tile raw_page_bytes: ctm.Constexpr = raw_n_tile_bytes * SMEM_P4_N_OUT_TILES @@ -2476,7 +2398,7 @@ def _issue_raw_k_pages_to_v_staging( raw_stage_offset = ctm.Int32( logical_page_idx * raw_page_bytes + n_tile_idx * raw_n_tile_bytes ) - prims.cp_async_bulk_tensor_shared_cta_global( + _cp_async.cp_async_bulk_tensor_shared_cta_global( sRawV.subview(raw_stage_offset), tma_k_v_raw_ptr, (packed_dim_begin, ctm.Int32(0), physical_page), @@ -2515,7 +2437,7 @@ def _transpose_raw_k_staging_to_v_stage( ) while not prims.mbarrier_try_wait_parity(raw_v_tma_mbar, raw_phase, time_limit=10000000): pass - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) raw_ptr = sRawV.data_ptr() v_word_ptr = sV.data_ptr() @@ -2633,7 +2555,7 @@ def _transpose_raw_k_staging_to_v_stage( ).bitcast(cutlass.Uint8) raw_fragment.store(packed_fragment) cute.copy(stsm_atom, raw_fragment, pv_lane) - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) prims.barrier( barrier_id=SMEM_P4_RAW_V_READY_BAR_ID, number_of_threads=SMEM_P4_RAW_V_READY_BAR_THREADS, @@ -2649,20 +2571,12 @@ def _load_vsf_tile_stage_only( sVSF, cta_rank: ctm.Int32, kv_tile_idx: ctm.Int32, - direct_physical_page0: ctm.Int32, - direct_physical_page1: ctm.Int32, v_page_offset: ctm.Int32, stage: ctm.Int32, - use_smem_page_plan: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr = False, ) -> None: sSFB_stage = sVSF.subview(SMEM_P4_V_SFB_TMA_STAGE_BYTES * stage) - physical_page0 = direct_physical_page0 - physical_page1 = direct_physical_page1 - if cutlass.const_expr(use_smem_page_plan): - physical_page0, physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) + physical_page0, physical_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx) if cutlass.const_expr(use_consecutive_page_pair): physical_page1 = physical_page0 + ctm.Int32(1) if cta_rank == ctm.Int32(0): @@ -2705,7 +2619,7 @@ def _load_vsf_tile_stage_only( @cute.jit def _fence_async_shared_cta() -> None: - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) @cute.jit @@ -2832,11 +2746,7 @@ def _softmax_exp2_pair_packed_f16x2( (score0, score1), (softmax_scale_log2, softmax_scale_log2), (p_bias, p_bias) ) shifted_h2 = _pack_f32_pair_to_f16x2(shifted[0], shifted[1]) - return cute_inline_ptx( - "ex2.approx.f16x2 {$w0}, {$r0};", - write_only_types=[ctm.Int32], - read_only_args=[shifted_h2], - ) + return ptx.ex2_f16x2(shifted_h2) @cute.jit @@ -4119,28 +4029,28 @@ def _smem_p4_p4_materialize_group_col( @cute.jit def _float_to_ordered_u32_for_atomic_max(value: ctm.Float32) -> ctm.Uint32: - bits = prims.mov_b32(value, target_type=ctm.Int32) + bits = ptx.mov_b32(value, target_type=ctm.Int32) sign_mask = bits >> ctm.Int32(31) | ctm.Int32(2147483648) encoded = bits ^ sign_mask - return prims.mov_b32(encoded, target_type=ctm.Uint32) + return ptx.mov_b32(encoded, target_type=ctm.Uint32) @cute.jit def _ordered_u32_to_float_after_atomic_max(value: ctm.Uint32) -> ctm.Float32: - encoded = prims.mov_b32(value, target_type=ctm.Int32) + encoded = ptx.mov_b32(value, target_type=ctm.Int32) sign_mask = ~(encoded >> ctm.Int32(31)) | ctm.Int32(2147483648) bits = encoded ^ sign_mask - return prims.mov_b32(bits, target_type=ctm.Float32) + return ptx.mov_b32(bits, target_type=ctm.Float32) @cute.jit def _smem_atomic_max_ordered_u32(pointer, value: ctm.Uint32) -> None: - prims.atomicrmw( - prims.AtomicOp.MAX, + ptx.atom( + AtomicOpKind.MAX, pointer, value, - syncscope=prims.MemScope.CTA, - space=prims.SharedSpace.shared_cta, + syncscope=MemScopeKind.CTA, + space=SharedSpace.shared_cta, ) @@ -4189,19 +4099,6 @@ def _prepare_p4_n256_score_half_tmem_addresses( return tuple(score_tmem_addresses) -@cute.jit -def _tcgen05_ld_red_x16_max_f32(tmem) -> tuple: - tmem_addr = tmem.toint(ctm.Int32) - return cute_inline_ptx( - "tcgen05.ld.red.sync.aligned.32x32b.x16.max.f32 " - "{{$w0}, {$w1}, {$w2}, {$w3}, {$w4}, {$w5}, {$w6}, {$w7}, " - "{$w8}, {$w9}, {$w10}, {$w11}, {$w12}, {$w13}, {$w14}, {$w15}}, " - "{$w16}, [{$r0}];", - write_only_types=[ctm.Int32] * 17, - read_only_args=[tmem_addr], - ) - - @cute.jit def _load_p4_n256_score_half_from_tmem( sAtomicRunningRowMax, @@ -4239,7 +4136,13 @@ def _load_p4_n256_score_half_from_tmem( pending_score_groups.append(group_scores) pending_group_stats.append(None) else: - regs = _tcgen05_ld_red_x16_max_f32(tmem) + regs = ptx.tcgen05_ld_red( + ptx.Tcgen05LdStShape.SHAPE_32X32B, + tmem, + num=SF_VEC_SIZE, + red_op="max", + type_="f32", + ) pending_score_groups.append(regs) pending_group_stats.append(regs[SF_VEC_SIZE]) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) @@ -4704,7 +4607,7 @@ def _finalize_p4_n256_score_half_psf_source( ) sPSF[psf_source_word_offset] = pv_sf_word cute.nvgpu.cfence() - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) cute.nvgpu.cfence() return (denominator_sf_word, prequant_row_sum) @@ -4816,7 +4719,7 @@ def _runtime_finalize_final_p4_n256_score_half_psf_source( ) sPSF[psf_source_word_offset] = pv_sf_word cute.nvgpu.cfence() - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) cute.nvgpu.cfence() return (denominator_sf_word, prequant_row_sum) @@ -5074,7 +4977,6 @@ def _store_final_o_from_tmem( final_row_sum: ctm.Float32, final_stat_scale: ctm.Float32, output_normalizer: ctm.Float32, - write_softmax_stats: ctm.Constexpr = False, producer_warp_base: ctm.Constexpr = 0, ) -> None: gC_arr = ctm.make_array_view(mC_mnl) @@ -5096,10 +4998,6 @@ def _store_final_o_from_tmem( + local_row ) batch_offset = bidz * m * n - if cutlass.const_expr(write_softmax_stats): - if col_band == ctm.Int32(0) and bidy == ctm.Int32(0): - mRowMax_ml[row, bidz] = final_row_max * final_stat_scale - mRowSum_ml[row, bidz] = final_row_sum n_tile_total = n // ctm.Int32(OUT_DIM) subtile_cols: ctm.Constexpr = 32 subtiles_per_n_tile: ctm.Constexpr = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N // subtile_cols @@ -5967,13 +5865,14 @@ def _runtime_t336_producer_tile( ) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - if common_source_ready_pred: - prims.mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - count=1, - scope=prims.MemScope.CLUSTER, - relaxed=True, - ) + ptx.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + 1, + sem=MBarrierArriveSem.RELAXED, + scope=MBarrierArriveScope.CLUSTER, + space=MBarrierSpace.SHARED_CLUSTER, + pred=common_source_ready_pred, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -6240,13 +6139,14 @@ def _runtime_t336_producer_tile_v23( warp_rebase = cute.arch.vote_any_sync(lane_rebase) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - if common_source_ready_pred: - prims.mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - count=1, - scope=prims.MemScope.CLUSTER, - relaxed=True, - ) + ptx.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + 1, + sem=MBarrierArriveSem.RELAXED, + scope=MBarrierArriveScope.CLUSTER, + space=MBarrierSpace.SHARED_CLUSTER, + pred=common_source_ready_pred, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -6383,7 +6283,6 @@ def _run_mla_decode_body( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, - mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_ptr, @@ -6412,12 +6311,9 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, - use_helix_kv_bounds: ctm.Constexpr, - use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, - write_softmax_stats: ctm.Constexpr = False, ) -> None: pv_psf_rescale = ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) * pv_output_scale warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -6428,14 +6324,11 @@ def _run_mla_decode_body( page_batch = ctm.Int32(0) valid_k_for_l = ctm.Int32(0) page_batch = cute.arch.make_warp_uniform(bidz // ctm.Int32(query_len_per_seq)) - if cutlass.const_expr(use_helix_kv_bounds): - valid_k_for_l = mHelixKvBounds_l[bidz] - else: - query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) - valid_k_for_l = ctm.max( - mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), - ctm.Int32(0), - ) + query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) + valid_k_for_l = ctm.max( + mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), + ctm.Int32(0), + ) valid_k_for_l = cute.arch.make_warp_uniform(valid_k_for_l) csr_page_begin = ctm.Int32(0) csr_page_count = ctm.Int32(0) @@ -6735,10 +6628,9 @@ def _run_mla_decode_body( _stage_smem_p4_runtime_scale_pair( mQGlobalScale, mKvGlobalScale, sRuntimeScalePair, softmax_scale_log2 ) - if cutlass.const_expr(use_smem_page_plan): - _stage_smem_p4_page_id_plan( - mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count - ) + _stage_smem_p4_page_id_plan( + mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count + ) cute.arch.cluster_arrive() cute.arch.cluster_wait() runtime_scale_pair = sRuntimeScalePair.data_ptr().load(count=2, alignment=8) @@ -6788,14 +6680,7 @@ def _run_mla_decode_body( qk0_handle = qk_smem_producer.acquire_and_advance() qk0_tma_mbar = ctm.Array(qk0_handle.barrier, shape=1) if is_leader_cta: - if cutlass.const_expr(use_smem_page_plan): - qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(0) - ) - else: - qk0_page0, qk0_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, ctm.Int32(0), csr_page_begin, csr_page_count - ) + qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, ctm.Int32(0)) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6829,14 +6714,9 @@ def _run_mla_decode_body( ) qk_prefix_tma_mbar = ctm.Array(qk_prefix_handle.barrier, shape=1) if is_leader_cta: - if cutlass.const_expr(use_smem_page_plan): - qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) - else: - qk_prefix_page0, qk_prefix_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count - ) + qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6872,14 +6752,9 @@ def _run_mla_decode_body( expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] ) qk_tma_mbar = ctm.Array(qk_handle.barrier, shape=1) - if cutlass.const_expr(use_smem_page_plan): - qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) - else: - qk_steady_page0, qk_steady_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count - ) + qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) _load_runtime_t336_qk_tile_dual_steady( mPageTable_pl, tma_k_page_ptr, @@ -6912,14 +6787,9 @@ def _run_mla_decode_body( ) qk15_tma_mbar = ctm.Array(qk15_handle.barrier, shape=1) if is_leader_cta: - if cutlass.const_expr(use_smem_page_plan): - qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(15) - ) - else: - qk15_page0, qk15_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, ctm.Int32(15), csr_page_begin, csr_page_count - ) + qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(15) + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6954,22 +6824,16 @@ def _run_mla_decode_body( v_tma_mbar = ctm.Array(v_handle.barrier, shape=1) v_stage = kv_tile_idx % ctm.Int32(SMEM_P4_V_PIPELINE_STAGES) raw_v_stage_mbar = raw_v_tma_mbar.subview(v_stage) - v_physical_page0 = ctm.Int32(0) - v_physical_page1 = ctm.Int32(0) - if cutlass.const_expr(not use_smem_page_plan): - v_physical_page0, v_physical_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count - ) _issue_raw_k_pages_to_v_staging( tma_k_v_raw_ptr, raw_v_stage_mbar, sRawV, + sV, sPageIdPlan, + tidx, cta_rank, kv_tile_idx, - v_physical_page0, - v_physical_page1, - use_smem_page_plan, + v_stage, ) prims.barrier( barrier_id=SMEM_P4_RAW_V_READY_BAR_ID, @@ -6977,7 +6841,7 @@ def _run_mla_decode_body( ) prims.fence_proxy( kind=prims.Proxy.ASYNC_SHARED, - space=prims.SharedSpace.shared_cta, + space=SharedSpace.shared_cta, ) _load_vsf_tile_stage_only( sPageIdPlan, @@ -6987,11 +6851,8 @@ def _run_mla_decode_body( sVSF, cta_rank, kv_tile_idx, - v_physical_page0, - v_physical_page1, v_page_offset, v_stage, - use_smem_page_plan, use_consecutive_page_pair=use_consecutive_page_pair, ) v_smem_producer.tail() @@ -6999,14 +6860,9 @@ def _run_mla_decode_body( if is_leader_cta: qk_prefix_end = min(kv_tiles, ctm.Int32(3)) for qk_prefix_li_idx in cutlass.range(1, qk_prefix_end, 1, unroll=1): - if cutlass.const_expr(use_smem_page_plan): - rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) - else: - rowmeta_page0, rowmeta_page1 = _load_global_page_native_tile_pair( - mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count - ) + rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) _issue_runtime_t336_qk_prefix_rank1( mPageTable_pl, tma_k_page_ptr, @@ -7424,19 +7280,8 @@ def _run_mla_decode_body( final_anchor_row_sum = sFinalAnchorRowSum[final_row_state_local_row] final_stat_scale = ctm.Float32(1.0) final_row_sum = final_anchor_row_sum - final_row_max = running_row_max - if cutlass.const_expr(write_softmax_stats): - final_stat_scale = softmax_scale_log2 / ctm.Float32(LOG2_E) - final_row_sum = final_anchor_row_sum * cute.exp2( - (running_row_anchor - running_row_max) * softmax_scale_log2, - fastmath=True, - ) - if valid_k_for_l == ctm.Int32(0): - final_stat_scale = ctm.Float32(1.0) - final_row_max = ctm.Float32(-ctm.Float32.inf) - final_row_sum = ctm.Float32(0.0) output_normalizer = ctm.Float32(0.0) - if valid_k_for_l != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): + if final_anchor_row_sum != ctm.Float32(0.0): output_normalizer = cute.arch.rcp_approx(final_anchor_row_sum) * pv_output_scale last_pv_li_idx = stream_li_total - ctm.Int32(1) last_pv_slot = last_pv_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -7462,11 +7307,10 @@ def _run_mla_decode_body( bidz, m, n, - final_row_max=final_row_max, + final_row_max=running_row_max, final_row_sum=final_row_sum, final_stat_scale=final_stat_scale, output_normalizer=output_normalizer, - write_softmax_stats=write_softmax_stats, producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, ) prims.barrier(barrier_id=O_STORE_BAR_ID, number_of_threads=O_STORE_BAR_THREADS) @@ -7485,7 +7329,6 @@ def kernel( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, - mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_desc: ctm.GridConstant[cuda_tma.TensorMap], @@ -7514,18 +7357,14 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, - use_helix_kv_bounds: ctm.Constexpr, - use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, - write_softmax_stats: ctm.Constexpr, ) -> None: _run_mla_decode_body( mPageTable_pl, mPageIndptr_s, mValidK_l, - mHelixKvBounds_l, mQGlobalScale, mKvGlobalScale, tma_q_desc.get_ptr(), @@ -7554,12 +7393,9 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, - use_helix_kv_bounds, - use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, - write_softmax_stats, ) @@ -7580,7 +7416,6 @@ def _make_fused_ptrs( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, - helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7600,7 +7435,6 @@ def _make_fused_ptrs( make_ptr(cutlass.Uint8, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=32), make_ptr(cutlass.Int32, page_table_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr(cutlass.Int32, valid_k_data_ptr, cute.AddressSpace.gmem, assumed_align=4), - make_ptr(cutlass.Int32, helix_kv_bounds_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr( _cutlass_output_dtype(output_dtype), c_data_ptr, @@ -7736,7 +7570,6 @@ def _compile_fused( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, - helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7758,25 +7591,13 @@ def _compile_fused( ksf_page_stride_bytes: int = 0, vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, - write_softmax_stats: bool = False, - use_helix_kv_bounds: bool = False, ) -> Callable: - if type(kv) is not int: - raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") - is_bucketed_runtime_profile = ( - SMEM_P4_PAGE_PLAN_PROFILE_KV < kv <= SMEM_P4_RUNTIME_MAX_KV - and _select_runtime_kv_profile(kv) == kv - ) - if kv != SMEM_P4_PAGE_PLAN_PROFILE_KV and not is_bucketed_runtime_profile: - raise ValueError( - "runtime-KV compile requires the fixed page-plan profile or a geometric runtime " - f"profile in ({SMEM_P4_PAGE_PLAN_PROFILE_KV}, {SMEM_P4_RUNTIME_MAX_KV}], got {kv}" - ) + if kv != SMEM_P4_RUNTIME_MAX_KV: + raise ValueError(f"runtime-KV compile requires fixed K={SMEM_P4_RUNTIME_MAX_KV}, got {kv}") use_ksf_gather4 = ( not use_consecutive_page_pair and ksf_page_stride_bytes == page_size * TRTLLM_K_SF_GROUPS ) cache_key = ( - kv, n, page_size, use_mixed_imlp, @@ -7784,8 +7605,6 @@ def _compile_fused( query_len_per_seq, use_consecutive_page_pair, use_ksf_gather4, - write_softmax_stats, - use_helix_kv_bounds, ) cached = _FUSED_COMPILE_CACHE.get(cache_key) if cached is not None: @@ -7800,7 +7619,6 @@ def _compile_fused( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, - helix_kv_bounds_data_ptr, c_data_ptr, accum_data_ptr, row_max_data_ptr, @@ -7829,11 +7647,8 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, - use_helix_kv_bounds=use_helix_kv_bounds, - use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, - write_softmax_stats=write_softmax_stats, options="--opt-level 2 --ptxas-options '--uumn'", ) _FUSED_COMPILE_CACHE[cache_key] = compiled @@ -7865,9 +7680,6 @@ def run_trtllm_fp4_mla_decode_page_native( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, - softmax_row_max: torch.Tensor | None = None, - softmax_row_sum: torch.Tensor | None = None, - helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7882,7 +7694,12 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"page-native decode requires page_size={TRTLLM_PAGE_SIZE}, got {page_size}" ) - physical_k = _select_runtime_kv_profile(max_kv_len) + if max_kv_len <= 0: + raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") + if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: + raise ValueError( + f"max_kv_len exceeds the fixed runtime-KV profile: {max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" + ) for name, value in ( ("assume_valid_k_prefix_tiles", assume_valid_k_prefix_tiles), ( @@ -7908,29 +7725,12 @@ def run_trtllm_fp4_mla_decode_page_native( del assume_valid_k_prefix_tiles, partition_runtime_valid_k required_consecutive_tiles = _ceil_div(max_kv_len, KV_TILE) use_consecutive_page_pair = assume_consecutive_page_prefix_tiles >= required_consecutive_tiles + physical_k = SMEM_P4_RUNTIME_MAX_KV if q_internal.dtype != torch.uint8 or q_internal.dim() != 3: raise TypeError( f"q_internal must be a uint8 [M, Q640/2, L] tensor, got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" ) physical_m, q_bytes, l_batch = q_internal.shape - if (softmax_row_max is None) != (softmax_row_sum is None): - raise ValueError("softmax_row_max and softmax_row_sum must be provided together") - write_softmax_stats = softmax_row_max is not None - if write_softmax_stats: - expected_stats_shape = (l_batch, physical_m) - for name, tensor in ( - ("softmax_row_max", softmax_row_max), - ("softmax_row_sum", softmax_row_sum), - ): - if ( - tensor.dtype != torch.float32 - or tensor.shape != expected_stats_shape - or not tensor.is_contiguous() - ): - raise ValueError( - f"{name} must be contiguous float32 with shape {expected_stats_shape}" - ) - _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) if l_batch <= 0 or l_batch > CUDA_GRID_Z_MAX: raise ValueError(f"queries must be in [1, {CUDA_GRID_Z_MAX}], got {l_batch}") if q_batch_capacity is None: @@ -7982,14 +7782,9 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"src_page_ids must be a 1D int32 physical-page list, got dtype={src_page_ids.dtype} shape={tuple(src_page_ids.shape)}" ) - src_page_id_count = src_page_ids.numel() - if src_page_id_count == 0 or src_page_ids.stride(0) != 1: - raise ValueError( - f"src_page_ids must be non-empty and contiguous, got numel={src_page_id_count} stride={src_page_ids.stride()}" - ) - if src_page_id_count > INT32_MAX: + if src_page_ids.numel() == 0 or src_page_ids.stride(0) != 1: raise ValueError( - f"src_page_ids exceeds the Int32 element limit: {src_page_id_count} > {INT32_MAX}" + f"src_page_ids must be non-empty and contiguous, got numel={src_page_ids.numel()} stride={src_page_ids.stride()}" ) if ( paged_kv_indptr_decode.dtype != torch.int32 @@ -8003,16 +7798,6 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"valid_k must be contiguous int32 [{num_sequences}], got dtype={valid_k.dtype} shape={tuple(valid_k.shape)} stride={valid_k.stride()}" ) - if helix_kv_bounds is not None and ( - helix_kv_bounds.dtype != torch.int32 - or helix_kv_bounds.shape != (l_batch,) - or helix_kv_bounds.stride(0) != 1 - ): - raise ValueError( - "helix_kv_bounds must be contiguous int32 with shape " - f"[{l_batch}], got dtype={helix_kv_bounds.dtype} " - f"shape={tuple(helix_kv_bounds.shape)} stride={helix_kv_bounds.stride()}" - ) cache_layout = _kv_cache_3d_layout(kv_cache, page_size) _validate_tensor_pointer_alignment("kv_cache", kv_cache, alignment_bytes=16) if ( @@ -8027,25 +7812,21 @@ def run_trtllm_fp4_mla_decode_page_native( ) num_cache_pages = cache_layout.num_pages page_table_capacity = num_sequences * (physical_k // page_size) - if page_table_capacity > INT32_MAX: - raise ValueError( - "bucketed CSR capacity exceeds the Int32 layout limit: " - f"{page_table_capacity} > {INT32_MAX}" - ) - if src_page_id_count > page_table_capacity: + if src_page_ids.numel() > page_table_capacity: raise ValueError( - f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_id_count} > {page_table_capacity}" + f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_ids.numel()} > {page_table_capacity}" ) if not isinstance(v_page_offset, int): raise TypeError(f"v_page_offset must be an int, got {type(v_page_offset).__name__}") if v_sf.dim() == 0: raise ValueError("v_sf must expose a physical-page dimension") num_v_cache_pages = int(v_sf.shape[0]) - if v_page_offset < 0 or v_page_offset > INT32_MAX: - raise ValueError(f"v_page_offset must be in [0, {INT32_MAX}], got {v_page_offset}") - if num_v_cache_pages > INT32_MAX: + int32_max = torch.iinfo(torch.int32).max + if v_page_offset < 0 or v_page_offset > int32_max: + raise ValueError(f"v_page_offset must be in [0, {int32_max}], got {v_page_offset}") + if num_v_cache_pages > int32_max: raise ValueError( - f"v_sf exceeds the Int32 physical-page limit: {num_v_cache_pages} > {INT32_MAX}" + f"v_sf exceeds the Int32 physical-page limit: {num_v_cache_pages} > {int32_max}" ) if v_page_offset + num_cache_pages > num_v_cache_pages: raise ValueError( @@ -8104,189 +7885,172 @@ def run_trtllm_fp4_mla_decode_page_native( q_global_scale, kv_global_scale, ) - if write_softmax_stats: - tensors += (softmax_row_max, softmax_row_sum) - if helix_kv_bounds is not None: - tensors += (helix_kv_bounds,) if device.type != "cuda" or any((tensor.device != device for tensor in tensors)): raise ValueError("all page-native decode tensors must share one CUDA device") if q_global_scale.dtype != torch.float32 or q_global_scale.numel() != 1: raise TypeError("q_global_scale must be a scalar FP32 CUDA tensor") if kv_global_scale.dtype != torch.float32 or kv_global_scale.numel() != 1: raise TypeError("kv_global_scale must be a scalar FP32 CUDA tensor") - with _launch_stream() as stream: - device_index = torch.cuda.current_device() - context_result, current_context = cuda.cuCtxGetCurrent() - if context_result != cuda.CUresult.CUDA_SUCCESS: + stream = _current_cu_stream() + device_index = torch.cuda.current_device() + context_result, current_context = cuda.cuCtxGetCurrent() + if context_result != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError( + f"Failed to query the current CUDA context for FP4 MLA launch: {context_result}." + ) + q_data_ptr = q_internal.data_ptr() + k_data_ptr = kv_cache.data_ptr() + q_sf_data_ptr = q_sf_internal.data_ptr() + k_sf_data_ptr = sf_cache.data_ptr() + # Legacy V tensor-map operands are dead in this specialization. Reuse the + # canonical KV pointer so the compiled call carries no sidecar allocation. + b_data_ptr = kv_cache.data_ptr() + scratch_ptr = output.data_ptr() + sfb_data_ptr = v_sf.data_ptr() + page_table_data_ptr = src_page_ids.data_ptr() + valid_k_data_ptr = valid_k.data_ptr() + c_data_ptr = output.data_ptr() + page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() + q_global_scale_data_ptr = q_global_scale.data_ptr() + kv_global_scale_data_ptr = kv_global_scale.data_ptr() + kv_page_stride_bytes = cache_layout.stride_page + ksf_page_stride_bytes = int(sf_cache.stride(0)) + vsf_page_stride_bytes = int(v_sf.stride(0)) + softmax_scale_log2 = sm_scale * LOG2_E + fused = _compile_fused( + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + scratch_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + c_data_ptr, + scratch_ptr, + scratch_ptr, + scratch_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + physical_m, + TRTLLM_V_HEAD_DIM, + physical_k, + l_batch, + page_size, + use_mixed_imlp=enable_mxi_imlp, + output_dtype=output.dtype, + stream=stream, + num_cache_pages=num_cache_pages, + query_len_per_seq=query_len_per_seq, + kv_page_stride_bytes=kv_page_stride_bytes, + ksf_page_stride_bytes=ksf_page_stride_bytes, + vsf_page_stride_bytes=vsf_page_stride_bytes, + use_consecutive_page_pair=use_consecutive_page_pair, + ) + supports_prepared = _class_defines_callables( + fused, "to", "generate_execution_args", "run_compiled_program" + ) + prepared_key = ( + fused, + device_index, + int(current_context), + int(stream), + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + c_data_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + physical_m, + l_batch, + q_batch_capacity, + num_cache_pages, + num_v_cache_pages, + v_page_offset, + kv_page_stride_bytes, + ksf_page_stride_bytes, + vsf_page_stride_bytes, + softmax_scale_log2, + output.dtype, + enable_mxi_imlp, + use_consecutive_page_pair, + ) + if supports_prepared: + prepared = _get_prepared_fused_call(prepared_key) + if prepared is not None: + prepared.run() + return + ptrs = _make_fused_ptrs( + q_data_ptr, + k_data_ptr, + q_sf_data_ptr, + k_sf_data_ptr, + b_data_ptr, + scratch_ptr, + sfb_data_ptr, + page_table_data_ptr, + valid_k_data_ptr, + c_data_ptr, + scratch_ptr, + scratch_ptr, + scratch_ptr, + page_indptr_data_ptr, + q_global_scale_data_ptr, + kv_global_scale_data_ptr, + output.dtype, + ) + runtime_args = ( + *ptrs, + (TRTLLM_V_HEAD_DIM, physical_k), + ctm.Int32(physical_m), + ctm.Int32(l_batch), + ctm.Int32(q_batch_capacity), + ctm.Int32(num_cache_pages), + ctm.Int32(num_v_cache_pages), + ctm.Int32(v_page_offset), + ctm.Int64(kv_page_stride_bytes), + ctm.Int64(ksf_page_stride_bytes), + ctm.Int64(vsf_page_stride_bytes), + ctm.Float32(softmax_scale_log2), + ctm.Float32(1.0), + stream, + ) + if not supports_prepared: + fused(*runtime_args) + return + executor_key = (fused, device_index, int(current_context)) + executor = _get_fused_executor(executor_key) + if executor is None: + if torch.cuda.is_current_stream_capturing(): raise RuntimeError( - f"Failed to query the current CUDA context for FP4 MLA launch: {context_result}." + "FP4 MLA CUDA Graph capture requires an eager warmup for the compiled kernel, device, and CUDA context." ) - q_data_ptr = q_internal.data_ptr() - k_data_ptr = kv_cache.data_ptr() - q_sf_data_ptr = q_sf_internal.data_ptr() - k_sf_data_ptr = sf_cache.data_ptr() - # Legacy V tensor-map operands are dead in this specialization. Reuse the - # canonical KV pointer so the compiled call carries no sidecar allocation. - b_data_ptr = kv_cache.data_ptr() - scratch_ptr = output.data_ptr() - row_max_data_ptr = softmax_row_max.data_ptr() if write_softmax_stats else scratch_ptr - row_sum_data_ptr = softmax_row_sum.data_ptr() if write_softmax_stats else scratch_ptr - sfb_data_ptr = v_sf.data_ptr() - page_table_data_ptr = src_page_ids.data_ptr() - valid_k_data_ptr = valid_k.data_ptr() - helix_kv_bounds_data_ptr = ( - helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else valid_k_data_ptr - ) - c_data_ptr = output.data_ptr() - page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() - q_global_scale_data_ptr = q_global_scale.data_ptr() - kv_global_scale_data_ptr = kv_global_scale.data_ptr() - kv_page_stride_bytes = cache_layout.stride_page - ksf_page_stride_bytes = int(sf_cache.stride(0)) - vsf_page_stride_bytes = int(v_sf.stride(0)) - softmax_scale_log2 = sm_scale * LOG2_E - fused = _compile_fused( - q_data_ptr, - k_data_ptr, - q_sf_data_ptr, - k_sf_data_ptr, - b_data_ptr, - scratch_ptr, - sfb_data_ptr, - page_table_data_ptr, - valid_k_data_ptr, - helix_kv_bounds_data_ptr, - c_data_ptr, - scratch_ptr, - row_max_data_ptr, - row_sum_data_ptr, - page_indptr_data_ptr, - q_global_scale_data_ptr, - kv_global_scale_data_ptr, - physical_m, - TRTLLM_V_HEAD_DIM, - physical_k, - l_batch, - page_size, - use_mixed_imlp=enable_mxi_imlp, - output_dtype=output.dtype, - stream=stream, - num_cache_pages=num_cache_pages, - query_len_per_seq=query_len_per_seq, - kv_page_stride_bytes=kv_page_stride_bytes, - ksf_page_stride_bytes=ksf_page_stride_bytes, - vsf_page_stride_bytes=vsf_page_stride_bytes, - use_consecutive_page_pair=use_consecutive_page_pair, - write_softmax_stats=write_softmax_stats, - use_helix_kv_bounds=helix_kv_bounds is not None, - ) - supports_prepared = _class_defines_callables( - fused, "to", "generate_execution_args", "run_compiled_program" - ) - prepared_key = ( - fused, - device_index, - int(current_context), - int(stream), - q_data_ptr, - k_data_ptr, - q_sf_data_ptr, - k_sf_data_ptr, - b_data_ptr, - sfb_data_ptr, - page_table_data_ptr, - valid_k_data_ptr, - helix_kv_bounds_data_ptr, - c_data_ptr, - row_max_data_ptr, - row_sum_data_ptr, - page_indptr_data_ptr, - q_global_scale_data_ptr, - kv_global_scale_data_ptr, - physical_m, - l_batch, - q_batch_capacity, - num_cache_pages, - num_v_cache_pages, - v_page_offset, - kv_page_stride_bytes, - ksf_page_stride_bytes, - vsf_page_stride_bytes, - softmax_scale_log2, - output.dtype, - enable_mxi_imlp, - use_consecutive_page_pair, - ) - if supports_prepared: - prepared = _get_prepared_fused_call(prepared_key) - if prepared is not None: - prepared.run() - return - ptrs = _make_fused_ptrs( - q_data_ptr, - k_data_ptr, - q_sf_data_ptr, - k_sf_data_ptr, - b_data_ptr, - scratch_ptr, - sfb_data_ptr, - page_table_data_ptr, - valid_k_data_ptr, - helix_kv_bounds_data_ptr, - c_data_ptr, - scratch_ptr, - row_max_data_ptr, - row_sum_data_ptr, - page_indptr_data_ptr, - q_global_scale_data_ptr, - kv_global_scale_data_ptr, - output.dtype, - ) - runtime_args = ( - *ptrs, - (TRTLLM_V_HEAD_DIM, physical_k), - ctm.Int32(physical_m), - ctm.Int32(l_batch), - ctm.Int32(q_batch_capacity), - ctm.Int32(num_cache_pages), - ctm.Int32(num_v_cache_pages), - ctm.Int32(v_page_offset), - ctm.Int64(kv_page_stride_bytes), - ctm.Int64(ksf_page_stride_bytes), - ctm.Int64(vsf_page_stride_bytes), - ctm.Float32(softmax_scale_log2), - ctm.Float32(1.0), - stream, - ) - if not supports_prepared: + candidate = fused.to(device_index) + if not _class_defines_callables( + candidate, "generate_execution_args", "run_compiled_program" + ): fused(*runtime_args) return - executor_key = (fused, device_index, int(current_context)) - executor = _get_fused_executor(executor_key) - if executor is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "FP4 MLA CUDA Graph capture requires an eager warmup for the " - "compiled kernel, device, and CUDA context." - ) - candidate = fused.to(device_index) - if not _class_defines_callables( - candidate, "generate_execution_args", "run_compiled_program" - ): - fused(*runtime_args) - return - executor = _cache_fused_executor(executor_key, candidate) - execution_args, adapted_args = executor.generate_execution_args(*runtime_args) - prepared = _cache_prepared_fused_call( - prepared_key, - _PreparedFusedCall( - executor=executor, - runtime_args=runtime_args, - execution_args=execution_args, - adapted_args=adapted_args, - ), - ) - prepared.run() + executor = _cache_fused_executor(executor_key, candidate) + execution_args, adapted_args = executor.generate_execution_args(*runtime_args) + prepared = _cache_prepared_fused_call( + prepared_key, + _PreparedFusedCall( + executor=executor, + runtime_args=runtime_args, + execution_args=execution_args, + adapted_args=adapted_args, + ), + ) + prepared.run() def run_trtllm_fp4_mla_decode_page_native_from_raw( @@ -8315,9 +8079,6 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, - softmax_row_max: torch.Tensor | None = None, - softmax_row_sum: torch.Tensor | None = None, - helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -8369,7 +8130,4 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles=assume_consecutive_page_prefix_tiles, partition_runtime_valid_k=partition_runtime_valid_k, enable_mxi_imlp=enable_mxi_imlp, - softmax_row_max=softmax_row_max, - softmax_row_sum=softmax_row_sum, - helix_kv_bounds=helix_kv_bounds, ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py index 7f9463d190f3..26f24caedaab 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py @@ -80,40 +80,19 @@ def _compile_cutedsl(*args, **kwargs): _EXPLICIT_TORCH_STREAM: torch.cuda.Stream | None = None -@contextlib.contextmanager -def _launch_stream(): - """Yield the CUDA driver stream these kernels should launch on. - - cuda2ctl capture cannot reliably query the default null stream, so SMART - wrappers can ask for an explicit one through DKG_MLA_EXPLICIT_STREAM. The - substitution must be bracketed rather than assigned: the caller's stream - already carries the RoPE, Q quantization and KV-cache writes these kernels - read, and it is also the stream their consumers read the output from. - Entering waits on the caller's work, leaving publishes ours back to it, and - ``torch.cuda.stream`` restores the thread's current stream on the way out. - A bare ``set_stream`` does none of the three. - - Switching the current stream is illegal while a CUDA graph is capturing, so - the knob is ignored under capture and the caller's stream is used as-is. - """ - entry_stream = torch.cuda.current_stream() - if os.environ.get("DKG_MLA_EXPLICIT_STREAM") != "1" or torch.cuda.is_current_stream_capturing(): - yield cuda.CUstream(entry_stream.cuda_stream) - return +def _current_cu_stream() -> cuda.CUstream: + """Return the active PyTorch stream as a CUDA driver stream. + cuda2ctl capture cannot reliably query the default null stream. Keep the + default behavior for normal pytest/AModel runs, but allow SMART wrappers to + request an explicit stream through the environment. + """ global _EXPLICIT_TORCH_STREAM - if _EXPLICIT_TORCH_STREAM is None: - _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() - start_event = torch.cuda.Event() - done_event = torch.cuda.Event() - start_event.record(entry_stream) - with torch.cuda.stream(_EXPLICIT_TORCH_STREAM): - _EXPLICIT_TORCH_STREAM.wait_event(start_event) - try: - yield cuda.CUstream(_EXPLICIT_TORCH_STREAM.cuda_stream) - finally: - done_event.record(_EXPLICIT_TORCH_STREAM) - entry_stream.wait_event(done_event) + if os.environ.get("DKG_MLA_EXPLICIT_STREAM") == "1": + if _EXPLICIT_TORCH_STREAM is None: + _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() + torch.cuda.set_stream(_EXPLICIT_TORCH_STREAM) + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) @dataclass(frozen=True) @@ -801,44 +780,44 @@ def fp4_mla_repack_v_cache( ) if resolve_generation_pages and not use_tma_fast_path: raise ValueError("generation-aware V repack requires the PAGE128/BLOCKV128 TMA path") - with _launch_stream() as stream: - repack_fn = _compile_fp4_mla_v_repack( - v_packed.data_ptr(), - kv_cache.data_ptr(), - page_ids_data_ptr, - page_indptr_data_ptr, - kv_lens_data_ptr, - generation_lens_data_ptr, - page_size, - v_head_dim, - use_page_ids, - resolve_generation_pages, - max_touched_pages, - block_v, - use_tma_fast_path, - stream, - ) - ptrs = _make_v_repack_ptrs( - v_packed.data_ptr(), - kv_cache.data_ptr(), - page_ids_data_ptr, - page_indptr_data_ptr, - kv_lens_data_ptr, - generation_lens_data_ptr, - ) - repack_fn( - *ptrs, - stream, - ctm.Int32(layout.num_pages), - ctm.Int64(layout.stride_page), - ctm.Int64(layout.stride_token), - ctm.Int64(layout.stride_packed_dim), - ctm.Int64(v_packed.stride(0)), - ctm.Int64(v_packed.stride(1)), - ctm.Int64(page_ids_stride), - ctm.Int32(num_page_ids), - ctm.Int32(num_generation_sequences), - ) + stream = _current_cu_stream() + repack_fn = _compile_fp4_mla_v_repack( + v_packed.data_ptr(), + kv_cache.data_ptr(), + page_ids_data_ptr, + page_indptr_data_ptr, + kv_lens_data_ptr, + generation_lens_data_ptr, + page_size, + v_head_dim, + use_page_ids, + resolve_generation_pages, + max_touched_pages, + block_v, + use_tma_fast_path, + stream, + ) + ptrs = _make_v_repack_ptrs( + v_packed.data_ptr(), + kv_cache.data_ptr(), + page_ids_data_ptr, + page_indptr_data_ptr, + kv_lens_data_ptr, + generation_lens_data_ptr, + ) + repack_fn( + *ptrs, + stream, + ctm.Int32(layout.num_pages), + ctm.Int64(layout.stride_page), + ctm.Int64(layout.stride_token), + ctm.Int64(layout.stride_packed_dim), + ctm.Int64(v_packed.stride(0)), + ctm.Int64(v_packed.stride(1)), + ctm.Int64(page_ids_stride), + ctm.Int32(num_page_ids), + ctm.Int32(num_generation_sequences), + ) def fp4_mla_repack_v_cache_reference( diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py index f844195e0df2..5295c6f788d0 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py @@ -286,96 +286,6 @@ def _fp4_mla_swizzled_sf_offset( # FP4 conversion and cache kernels -@triton.jit( - do_not_specialize=[ - "num_page_ids", - "num_pages", - "num_layers", - "local_layer", - ], - do_not_specialize_on_alignment=[ - "num_page_ids", - "num_pages", - "num_layers", - "local_layer", - ], -) -def _fp4_mla_rebuild_v_scale_from_k_scale_kernel( - k_sf_ptr, - v_sf_ptr, - page_ids_ptr, - page_valid_tokens_ptr, - num_page_ids, - num_pages, - num_layers, - local_layer, - page_size, - k_sf_s0, - v_sf_s0, - v_sf_s1, - V_HEAD_D: tl.constexpr, - HP_BLOCK: tl.constexpr, - SF_PER_TOKEN: tl.constexpr, - SF_PER_PAGE: tl.constexpr, - BLOCK_TOKEN_GROUPS: tl.constexpr, -): - """Rebuild dim-major MLA V scales from transferred token-major K scales. - - Context quantization uses one shared scale for every 16x16 compressed-V - tile. That byte is repeated across the tile's 16 K rows and 16 V rows, - so the process-local V layout can be reconstructed exactly from the first - valid K row of every token tile. - """ - page_work_idx = tl.program_id(0) - dim_block = tl.program_id(1) - if page_work_idx >= num_page_ids: - return - if (local_layer < 0) | (local_layer >= num_layers): - return - - physical_page = tl.load(page_ids_ptr + page_work_idx).to(tl.int64) - if (physical_page < 0) | (physical_page >= num_pages): - return - - valid_tokens = tl.load(page_valid_tokens_ptr + page_work_idx).to(tl.int32) - valid_tokens = tl.maximum(0, tl.minimum(valid_tokens, page_size)) - token_groups = tl.arange(0, BLOCK_TOKEN_GROUPS) - token_group_valid = (token_groups * HP_BLOCK < valid_tokens) & ( - token_groups * HP_BLOCK < page_size - ) - - # Scale tensors are passed as uint8 views so this is a bit-exact copy of - # the FP8 E4M3 encoding, including signed zero if it is ever produced. - k_rows = token_groups * HP_BLOCK - k_sf_offsets = _fp4_mla_swizzled_sf_offset( - k_rows, - dim_block, - SF_PER_TOKEN, - ) - scale_bits = tl.load( - k_sf_ptr + physical_page * k_sf_s0 + k_sf_offsets, - mask=token_group_valid, - other=0, - ) - - dims = dim_block * HP_BLOCK + tl.arange(0, HP_BLOCK) - dim_valid = dims < V_HEAD_D - safe_dims = tl.where(dim_valid, dims, 0) - v_sf_offsets = _fp4_mla_swizzled_sf_offset( - safe_dims[:, None], - token_groups[None, :], - SF_PER_PAGE, - ) - v_sf_base = tl.cast(local_layer, tl.int64) * tl.cast( - v_sf_s0, tl.int64 - ) + physical_page * tl.cast(v_sf_s1, tl.int64) - tl.store( - v_sf_ptr + v_sf_base + v_sf_offsets.to(tl.int64), - scale_bits[None, :], - mask=dim_valid[:, None] & (token_groups[None, :] * HP_BLOCK < page_size), - ) - - @triton.jit def _fp4_e2m1_to_f32(nibble): magnitude = nibble & 0x7 @@ -404,130 +314,6 @@ def _fp4_e2m1_to_f32(nibble): return tl.where(sign, -value, value) -@triton.jit -def _fp4_mla_chunked_cache_gather_kernel( - compressed_kv_ptr, - k_pe_ptr, - kv_cache_ptr, - sf_cache_ptr, - page_ids_ptr, - cu_chunked_seq_len_ptr, - chunked_global_offset_ptr, - global_scale_ptr, - max_chunk_len, - page_table_stride, - num_pages, - page_size, - kv_s0, - kv_s2, - kv_s4, - sf_s0, - compressed_kv_s0, - k_pe_s0, - KV_LORA_RANK: tl.constexpr, - QK_ROPE_HEAD_DIM: tl.constexpr, - FP4_BLOCK: tl.constexpr, - SF_PER_TOKEN: tl.constexpr, - TOKEN_BLOCK: tl.constexpr, -): - """Gather one logical MLA prefix chunk from the canonical FP4 K view.""" - token_block = tl.program_id(0) - batch_idx = tl.program_id(1) - dim_block = tl.program_id(2) - - local_tokens = token_block * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK) - chunk_start = tl.load(cu_chunked_seq_len_ptr + batch_idx).to(tl.int64) - chunk_end = tl.load(cu_chunked_seq_len_ptr + batch_idx + 1).to(tl.int64) - chunk_len = chunk_end - chunk_start - valid_tokens = (local_tokens < chunk_len) & (local_tokens < max_chunk_len) - - logical_positions = tl.load(chunked_global_offset_ptr + batch_idx).to( - tl.int64 - ) + local_tokens.to(tl.int64) - logical_page = logical_positions // page_size - page_position = logical_positions - logical_page * page_size - page_table_offsets = batch_idx * page_table_stride + logical_page - valid_tokens = valid_tokens & (logical_page >= 0) & (logical_page < page_table_stride) - physical_pages = tl.load(page_ids_ptr + page_table_offsets, mask=valid_tokens, other=0).to( - tl.int64 - ) - valid_tokens = valid_tokens & (physical_pages >= 0) & (physical_pages < num_pages) - - dims = dim_block * FP4_BLOCK + tl.arange(0, FP4_BLOCK) - head_dim = KV_LORA_RANK + QK_ROPE_HEAD_DIM - valid_dims = dims < head_dim - packed_cols = dims // 2 - packed = tl.load( - kv_cache_ptr - + physical_pages[:, None] * kv_s0 - + page_position[:, None] * kv_s2 - + packed_cols[None, :] * kv_s4, - mask=valid_tokens[:, None] & valid_dims[None, :], - other=0, - ).to(tl.uint8) - nibbles = tl.where((dims[None, :] & 1) == 0, packed & 0x0F, packed >> 4) - - sf_offsets = _fp4_mla_swizzled_sf_offset( - page_position, - dim_block, - SF_PER_TOKEN, - ) - scale = tl.load( - sf_cache_ptr + physical_pages * sf_s0 + sf_offsets, - mask=valid_tokens, - other=0.0, - ).to(tl.float32) - global_scale = tl.load(global_scale_ptr).to(tl.float32) - values = _fp4_e2m1_to_f32(nibbles) * scale[:, None] / global_scale - - if dim_block * FP4_BLOCK >= KV_LORA_RANK: - residual_group = dim_block - KV_LORA_RANK // FP4_BLOCK - residual_packed_cols = ( - head_dim // 2 + residual_group * (FP4_BLOCK // 2) + dims % FP4_BLOCK // 2 - ) - residual_packed = tl.load( - kv_cache_ptr - + physical_pages[:, None] * kv_s0 - + page_position[:, None] * kv_s2 - + residual_packed_cols[None, :] * kv_s4, - mask=valid_tokens[:, None] & valid_dims[None, :], - other=0, - ).to(tl.uint8) - residual_nibbles = tl.where( - (dims[None, :] & 1) == 0, - residual_packed & 0x0F, - residual_packed >> 4, - ) - residual_sf_col = head_dim // FP4_BLOCK + residual_group - residual_sf_offsets = _fp4_mla_swizzled_sf_offset( - page_position, - residual_sf_col, - SF_PER_TOKEN, - ) - residual_scale = tl.load( - sf_cache_ptr + physical_pages * sf_s0 + residual_sf_offsets, - mask=valid_tokens, - other=0.0, - ).to(tl.float32) - values += _fp4_e2m1_to_f32(residual_nibbles) * residual_scale[:, None] / global_scale - - output_tokens = chunk_start + local_tokens - if dim_block * FP4_BLOCK < KV_LORA_RANK: - output_dims = dims - tl.store( - compressed_kv_ptr + output_tokens[:, None] * compressed_kv_s0 + output_dims[None, :], - values, - mask=valid_tokens[:, None] & (output_dims[None, :] < KV_LORA_RANK), - ) - else: - output_dims = dims - KV_LORA_RANK - tl.store( - k_pe_ptr + output_tokens[:, None] * k_pe_s0 + output_dims[None, :], - values, - mask=valid_tokens[:, None] & (output_dims[None, :] < QK_ROPE_HEAD_DIM), - ) - - @triton.jit def _fp4_mla_floor_half_up(abs_value, multiplier, bias): return tl.inline_asm_elementwise( @@ -618,7 +404,6 @@ def _fp4_mla_context_cache_update_kernel( v_sf_ptr, v_packed_ptr, latent_cache_ptr, - q_context_ptr, global_scale_ptr, rotary_cos_sin_ptr, hp_pool_ptr, @@ -645,9 +430,6 @@ def _fp4_mla_context_cache_update_kernel( sf_s0, lc_s0, lc_s1, - q_s0, - q_s1, - q_s2, vsf_s0, vsf_s1, v_packed_s0, @@ -665,17 +447,12 @@ def _fp4_mla_context_cache_update_kernel( STORE_K_RESIDUAL: tl.constexpr, ROPE_DIM: tl.constexpr, APPLY_K_ROPE: tl.constexpr, - APPLY_Q_ROPE: tl.constexpr, - NUM_DIM_BLOCKS: tl.constexpr, - NUM_Q_HEADS: tl.constexpr, - Q_NOPE_DIM: tl.constexpr, - BLOCK_Q_HEADS: tl.constexpr, POOL_HEAD_D: tl.constexpr, STORE_HP_TAIL: tl.constexpr, WRITE_V_PACKED: tl.constexpr, ): token_idx = tl.program_id(0) - work_idx = tl.program_id(1) + dim_block = tl.program_id(1) if (local_layer < 0) | (local_layer >= num_layers): return if token_idx >= num_tokens: @@ -689,47 +466,6 @@ def _fp4_mla_context_cache_update_kernel( position = tl.load(positions_ptr + metadata_token_idx).to(tl.int64) if (batch_idx < 0) | (batch_idx + 1 >= indptr_len) | (position < 0): return - - if APPLY_Q_ROPE and work_idx >= NUM_DIM_BLOCKS: - q_head_block = work_idx - NUM_DIM_BLOCKS - q_heads = q_head_block * BLOCK_Q_HEADS + tl.arange(0, BLOCK_Q_HEADS) - q_head_mask = q_heads < NUM_Q_HEADS - pair_offsets = tl.arange(0, ROPE_DIM // 2) - rotary_offsets = position * (ROPE_DIM * 2) + pair_offsets * 2 - cos = tl.load(rotary_cos_sin_ptr + rotary_offsets).to(tl.float32) - sin = tl.load(rotary_cos_sin_ptr + rotary_offsets + 1).to(tl.float32) - q_base = token_idx * q_s0 + q_heads[:, None].to(tl.int64) * q_s1 + Q_NOPE_DIM * q_s2 - q_even_offsets = q_base + (pair_offsets[None, :] * 2).to(tl.int64) * q_s2 - q_odd_offsets = q_even_offsets + q_s2 - q_even = tl.load( - q_context_ptr + q_even_offsets, - mask=q_head_mask[:, None], - other=0.0, - ).to(tl.float32) - q_odd = tl.load( - q_context_ptr + q_odd_offsets, - mask=q_head_mask[:, None], - other=0.0, - ).to(tl.float32) - q_roped_even, q_roped_odd = _fp4_mla_rope_fp32( - q_even, - q_odd, - cos[None, :], - sin[None, :], - ) - tl.store( - q_context_ptr + q_even_offsets, - q_roped_even.to(tl.bfloat16), - mask=q_head_mask[:, None], - ) - tl.store( - q_context_ptr + q_odd_offsets, - q_roped_odd.to(tl.bfloat16), - mask=q_head_mask[:, None], - ) - return - - dim_block = work_idx if position % HP_BLOCK != 0: return @@ -815,24 +551,6 @@ def _fp4_mla_context_cache_update_kernel( roped_even, roped_odd = _fp4_mla_rope_fp32(even_values, odd_values, cos, sin) roped_even = roped_even.to(tl.bfloat16).to(tl.float32) roped_odd = roped_odd.to(tl.bfloat16).to(tl.float32) - if APPLY_Q_ROPE: - # Only chunked prefill consumes the rotated current K from - # latent_cache. Normal context overlaps this cache update with - # FP8 attention on another stream, so its input must stay read-only. - tl.store( - latent_cache_ptr - + safe_token_candidates[:, None] * lc_s0 - + safe_even_d[None, :] * lc_s1, - roped_even, - mask=rotary_mask, - ) - tl.store( - latent_cache_ptr - + safe_token_candidates[:, None] * lc_s0 - + safe_odd_d[None, :] * lc_s1, - roped_odd, - mask=rotary_mask, - ) even_values = tl.where(rotary_mask, roped_even, even_values) odd_values = tl.where(rotary_mask, roped_odd, odd_values) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py index 38814cc4d902..98bc954cc56a 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py @@ -1024,533 +1024,6 @@ def _fp4_mla_attention_page_stats_kernel( tl.store(page_sum_ptr + out_offsets, page_sum, mask=mask_h) -@triton.jit -def _fp4_mla_attention_page_stats_grouped_kernel( - page_max_ptr, - page_sum_ptr, - p_fp4_ptr, - p_sf_ptr, - q_fp4_ptr, - q_sf_ptr, - kv_cache_ptr, - sf_cache_ptr, - global_scale_ptr, - q_global_scale_ptr, - src_page_ids_ptr, - paged_kv_indptr_decode_ptr, - kv_lens_ptr, - page_ids_len, - num_pages, - q_fp4_s0, - q_fp4_s1, - kv_s0, - kv_s2, - kv_s4, - sf_s0, - page_stats_s0, - page_stats_s1, - p_s0, - p_s1, - p_num_rows, - q_num_rows, - sm_scale, - NUM_HEADS: tl.constexpr, - Q_HEAD_D: tl.constexpr, - K_HEAD_D: tl.constexpr, - Q_RESIDUAL_D: tl.constexpr, - PAGE_SIZE: tl.constexpr, - FP4_BLOCK: tl.constexpr, - Q_SF_PER_TOKEN: tl.constexpr, - K_SF_PER_TOKEN: tl.constexpr, - SF_PER_PAGE: tl.constexpr, - P_GLOBAL_SCALE: tl.constexpr, - QUERY_LEN_PER_SEQ: tl.constexpr, - MAX_PAGES: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_T: tl.constexpr, - BLOCK_K: tl.constexpr, - FULL_BLOCK_END: tl.constexpr, - TAIL_BLOCK_K: tl.constexpr, - GROUP_PAGES: tl.constexpr, - ASSUME_FULL_PAGES: tl.constexpr, - ASSUME_VALID_PAGES: tl.constexpr, - PAGE_LOOP_STAGES: tl.constexpr, - occupancy: tl.constexpr = 1, -): - """Grouped page-stats QK + softmax-stats + FP4 P pack. - - Functionally identical to ``_fp4_mla_attention_page_stats_kernel`` for the - perfect decode shape (NUM_HEADS == BLOCK_H, TMA + PACK_PROBS, the standard - 640/576/64 residual-Q layout) but each program owns one - ``(query, head_block, page_group)`` and walks ``GROUP_PAGES`` pages in a - pipelined loop. Q (and its scales) and the TMA descriptors are loaded once - and reused across the group, eliminating the per-page Q reload and the tiny - per-CTA prologue that made the one-page-per-CTA kernel work-bound at long - context. Per-page outputs (page_max/page_sum and packed P) are written - exactly as the one-page kernel writes them, so every downstream stage is - unchanged. - """ - query_idx = tl.program_id(0) - head_block = tl.program_id(1) - page_group = tl.program_id(2) - seq_idx = query_idx // QUERY_LEN_PER_SEQ - query_offset = query_idx - seq_idx * QUERY_LEN_PER_SEQ - - head_start = head_block * BLOCK_H - offs_h = head_start + tl.arange(0, BLOCK_H) - offs_t = tl.arange(0, BLOCK_T) - q_row_base = query_idx * NUM_HEADS - - if ASSUME_FULL_PAGES: - kv_len = 0 - else: - kv_len = tl.load(kv_lens_ptr + seq_idx) - (QUERY_LEN_PER_SEQ - 1 - query_offset) - kv_len = tl.maximum(kv_len, 0) - page_table_start = tl.load(paged_kv_indptr_decode_ptr + seq_idx).to(tl.int64) - - q_gscale = tl.load(q_global_scale_ptr) - - residual_groups = Q_RESIDUAL_D // FP4_BLOCK - non_residual_groups = K_HEAD_D // FP4_BLOCK - residual_groups - - # ---- Hoisted, page-independent index tensors and Q tiles ---- - # Main window (q_start == 0): the first FULL_BLOCK_END elements sit entirely - # in the non-residual region, so Q and K map 1:1 and load contiguously. - # Q rows are global rows (query_idx * NUM_HEADS + head); the swizzle's - # row_group term selects this query's scale/tail block, so q_row_base must - # be folded in (the main q_vals descriptor load already does this via its - # row coordinate). Keep int32 to match the one-page kernel -- int64 swizzle - # math is emulated and was measured ~2x slower; the max global row index - # (num_queries * NUM_HEADS * stride) stays well within int32. - q_rows = q_row_base + offs_h - scale_offsets = tl.arange(0, BLOCK_K // FP4_BLOCK) - q_sf_cols = scale_offsets - q_sf_offsets = _fp4_mla_swizzled_sf_offset(q_rows[:, None], q_sf_cols[None, :], Q_SF_PER_TOKEN) - k_sf_offsets_main = _fp4_mla_swizzled_sf_offset( - offs_t[:, None], q_sf_cols[None, :], K_SF_PER_TOKEN - ) - q_scales = tl.load(q_sf_ptr + q_sf_offsets) - - # Tail window (q_start == FULL_BLOCK_END): the residual-Q groups, each of - # which maps onto a duplicated K residual group. - tail_packed_offsets = tl.arange(0, TAIL_BLOCK_K // 2) - tail_scale_offsets = tl.arange(0, TAIL_BLOCK_K // FP4_BLOCK) - qt_elem = FULL_BLOCK_END + tail_packed_offsets * 2 - qt_group = qt_elem // FP4_BLOCK - kt_group = tl.where( - qt_group < non_residual_groups, - qt_group, - non_residual_groups + (qt_group - non_residual_groups) // 2, - ) - byte_t = (qt_elem % FP4_BLOCK) // 2 - packed_qt_cols = FULL_BLOCK_END // 2 + tail_packed_offsets - packed_kt_cols = kt_group * (FP4_BLOCK // 2) + byte_t - qt_sf_cols = FULL_BLOCK_END // FP4_BLOCK + tail_scale_offsets - kt_sf_cols = tl.where( - qt_sf_cols < non_residual_groups, - qt_sf_cols, - non_residual_groups + (qt_sf_cols - non_residual_groups) // 2, - ) - qt_sf_offsets = _fp4_mla_swizzled_sf_offset( - q_rows[:, None], qt_sf_cols[None, :], Q_SF_PER_TOKEN - ) - kt_sf_offsets = _fp4_mla_swizzled_sf_offset( - offs_t[:, None], kt_sf_cols[None, :], K_SF_PER_TOKEN - ) - q_tail_scales = tl.load(q_sf_ptr + qt_sf_offsets) - q_tail_vals = tl.load( - q_fp4_ptr + q_rows[:, None] * q_fp4_s0 + packed_qt_cols[None, :] * q_fp4_s1 - ) - - tl.assume(q_fp4_s0 % 8 == 0) - tl.assume(q_fp4_s1 == 1) - tl.assume(kv_s0 % 8 == 0) - tl.assume(kv_s2 % 8 == 0) - tl.assume(kv_s4 == 1) - tl.assume(p_s0 % 8 == 0) - tl.assume(p_s1 == 1) - q_desc = tl.make_tensor_descriptor( - q_fp4_ptr, - shape=[q_num_rows, Q_HEAD_D // 2], - strides=[q_fp4_s0, q_fp4_s1], - block_shape=[BLOCK_H, BLOCK_K // 2], - ) - k_desc = tl.make_tensor_descriptor( - kv_cache_ptr, - shape=[num_pages, BLOCK_T, K_HEAD_D // 2], - strides=[kv_s0, kv_s2, kv_s4], - block_shape=[1, BLOCK_T, BLOCK_K // 2], - ) - p_desc = tl.make_tensor_descriptor( - p_fp4_ptr, - shape=[p_num_rows, PAGE_SIZE // 2], - strides=[p_s0, p_s1], - block_shape=[BLOCK_H, PAGE_SIZE // 2], - ) - q_vals = q_desc.load([(q_row_base + head_start).to(tl.int32), 0]) - - scale_cols = tl.arange(0, SF_PER_PAGE) - byte_offsets = tl.arange(0, FP4_BLOCK // 2) - byte_cols = scale_cols[:, None] * (FP4_BLOCK // 2) + byte_offsets[None, :] - - page_lo = page_group * GROUP_PAGES - page_hi = page_lo + GROUP_PAGES - for page_rel in tl.range(page_lo, page_hi, num_stages=PAGE_LOOP_STAGES): - if page_rel < MAX_PAGES: - page_start = page_rel * PAGE_SIZE - page_max = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) - page_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) - if ASSUME_FULL_PAGES or page_start < kv_len: - compact_page = page_table_start + page_rel - if ASSUME_VALID_PAGES: - physical_page = tl.load(src_page_ids_ptr + compact_page).to(tl.int64) - safe_physical_page = physical_page - else: - valid_compact_page = (compact_page >= 0) & (compact_page < page_ids_len) - safe_compact_page = tl.where(valid_compact_page, compact_page, 0) - physical_page = tl.load( - src_page_ids_ptr + safe_compact_page, mask=valid_compact_page, other=-1 - ).to(tl.int64) - valid_physical_page = ( - valid_compact_page & (physical_page >= 0) & (physical_page < num_pages) - ) - safe_physical_page = tl.where(valid_physical_page, physical_page, 0) - - global_scale = tl.load(global_scale_ptr) - qk_scale = sm_scale / (q_gscale * global_scale) - - k_vals = k_desc.load([safe_physical_page.to(tl.int32), 0, 0]) - k_vals = tl.reshape(k_vals, (BLOCK_T, BLOCK_K // 2)) - if not ASSUME_VALID_PAGES: - k_vals = tl.where(valid_physical_page, k_vals, 0) - k_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + k_sf_offsets_main) - scores = tl.dot_scaled( - q_vals, - q_scales, - "e2m1", - k_vals.T, - k_scales, - "e2m1", - fast_math=True, - rhs_k_pack=True, - ) - - kt_ptrs = ( - kv_cache_ptr - + safe_physical_page * kv_s0 - + offs_t[:, None].to(tl.int64) * kv_s2 - + packed_kt_cols[None, :] * kv_s4 - ) - if ASSUME_VALID_PAGES: - kt_vals = tl.load(kt_ptrs) - else: - kt_vals = tl.load(kt_ptrs, mask=valid_physical_page, other=0) - kt_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + kt_sf_offsets) - scores = tl.dot_scaled( - q_tail_vals, - q_tail_scales, - "e2m1", - kt_vals.T, - kt_scales, - "e2m1", - acc=scores, - fast_math=True, - rhs_k_pack=True, - ) - - if ASSUME_FULL_PAGES: - scores = scores * qk_scale - page_max = tl.max(scores, axis=1) - exp_scores = tl.math.exp2((scores - page_max[:, None]) * _LOG2_E) - page_sum = tl.sum(exp_scores, axis=1) - else: - valid_t = page_start + offs_t < kv_len - scores = tl.where(valid_t[None, :], scores * qk_scale, -float("inf")) - page_max = tl.max(scores, axis=1) - exp_scores = tl.math.exp2((scores - page_max[:, None]) * _LOG2_E) - exp_scores = tl.where(valid_t[None, :], exp_scores, 0.0) - page_sum = tl.sum(exp_scores, axis=1) - - grouped_probs = tl.reshape(exp_scores, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK)) - amax = tl.max(grouped_probs, axis=2) - inv_local_scale = tl.where(amax > 0.0, 6.0 / amax, 1.0) - stored_scale = tl.where( - amax > 0.0, - tl.minimum(amax * (P_GLOBAL_SCALE / 6.0), 448.0), - 1.0, - ) - scaled_probs = grouped_probs * tl.reshape( - inv_local_scale, (BLOCK_H, SF_PER_PAGE, 1) - ) - pairs = tl.reshape(scaled_probs, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK // 2, 2)) - even_probs, odd_probs = tl.split(pairs) - packed = _fp4_e2m1_quantize_packed(even_probs, odd_probs) - - p_page = query_idx * MAX_PAGES + page_rel - if ASSUME_VALID_PAGES and NUM_HEADS == 128 and BLOCK_H == 128: - sf_offsets = _fp4_mla_swizzled_sf_offset_row_block( - p_page, offs_h[:, None], scale_cols[None, :], SF_PER_PAGE - ) - else: - p_rows = (p_page * NUM_HEADS + offs_h).to(tl.int64) - sf_offsets = _fp4_mla_swizzled_sf_offset( - p_rows[:, None], scale_cols[None, :], SF_PER_PAGE - ) - if ASSUME_VALID_PAGES: - tl.store(p_sf_ptr + sf_offsets, stored_scale) - p_desc.store( - [(p_page * NUM_HEADS + head_start).to(tl.int32), 0], - tl.reshape(packed, (BLOCK_H, PAGE_SIZE // 2)), - ) - else: - tl.store(p_sf_ptr + sf_offsets, stored_scale, mask=valid_compact_page) - p_rows = (p_page * NUM_HEADS + offs_h).to(tl.int64) - tl.store( - p_fp4_ptr + p_rows[:, None, None] * p_s0 + byte_cols[None, :, :] * p_s1, - packed, - mask=valid_compact_page, - ) - - out_offsets = query_idx * page_stats_s0 + page_rel * page_stats_s1 + offs_h - tl.store(page_max_ptr + out_offsets, page_max) - tl.store(page_sum_ptr + out_offsets, page_sum) - - -@triton.jit -def _fp4_mla_attention_page_stats_mtp_kernel( - page_max_ptr, - page_sum_ptr, - p_fp4_ptr, - p_sf_ptr, - q_fp4_ptr, - q_sf_ptr, - kv_cache_ptr, - sf_cache_ptr, - global_scale_ptr, - q_global_scale_ptr, - src_page_ids_ptr, - paged_kv_indptr_decode_ptr, - kv_lens_ptr, - page_ids_len, - num_pages, - q_fp4_s0, - q_fp4_s1, - kv_s0, - kv_s2, - kv_s4, - sf_s0, - page_stats_s0, - page_stats_s1, - p_s0, - p_s1, - p_num_rows, - q_num_rows, - sm_scale, - NUM_HEADS: tl.constexpr, - Q_HEAD_D: tl.constexpr, - K_HEAD_D: tl.constexpr, - Q_RESIDUAL_D: tl.constexpr, - PAGE_SIZE: tl.constexpr, - FP4_BLOCK: tl.constexpr, - Q_SF_PER_TOKEN: tl.constexpr, - K_SF_PER_TOKEN: tl.constexpr, - SF_PER_PAGE: tl.constexpr, - P_GLOBAL_SCALE: tl.constexpr, - QUERY_LEN_PER_SEQ: tl.constexpr, - MAX_PAGES: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_T: tl.constexpr, - BLOCK_K: tl.constexpr, - FULL_BLOCK_END: tl.constexpr, - TAIL_BLOCK_K: tl.constexpr, - occupancy: tl.constexpr = 1, -): - """MTP-fused page-stats: one CTA owns (seq, head_block, page) and processes - all QUERY_LEN_PER_SEQ linear-MTP query rows of the sequence, loading the - page's K (and K scales) once and reusing it across the q_len QK matmuls. - - The per-query-row kernel reloads K once per query row (q_len times per - page); at decode the QK is load-latency bound (one K load feeds one MMA), - so amortizing the K load over q_len rows lifts the load:MMA ratio. Per-row - outputs are written identically to the one-page kernel's masked path - (ASSUME_FULL_PAGES/VALID_PAGES are always False for q_len>1), so all - downstream stages are unchanged. Restricted to the perfect decode shape. - """ - seq_idx = tl.program_id(0) - head_block = tl.program_id(1) - page_rel = tl.program_id(2) - head_start = head_block * BLOCK_H - offs_h = head_start + tl.arange(0, BLOCK_H) - offs_t = tl.arange(0, BLOCK_T) - page_start = page_rel * PAGE_SIZE - - kv_len_base = tl.load(kv_lens_ptr + seq_idx) - page_table_start = tl.load(paged_kv_indptr_decode_ptr + seq_idx).to(tl.int64) - q_gscale = tl.load(q_global_scale_ptr) - - residual_groups = Q_RESIDUAL_D // FP4_BLOCK - non_residual_groups = K_HEAD_D // FP4_BLOCK - residual_groups - - # ---- r-independent index tensors (main + residual-tail column maps) ---- - scale_offsets = tl.arange(0, BLOCK_K // FP4_BLOCK) - q_sf_cols = scale_offsets - k_sf_offsets_main = _fp4_mla_swizzled_sf_offset( - offs_t[:, None], q_sf_cols[None, :], K_SF_PER_TOKEN - ) - tail_packed_offsets = tl.arange(0, TAIL_BLOCK_K // 2) - tail_scale_offsets = tl.arange(0, TAIL_BLOCK_K // FP4_BLOCK) - qt_elem = FULL_BLOCK_END + tail_packed_offsets * 2 - qt_group = qt_elem // FP4_BLOCK - kt_group = tl.where( - qt_group < non_residual_groups, - qt_group, - non_residual_groups + (qt_group - non_residual_groups) // 2, - ) - byte_t = (qt_elem % FP4_BLOCK) // 2 - packed_qt_cols = FULL_BLOCK_END // 2 + tail_packed_offsets - packed_kt_cols = kt_group * (FP4_BLOCK // 2) + byte_t - qt_sf_cols = FULL_BLOCK_END // FP4_BLOCK + tail_scale_offsets - kt_sf_cols = tl.where( - qt_sf_cols < non_residual_groups, - qt_sf_cols, - non_residual_groups + (qt_sf_cols - non_residual_groups) // 2, - ) - kt_sf_offsets = _fp4_mla_swizzled_sf_offset( - offs_t[:, None], kt_sf_cols[None, :], K_SF_PER_TOKEN - ) - scale_cols = tl.arange(0, SF_PER_PAGE) - byte_offsets = tl.arange(0, FP4_BLOCK // 2) - byte_cols = scale_cols[:, None] * (FP4_BLOCK // 2) + byte_offsets[None, :] - - tl.assume(q_fp4_s0 % 8 == 0) - tl.assume(q_fp4_s1 == 1) - tl.assume(kv_s0 % 8 == 0) - tl.assume(kv_s2 % 8 == 0) - tl.assume(kv_s4 == 1) - tl.assume(p_s0 % 8 == 0) - tl.assume(p_s1 == 1) - q_desc = tl.make_tensor_descriptor( - q_fp4_ptr, - shape=[q_num_rows, Q_HEAD_D // 2], - strides=[q_fp4_s0, q_fp4_s1], - block_shape=[BLOCK_H, BLOCK_K // 2], - ) - k_desc = tl.make_tensor_descriptor( - kv_cache_ptr, - shape=[num_pages, BLOCK_T, K_HEAD_D // 2], - strides=[kv_s0, kv_s2, kv_s4], - block_shape=[1, BLOCK_T, BLOCK_K // 2], - ) - # p_desc = tl.make_tensor_descriptor( - # p_fp4_ptr, - # shape=[p_num_rows, PAGE_SIZE // 2], - # strides=[p_s0, p_s1], - # block_shape=[BLOCK_H, PAGE_SIZE // 2], - # ) - - # ---- Load this page's K once (shared across all query rows). ---- - compact_page = page_table_start + page_rel - valid_compact_page = (compact_page >= 0) & (compact_page < page_ids_len) - safe_compact_page = tl.where(valid_compact_page, compact_page, 0) - physical_page = tl.load( - src_page_ids_ptr + safe_compact_page, mask=valid_compact_page, other=-1 - ).to(tl.int64) - valid_physical_page = valid_compact_page & (physical_page >= 0) & (physical_page < num_pages) - safe_physical_page = tl.where(valid_physical_page, physical_page, 0) - global_scale = tl.load(global_scale_ptr) - qk_scale = sm_scale / (q_gscale * global_scale) - - k_vals = k_desc.load([safe_physical_page.to(tl.int32), 0, 0]) - k_vals = tl.reshape(k_vals, (BLOCK_T, BLOCK_K // 2)) - k_vals = tl.where(valid_physical_page, k_vals, 0) - k_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + k_sf_offsets_main) - kt_vals = tl.load( - kv_cache_ptr - + safe_physical_page * kv_s0 - + offs_t[:, None].to(tl.int64) * kv_s2 - + packed_kt_cols[None, :] * kv_s4, - mask=valid_physical_page, - other=0, - ) - kt_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + kt_sf_offsets) - - for r in tl.static_range(QUERY_LEN_PER_SEQ): - query_idx_r = seq_idx * QUERY_LEN_PER_SEQ + r - kv_len_r = tl.maximum(kv_len_base - (QUERY_LEN_PER_SEQ - 1 - r), 0) - q_row_base_r = query_idx_r * NUM_HEADS - q_rows_r = q_row_base_r + offs_h - page_max = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) - page_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) - if page_start < kv_len_r: - q_vals = q_desc.load([(q_row_base_r + head_start).to(tl.int32), 0]) - q_sf_offsets = _fp4_mla_swizzled_sf_offset( - q_rows_r[:, None], q_sf_cols[None, :], Q_SF_PER_TOKEN - ) - q_scales = tl.load(q_sf_ptr + q_sf_offsets) - scores = tl.dot_scaled( - q_vals, - q_scales, - "e2m1", - k_vals.T, - k_scales, - "e2m1", - fast_math=True, - rhs_k_pack=True, - ) - qt_sf_offsets = _fp4_mla_swizzled_sf_offset( - q_rows_r[:, None], qt_sf_cols[None, :], Q_SF_PER_TOKEN - ) - q_tail_scales = tl.load(q_sf_ptr + qt_sf_offsets) - q_tail_vals = tl.load( - q_fp4_ptr + q_rows_r[:, None] * q_fp4_s0 + packed_qt_cols[None, :] * q_fp4_s1 - ) - scores = tl.dot_scaled( - q_tail_vals, - q_tail_scales, - "e2m1", - kt_vals.T, - kt_scales, - "e2m1", - acc=scores, - fast_math=True, - rhs_k_pack=True, - ) - - valid_t = page_start + offs_t < kv_len_r - scores = tl.where(valid_t[None, :], scores * qk_scale, -float("inf")) - page_max = tl.max(scores, axis=1) - exp_scores = tl.math.exp2((scores - page_max[:, None]) * _LOG2_E) - exp_scores = tl.where(valid_t[None, :], exp_scores, 0.0) - page_sum = tl.sum(exp_scores, axis=1) - - grouped_probs = tl.reshape(exp_scores, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK)) - amax = tl.max(grouped_probs, axis=2) - inv_local_scale = tl.where(amax > 0.0, 6.0 / amax, 1.0) - stored_scale = tl.where( - amax > 0.0, tl.minimum(amax * (P_GLOBAL_SCALE / 6.0), 448.0), 1.0 - ) - scaled_probs = grouped_probs * tl.reshape(inv_local_scale, (BLOCK_H, SF_PER_PAGE, 1)) - pairs = tl.reshape(scaled_probs, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK // 2, 2)) - even_probs, odd_probs = tl.split(pairs) - packed = _fp4_e2m1_quantize_packed(even_probs, odd_probs) - - p_page = query_idx_r * MAX_PAGES + page_rel - p_rows = (p_page * NUM_HEADS + offs_h).to(tl.int64) - sf_offsets = _fp4_mla_swizzled_sf_offset( - p_rows[:, None], scale_cols[None, :], SF_PER_PAGE - ) - tl.store(p_sf_ptr + sf_offsets, stored_scale, mask=valid_compact_page) - tl.store( - p_fp4_ptr + p_rows[:, None, None] * p_s0 + byte_cols[None, :, :] * p_s1, - packed, - mask=valid_compact_page, - ) - - out_offsets = query_idx_r * page_stats_s0 + page_rel * page_stats_s1 + offs_h - tl.store(page_max_ptr + out_offsets, page_max) - tl.store(page_sum_ptr + out_offsets, page_sum) - - @triton.jit def _fp4_mla_attention_reduce_stats_kernel( max_ptr, diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 7ab975a02a33..b246312d7ed3 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -30,6 +30,7 @@ from ...speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned +from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType from tensorrt_llm.logger import logger @@ -40,6 +41,11 @@ from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from .fmha.manager import FmhaManager +from .fp4_mla import (FP4_MLA_KV_GLOBAL_SCALE, FP4_MLA_Q_GLOBAL_SCALE, + HP_BLOCK_SIZE, can_fuse_fp4_mla_q_quant, + configure_fp4_mla_device_page_table, + populate_fp4_mla_append_metadata, + scatter_fp4_mla_kv_cache) from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, KVCacheParams, MLAParams, PositionalEmbeddingParams, @@ -218,6 +224,59 @@ def effective_beam_width(self) -> int: draft_block_ids_per_seq: Optional[torch.Tensor] = None draft_kv_block_ids_per_seq: Optional[torch.Tensor] = None + # True during warmup forward passes (dummy requests, no real data). + is_warmup: bool = False + + # High-precision BF16 KV pool for MLA FP4 models. The FP4 MLA V2 manager + # exposes compact paged rings. + high_precision_kv_pool: Optional[torch.Tensor] = None + _fp4_mla_hp_page_indices: Optional[torch.Tensor] = None + fp4_mla_v_scale_pool: Optional[torch.Tensor] = None + _fp4_mla_q_global_scale: Optional[torch.Tensor] = None + _fp4_mla_kv_global_scale: Optional[torch.Tensor] = None + batch_indices: Optional[torch.Tensor] = None + positions: Optional[torch.Tensor] = None + fp4_mla_generation_kv_lens: Optional[torch.Tensor] = None + fp4_mla_generation_append_lens: Optional[torch.Tensor] = None + fp4_mla_generation_lengths_num_tokens: int = field(init=False, default=-1) + fp4_mla_generation_lengths_num_seqs: int = field(init=False, default=-1) + fp4_mla_generation_lengths_num_contexts: int = field(init=False, default=-1) + _fp4_mla_generation_lengths_capture_recorded: bool = field(init=False, + default=False, + repr=False) + _fp4_mla_generation_cache_scattered: bool = field(init=False, + default=False, + repr=False) + _fp4_mla_device_page_table: bool = field(init=False, + default=False, + repr=False) + _fp4_mla_device_page_table_valid: bool = field(init=False, + default=False, + repr=False) + fp4_mla_page_table_stride: int = field(init=False, default=0, repr=False) + fp4_mla_context_repack_max_touched_pages: int = field(init=False, + default=1, + repr=False) + _fp4_mla_prequantized_q: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) + _fp4_mla_prequantized_q_sf: Optional[torch.Tensor] = field(init=False, + default=None, + repr=False) + _fp4_mla_q_batch_capacity: Optional[int] = field(init=False, + default=None, + repr=False) + _fp4_mla_fp8_context_state: Optional[Tuple[Any, Any]] = field(init=False, + default=None, + repr=False, + compare=False) + _paged_kv_indptr: Optional[torch.Tensor] = None + paged_kv_indptr_decode: Optional[torch.Tensor] = None + _paged_kv_indices: Optional[torch.Tensor] = None + num_blocks: Optional[List[int]] = None + num_context_blocks: int = 0 + num_generation_blocks: int = 0 + # Pre-computed FlashMLA tile-scheduler metadata and num_splits. # Computed once per forward pass in TrtllmAttention.forward() and reused across layers. flash_mla_tile_scheduler_metadata: Optional[torch.Tensor] = None @@ -301,6 +360,33 @@ def tokens_per_block(self) -> Optional[int]: """ return self.kv_cache_manager.tokens_per_block if self.kv_cache_manager is not None else None + @property + def page_size(self) -> int: + """ + Number of tokens per cache page. + """ + assert self.kv_cache_manager is not None, "page_size requires a KV cache manager" + return self.kv_cache_manager.tokens_per_block + + @property + def paged_kv_indices(self) -> torch.Tensor: + """ + Flattened page table used by FP4 MLA helper kernels. + """ + if self._paged_kv_indices is None: + raise RuntimeError("paged_kv_indices is not allocated.") + total_blocks = self.num_context_blocks + self.num_generation_blocks + return self._paged_kv_indices[:total_blocks] + + @property + def paged_kv_indptr(self) -> torch.Tensor: + """ + Page-table indptr used by FP4 MLA helper kernels. + """ + if self._paged_kv_indptr is None: + raise RuntimeError("paged_kv_indptr is not allocated.") + return self._paged_kv_indptr[:self.num_seqs + 1] + @property def host_kv_cache_pool_pointers(self) -> Optional[torch.Tensor]: """ @@ -557,6 +643,110 @@ def _post_init_with_buffers(self, buffers) -> None: pin_memory=prefer_pinned(), ) + # Bind the native paged high-precision BF16 KV pool for MLA FP4 models. + if (self.kv_cache_manager is not None + and self.kv_cache_manager.kv_factor == 1 + and self.kv_cache_manager.dtype == DataType.NVFP4): + num_local_layers = self.kv_cache_manager.num_local_layers + head_dim = self.kv_cache_manager.head_dim + kv_factor = self.kv_cache_manager.kv_factor + hp_ring_size = int(self.kv_cache_manager._fp4_mla_hp_pool_size) + if hp_ring_size < HP_BLOCK_SIZE: + raise RuntimeError( + "FP4 MLA high-precision ring must retain at least one " + f"{HP_BLOCK_SIZE}-token quantization tile, got " + f"{hp_ring_size} slots.") + hp_pool = self.kv_cache_manager.get_fp4_mla_hp_pool() + expected_tail = hp_ring_size * head_dim + if (not isinstance(hp_pool, torch.Tensor) + or hp_pool.dtype != torch.bfloat16 + or hp_pool.device.type != 'cuda' or hp_pool.ndim != 4 + or hp_pool.shape[1] != num_local_layers + or hp_pool.shape[2] != kv_factor + or hp_pool.shape[3] != expected_tail): + raise RuntimeError( + "FP4 MLA V2 HP pool must be a CUDA BF16 tensor shaped " + "[pages, local_layers, kv_factor, ring * head_dim], got " + f"{getattr(hp_pool, 'shape', None)}.") + self.high_precision_kv_pool = hp_pool + self.batch_indices = self.get_empty( + buffers, + (self.max_num_tokens, ), + cache_name="fp4_mla_batch_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.positions = self.get_empty( + buffers, + (self.max_num_tokens, ), + cache_name="fp4_mla_positions", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.fp4_mla_generation_kv_lens = self.get_empty( + buffers, + (self.max_num_sequences, ), + cache_name="fp4_mla_generation_kv_lens", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.fp4_mla_generation_append_lens = self.get_empty( + buffers, + (self.max_num_sequences, ), + cache_name="fp4_mla_generation_append_lens", + dtype=torch.int32, + capture_graph=capture_graph, + ) + page_table_capacity = ( + self.max_num_sequences * + int(self.kv_cache_manager.max_blocks_per_seq)) + self._paged_kv_indices = self.get_empty( + buffers, + (page_table_capacity, ), + cache_name="fp4_mla_paged_kv_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self._fp4_mla_hp_page_indices = self.get_empty( + buffers, + (page_table_capacity, ), + cache_name="fp4_mla_hp_page_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self._paged_kv_indptr = self.get_empty( + buffers, + (self.max_num_sequences + 1, ), + cache_name="fp4_mla_paged_kv_indptr", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.paged_kv_indptr_decode = self.get_empty( + buffers, + (self.max_num_sequences + 1, ), + cache_name="fp4_mla_paged_kv_indptr_decode", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self._fp4_mla_q_global_scale = self.get_empty( + buffers, + (1, ), + cache_name="fp4_mla_q_global_scale", + dtype=torch.float32, + capture_graph=capture_graph, + ) + self._fp4_mla_q_global_scale.fill_(FP4_MLA_Q_GLOBAL_SCALE) + self._fp4_mla_kv_global_scale = self.get_empty( + buffers, + (1, ), + cache_name="fp4_mla_kv_global_scale", + dtype=torch.float32, + capture_graph=capture_graph, + ) + self._fp4_mla_kv_global_scale.fill_(FP4_MLA_KV_GLOBAL_SCALE) + self.fp4_mla_v_scale_pool = self.kv_cache_manager.get_mla_v_scale_pool( + ) + # Allocate static buffers for helix parallelism support. if self.enable_helix: self.helix_position_offsets = self.get_empty( @@ -625,6 +815,15 @@ def on_update_kv_lens(self): if self.enable_flash_mla: self._flash_mla_metadata_valid = False self._invalidate_mla_scheduler_buffers() + if getattr(self, '_fp4_mla_device_page_table', False): + self._fp4_mla_device_page_table_valid = False + self._update_fp4_mla_append_metadata() + + def _update_fp4_mla_append_metadata(self) -> None: + if self.high_precision_kv_pool is not None and self.num_tokens > 0: + self._invalidate_fp4_mla_generation_lengths() + if self._needs_fp4_mla_append_metadata(): + self._populate_fp4_mla_batch_indices_positions() def update_for_spec_dec(self) -> None: # MTP updates kv_lens_cuda in-place between sub-steps, which changes @@ -633,6 +832,14 @@ def update_for_spec_dec(self) -> None: if self.enable_flash_mla: self._flash_mla_metadata_valid = False self._invalidate_mla_scheduler_buffers() + if self.high_precision_kv_pool is None: + return + + num_seqs = self.num_seqs + self.prompt_lens_cuda_runtime = self.seq_lens_kv_cuda[:num_seqs] + if not torch.cuda.is_current_stream_capturing(): + self.prompt_lens_cpu_runtime = self.seq_lens_kv[:num_seqs] + self._update_fp4_mla_append_metadata() def _invalidate_mla_scheduler_buffers(self) -> None: # Spec-dec rewrites q_lens and kv_lens between sub-steps, so the cumulative @@ -640,6 +847,23 @@ def _invalidate_mla_scheduler_buffers(self) -> None: self._mla_scheduler_buffers_valid = False self._mla_ctx_cu_seqlens_valid = False + def restore_from_spec_dec(self) -> None: + # The spec-dec draft loop pointed the FP4 MLA length aliases at the + # temporary _seq_lens_cuda clone made by prepare_for_spec_dec. Rebind + # them to the restored stable buffers; otherwise the next forward's + # captured ops (CUDA graph) bake pointers to the clone, which is freed + # after capture, and every replay reads freed memory (garbage + # positions/kv lens -> corrupted KV writes and illegal memory access + # in the RoPE table lookup). + super().restore_from_spec_dec() + if self.high_precision_kv_pool is None: + return + num_seqs = self.num_seqs + self.kv_lens_cuda_runtime = self.kv_lens_cuda[:num_seqs] + self.prompt_lens_cuda_runtime = self.seq_lens_kv_cuda[:num_seqs] + if not torch.cuda.is_current_stream_capturing(): + self.prompt_lens_cpu_runtime = self.seq_lens_kv[:num_seqs] + def update_helix_param( self, helix_position_offsets: List[int], @@ -977,14 +1201,30 @@ def prepare(self) -> None: # tokens. Use the actual KV length (without extra tokens) for # kv_lens_runtime, which becomes host_past_key_value_lengths and # eventually mMaxSeqLenKv. + if self.high_precision_kv_pool is not None: + # FP4 MLA needs the per-forward append lengths rather than the + # original prompt lengths. Context scratch-cache metadata and MTP + # generation both consume these runtime views. + kv_lens_runtime = kv_lens[:self.num_seqs] + prompt_lens_cuda_runtime = self.seq_lens_kv_cuda[:self.num_seqs] + prompt_lens_cpu_runtime = self.seq_lens_kv[:self.num_seqs] + else: + kv_lens_runtime = kv_lens[:self.num_seqs] + prompt_lens_cuda_runtime = self.prompt_lens_cuda[:self.num_seqs] + prompt_lens_cpu_runtime = self.prompt_lens_cpu[:self.num_seqs] self._bind_runtime_views( kv_lens_cuda=self.kv_lens_cuda[:self.num_seqs], - kv_lens=kv_lens[:self.num_seqs], - prompt_lens_cuda=self.prompt_lens_cuda[:self.num_seqs], - prompt_lens_cpu=self.prompt_lens_cpu[:self.num_seqs], + kv_lens=kv_lens_runtime, + prompt_lens_cuda=prompt_lens_cuda_runtime, + prompt_lens_cpu=prompt_lens_cpu_runtime, host_request_types=self.host_request_types[:self.num_seqs], ) + if self.high_precision_kv_pool is not None: + self._invalidate_fp4_mla_generation_lengths() + self._configure_fp4_mla_page_metadata(kv_lens) + self._prepare_fp4_mla_append_metadata() + def prepare_encoder_decoder_from_precomputed_lengths( self, prompt_lens: torch.Tensor, kv_lens: torch.Tensor, context_kv_tokens: int, generation_kv_tokens: int, @@ -1035,6 +1275,7 @@ def prepare_encoder_decoder_from_precomputed_lengths( host_request_types=self.host_request_types[:num_seqs], ) + def prepare_encoder_only(self) -> None: """Fast path for encoder-only forward (eager + CUDA graph capture).""" extra_attrs = get_model_extra_attrs() @@ -1094,6 +1335,75 @@ def prepare_encoder_cuda_graph_replay(self, seq_lens: List[int], self._num_ctx_tokens = padded_num_tokens self.host_total_kv_lens[0] = padded_num_tokens + def _prepare_fp4_mla_append_metadata(self) -> None: + """Populate eager append metadata or defer it to CUDA graph replay. + + CUDA graph ``_forward_step`` captures ``on_update_kv_lens`` before the + model forward. Let that captured update own these buffers instead of + launching the same metadata kernel eagerly during input preparation. + + One-token generation is request-major in the fused FP4 update kernel, + so it derives the token position directly from the sequence length and + does not need these per-token buffers at capture or replay time. + """ + if (self.is_cuda_graph or self.num_tokens == 0 + or not self._needs_fp4_mla_append_metadata()): + return + self._populate_fp4_mla_batch_indices_positions() + + def _needs_fp4_mla_append_metadata(self) -> bool: + """Return whether a forward still consumes materialized token indices.""" + return self.num_contexts > 0 + + def _invalidate_fp4_mla_generation_lengths(self) -> None: + """Invalidate generation lengths before the next FP4 MLA forward step.""" + self.fp4_mla_generation_lengths_num_tokens = -1 + self.fp4_mla_generation_lengths_num_seqs = -1 + self.fp4_mla_generation_lengths_num_contexts = -1 + self._fp4_mla_generation_lengths_capture_recorded = False + + def _configure_fp4_mla_page_metadata(self, kv_lens: torch.Tensor) -> None: + """Configure fixed-stride device page metadata.""" + if self.kv_cache_manager is None or self.request_ids is None: + raise RuntimeError( + "FP4 MLA device page metadata requires a KV cache manager " + "and request IDs.") + assert self._paged_kv_indices is not None + assert self._paged_kv_indptr is not None + assert self.paged_kv_indptr_decode is not None + if not configure_fp4_mla_device_page_table(self, kv_lens): + raise RuntimeError( + "FP4 MLA requires fixed-stride device page metadata; the " + "current KV-cache manager or batch layout is unsupported.") + + def _populate_fp4_mla_batch_indices_positions(self) -> None: + """Populate FP4 MLA scatter/HP append metadata in one CUDA launch.""" + num_seqs = self.num_contexts + self.num_generations + if num_seqs == 0 or self.num_tokens == 0: + return + assert self.batch_indices is not None + assert self.positions is not None + assert self.kv_lens_cuda_runtime is not None + assert self.prompt_lens_cuda_runtime is not None + + append_lens = self.seq_lens_kv_cuda[:num_seqs] + if not append_lens.is_cuda: + raise RuntimeError( + "FP4 MLA append metadata requires CUDA sequence lengths.") + # Use the canonical tensors rather than the *_runtime aliases. A + # spec-dec sub-step can replace the canonical append lengths, and CUDA + # graph capture requires their stable, live storage. + populate_fp4_mla_append_metadata( + append_lens, + self.kv_lens_cuda[:num_seqs], + self.batch_indices, + self.positions, + num_tokens=self.num_tokens, + num_sequences=num_seqs, + num_contexts=self.num_contexts, + num_context_tokens=self.num_ctx_tokens, + ) + def prepare_flash_mla(self) -> None: self._flash_mla_metadata_valid = False # Request-specific fills and H2D copies must happen before replay, not @@ -1857,8 +2167,15 @@ def _compute_flash_mla_metadata(self, ) def _ensure_rope_table_size(self, required_max_positions: int) -> None: - if required_max_positions > self.rope_params.max_positions: - self.rope_params.max_positions = required_max_positions + floats_per_position = self.rope_params.dim * ( + 2 if self.rope_params.duplicate_data else 1) + table_max_positions = (self.rotary_cos_sin.numel() // + floats_per_position + if self.rotary_cos_sin is not None + and floats_per_position > 0 else 0) + if required_max_positions > table_max_positions: + self.rope_params.max_positions = max(required_max_positions, + self.rope_params.max_positions) self.rotary_inv_freq, self.rotary_cos_sin = ( self.rope_params.create_rope_const_params()) @@ -2259,6 +2576,49 @@ def forward( assert metadata.kv_cache_manager is None assert metadata.num_contexts == metadata.num_seqs + # Testing only: ``mla_rope_generation`` normally rotates q_pe, appends the + # new latent to the paged cache, and fills the trtllm-gen scheduler + # buffers (cumulative q/kv seqlens + the FMHA scheduler counter). When the + # harness sets ``skip_mla_rope_generation`` it feeds a pre-RoPE'd fused_q, + # so we skip only the RoPE and do the append + scheduler init here: the + # generation FMHA only reads the cache, and the fallback path needs the + # scheduler buffers (the flashinfer trtllm-gen decode kernel ignores them). + if (self.is_mla_enable and forward_args.skip_mla_rope_generation + and forward_args.attention_input_type + == AttentionInputType.generation_only): + num_ctx = metadata.num_contexts + n_gen = metadata.num_generations + # Use the GPU-resident length tensors (no host->device copy) so this + # stays CUDA-graph-capturable. + gen_q_lens = metadata.seq_lens_cuda[num_ctx:num_ctx + n_gen].to( + torch.int32) + gen_kv_lens = metadata.kv_lens_cuda_runtime[num_ctx:num_ctx + + n_gen].to(torch.int32) + cu_q = torch.zeros(n_gen + 1, dtype=torch.int32, device=q.device) + cu_kv = torch.zeros(n_gen + 1, dtype=torch.int32, device=q.device) + cu_q[1:] = torch.cumsum(gen_q_lens, dim=0).to( + torch.int32) * self.num_heads + cu_kv[1:] = torch.cumsum(gen_kv_lens, dim=0).to(torch.int32) + forward_args.cu_q_seqlens = cu_q + forward_args.cu_kv_seqlens = cu_kv + if forward_args.fmha_scheduler_counter is None: + forward_args.fmha_scheduler_counter = torch.zeros( + 1, dtype=torch.uint32, device=q.device) + else: + forward_args.fmha_scheduler_counter.zero_() + assert forward_args.latent_cache is not None + from .utils import append_mla_latent_cache + append_mla_latent_cache( + metadata.kv_cache_manager, + self.get_local_layer_idx(metadata), + metadata.request_ids, + metadata.seq_lens.tolist(), + metadata.kv_cache_params.num_cached_tokens_per_seq, + forward_args.latent_cache, + kv_layout=metadata.kv_layout, + seq_start=num_ctx, + ) + fmha = self._fmha_manager.select(self, q, k, v, metadata, forward_args) if fmha is None: @@ -2531,6 +2891,7 @@ def mla_rope_generation( kv_only: bool = False, kv_done_elsewhere: bool = False, quant_scale_qkv: Optional[torch.Tensor] = None, + fuse_fp4_q_quant: bool = False, ) -> None: """ fused_q (torch.Tensor): The tensor to store the fused q, with shape (num_tokens, num_heads, kv_lora_rank + qk_rope_head_dim) on GPU. @@ -2551,6 +2912,7 @@ def mla_rope_generation( kv_only (bool): Run the KV half only; the Q half already ran (q_b_layernorm folded the Q RoPE). Mutually exclusive with kv_done_elsewhere. kv_done_elsewhere (bool): Run the Q half only; the KV half already ran, hoisted onto an aux stream. Mutually exclusive with kv_only. With both halves done, do not call this at all. quant_scale_qkv (torch.Tensor): Non-None means q_nope in quant_q_buffer is already FP8, so the kernel drops the q_nope quantize rows from its grid. + fuse_fp4_q_quant (bool): Quantize Q in the fused FP4 RoPE/cache update. """ assert self.is_mla_enable and self.mla_params is not None @@ -2560,6 +2922,16 @@ def mla_rope_generation( # kernel reads it. self._ensure_rope_table_size(metadata.max_seq_len) + if self.has_fp4_kv_cache: + self._fp4_mla_rope_generation( + fused_q, + q_pe, + latent_cache, + metadata, + fuse_q_quant=fuse_fp4_q_quant, + ) + return + helix_tensor_params = [ metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] @@ -2618,3 +2990,71 @@ def mla_rope_generation( kv_done_elsewhere, quant_scale_qkv, ) + + def _fp4_mla_rope_generation( + self, + fused_q: torch.Tensor, + q_pe: torch.Tensor, + latent_cache: torch.Tensor, + metadata: TrtllmAttentionMetadata, + *, + fuse_q_quant: bool = False, + ) -> None: + """Apply fused generation RoPE, Q quantization, and cache update.""" + assert self.kv_lora_rank is not None + assert self.qk_rope_head_dim is not None + + if q_pe.shape[-1] != self.qk_rope_head_dim: + raise RuntimeError( + f"FP4 MLA q_pe last dimension must be {self.qk_rope_head_dim}, " + f"got {q_pe.shape[-1]}.") + if latent_cache.shape[-1] != self.kv_lora_rank + self.qk_rope_head_dim: + raise RuntimeError( + "FP4 MLA latent_cache last dimension must be " + f"{self.kv_lora_rank + self.qk_rope_head_dim}, got " + f"{latent_cache.shape[-1]}.") + if not fuse_q_quant: + raise RuntimeError( + "FP4 MLA generation requires fused Q quantization.") + if not self.can_fuse_fp4_mla_q_quant(fused_q, q_pe, latent_cache, + metadata): + raise RuntimeError( + "FP4 MLA fused Q quantization eligibility changed before " + "launch.") + if (self.rotary_cos_sin is None or q_pe.dtype != torch.bfloat16 + or fused_q.dtype != torch.bfloat16 + or latent_cache.dtype != torch.bfloat16 + or self.rotary_cos_sin.dtype != torch.float32): + raise RuntimeError( + "FP4 MLA generation requires fused BF16 RoPE/cache update " + "with an FP32 rotary table.") + + metadata._fp4_mla_generation_cache_scattered = False + hp_pool_updated = scatter_fp4_mla_kv_cache( + metadata, + latent_cache, + self.layer_idx, + token_offset=getattr(metadata, "num_ctx_tokens", 0), + phase="generation", + local_layer=self.get_local_layer_idx(metadata), + v_head_dim=self.kv_lora_rank, + rotary_cos_sin=self.rotary_cos_sin, + q_pe=q_pe, + q_rope_out=fused_q[..., self.kv_lora_rank:], + q_quant_input=fused_q, + ) + if not hp_pool_updated: + raise RuntimeError( + "Fused FP4 MLA RoPE/cache scatter did not update the HP pool.") + metadata._fp4_mla_generation_cache_scattered = True + + def can_fuse_fp4_mla_q_quant( + self, + fused_q: torch.Tensor, + q_pe: torch.Tensor, + latent_cache: torch.Tensor, + metadata: TrtllmAttentionMetadata, + ) -> bool: + return bool( + self.has_fp4_kv_cache + and can_fuse_fp4_mla_q_quant(metadata, fused_q, q_pe, latent_cache)) diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index e331d9ebcf24..5ed998391d17 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -1579,6 +1579,20 @@ def forward_absorption_generation( device=q.device, ) + fp4_mla = isinstance(self.mqa, TrtllmAttention) and self.mqa.has_fp4_kv_cache + if fp4_mla and ( + self.k_b_proj_trans.dtype != torch.bfloat16 + or latent_cache is None + or self.mapping.has_cp_helix() + ): + raise RuntimeError( + "FP4 MLA generation requires fused BF16 Q " + "quantization, RoPE, and cache update on the TRT-LLM " + "backend." + ) + fuse_fp4_q_quant = fp4_mla + fp4_rope_kwargs = {"fuse_fp4_q_quant": True} if fuse_fp4_q_quant else {} + def _mla_gen_rope(): if self.apply_rotary_emb: # Non-fused backends (Vanilla / FlashInfer) do not fuse RoPE @@ -1608,9 +1622,10 @@ def _mla_gen_rope(): mla_bmm1_scale, mla_bmm2_scale, quant_q_buffer, + **fp4_rope_kwargs, ) - rope_stream = self.aux_stream if not use_fp8_mla else None + rope_stream = None if fuse_fp4_q_quant else self.aux_stream if not use_fp8_mla else None if self.k_b_proj_trans.dtype == torch.bfloat16: # [num_heads, num_tokens, self.qk_nope_head_dim] q_nope_t = q_nope.transpose(0, 1) From d199b530126d70ddad8061867992a1be20ac42c9 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:40:17 -0700 Subject: [PATCH 02/21] [None][feat] integrate FP4 MLA KV cache manager V2 Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../batch_manager/kvCacheManager.cpp | 3 + .../backends/fp4_mla/cache_manager.py | 505 ++---------------- tensorrt_llm/_torch/pyexecutor/_util.py | 27 + .../_torch/pyexecutor/model_engine.py | 8 + .../_torch/pyexecutor/py_executor_creator.py | 91 +++- .../_torch/pyexecutor/resource_manager.py | 22 +- tensorrt_llm/_torch/speculative/mtp.py | 6 + 7 files changed, 192 insertions(+), 470 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 255452dd654b..43c7e6605ec2 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -3270,6 +3270,9 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case , mEnableBlockReuse{mSinkBubbleLength > 0 ? false : enableBlockReuse} { + TLLM_CHECK_WITH_INFO(dtype != tensorrt_llm::DataType::kFP4 || cacheType != CacheType::kSELFKONLY, + "NVFP4 SELFKONLY cache storage requires Fp4MlaKVCacheManagerV2; KVCacheManager V1 is not supported."); + // When num_layers < len(maxAttentionWindowVec), not all window sizes in the // repeating pattern are used. Update mMaxAttentionWindow to the actual // maximum window size that has been allocated in the block manager. diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index 55983fd1572c..4ab4e0458cee 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -4,13 +4,10 @@ import math from dataclasses import dataclass, replace -from types import MethodType -from typing import Iterable, List, Optional +from typing import List, Optional import torch -from tensorrt_llm._torch.disaggregation.resource.page import MapperKind -from tensorrt_llm._torch.kimi_k3_cache_policy import get_kimi_k3_bf16_kv_layer_ids from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned @@ -47,57 +44,22 @@ class Fp4MlaPageTableSpec: hp_is_paged: bool = True -class Fp4MlaV2CacheLayoutPolicy: - """FP4 MLA storage policy for a KV cache manager V2 lifecycle. +class Fp4MlaKVCacheManagerV2(KVCacheManagerV2): + """V2 manager for canonical FP4 MLA pages and compact HP sequence state. K, K block scales, V scales, and optional packed V share one full-history lifecycle. Each model layer also has a virtual sliding layer whose page is a compact ``16 + max_rewind`` BF16 ring. The HP role therefore follows V2 - allocation, rewind, prefix, and disaggregated-transfer lifecycles without - storing BF16 values for the full sequence. - - ``Fp4MlaKVCacheManagerV2`` applies this policy directly for pure MLA - models. Hybrid linear-attention managers compose the same policy with - their SSM lifecycle and forward only attention-layout operations here. + allocation, rewind, and prefix lifecycles without storing BF16 values for + the full sequence. """ def __init__(self, *args, **kwargs) -> None: kv_cache_config = args[0] if args else kwargs.get("kv_cache_config") dtype = kwargs.get("dtype", DataType.HALF) kv_cache_type = args[1] if len(args) > 1 else kwargs.get("kv_cache_type") - kv_cache_config = self.configure_manager( - self, - kv_cache_config, - kv_cache_type, - dtype=dtype, - tokens_per_block=kwargs.get("tokens_per_block"), - head_dim=kwargs.get("head_dim"), - pretrained_config=kwargs.get("pretrained_config"), - spec_config=kwargs.get("spec_config"), - is_disagg=kwargs.get("is_disagg", False), - ) - if args: - args = (kv_cache_config, *args[1:]) - else: - kwargs["kv_cache_config"] = kv_cache_config - - super().__init__(*args, **kwargs) - self.finalize_manager(self) - - @staticmethod - def configure_manager( - manager, - kv_cache_config, - kv_cache_type, - *, - dtype: DataType, - tokens_per_block: Optional[int], - head_dim: Optional[int], - pretrained_config=None, - spec_config=None, - is_disagg: bool = False, - ): - """Validate and install FP4 MLA layout state before V2 construction.""" + tokens_per_block = kwargs.get("tokens_per_block") + head_dim = kwargs.get("head_dim") if dtype != DataType.NVFP4 or kv_cache_type != CacheTypeCpp.SELFKONLY: raise ValueError("Fp4MlaKVCacheManagerV2 requires NVFP4 SELFKONLY cache storage.") if kv_cache_config is None or kv_cache_config.dtype not in ("auto", "nvfp4"): @@ -115,6 +77,10 @@ def configure_manager( config_updates["enable_partial_reuse"] = False if config_updates: kv_cache_config = kv_cache_config.model_copy(update=config_updates) + if args: + args = (kv_cache_config, *args[1:]) + else: + kwargs["kv_cache_config"] = kv_cache_config if "enable_partial_reuse" in config_updates: logger.info( "FP4 MLA KV cache manager V2 disables partial block reuse " @@ -128,198 +94,47 @@ def configure_manager( if not isinstance(head_dim, int) or head_dim <= FP4_MLA_K_RESIDUAL_DIM: raise ValueError(f"FP4 MLA V2 requires a positive scalar MLA head_dim, got {head_dim}.") - manager.mla_v_scale_head_dim = int( + pretrained_config = kwargs.get("pretrained_config") + self.mla_v_scale_head_dim = int( getattr(pretrained_config, "kv_lora_rank", head_dim - FP4_MLA_K_RESIDUAL_DIM) ) - if not 0 < manager.mla_v_scale_head_dim < head_dim: + if not 0 < self.mla_v_scale_head_dim < head_dim: raise ValueError( "FP4 MLA V2 requires kv_lora_rank in (0, head_dim), got " - f"{manager.mla_v_scale_head_dim} and {head_dim}." + f"{self.mla_v_scale_head_dim} and {head_dim}." ) - manager._bf16_mla_global_layer_ids = get_kimi_k3_bf16_kv_layer_ids(pretrained_config) - if manager._bf16_mla_global_layer_ids: - logger.info( - "Kimi K3 MLA layers using BF16 KV-cache fallback: " - f"{sorted(manager._bf16_mla_global_layer_ids)}." - ) - manager._fp4_mla_storage_backend = _fp4_mla_attention_backend() - if manager._fp4_mla_storage_backend not in _FP4_MLA_K_RESIDUAL_BACKENDS: + self._fp4_mla_storage_backend = _fp4_mla_attention_backend() + if self._fp4_mla_storage_backend not in _FP4_MLA_K_RESIDUAL_BACKENDS: raise ValueError( "Fp4MlaKVCacheManagerV2 supports only the triton and cutedsl " - f"backends, got {manager._fp4_mla_storage_backend!r}." + f"backends, got {self._fp4_mla_storage_backend!r}." ) - manager.fp4_mla_k_residual_dim = ( + self.fp4_mla_k_residual_dim = ( FP4_MLA_K_RESIDUAL_DIM - if manager._fp4_mla_storage_backend in _FP4_MLA_K_RESIDUAL_BACKENDS + if self._fp4_mla_storage_backend in _FP4_MLA_K_RESIDUAL_BACKENDS else 0 ) - manager.mla_v_head_dim = ( - manager.mla_v_scale_head_dim - if manager._fp4_mla_storage_backend == _FP4_MLA_CUTEDSL_BACKEND + self.mla_v_head_dim = ( + self.mla_v_scale_head_dim + if self._fp4_mla_storage_backend == _FP4_MLA_CUTEDSL_BACKEND and not _fp4_mla_cutedsl_fused_v_transpose_enabled() else None ) + spec_config = kwargs.get("spec_config") max_rewind_len = int(spec_config.tokens_per_gen_step - 1) if spec_config else 0 - manager._fp4_mla_hp_pool_size = HP_BLOCK_SIZE + max_rewind_len - manager._fp4_mla_view_cache = {} - manager._attention_cache_layout_policy = Fp4MlaV2CacheLayoutPolicy - Fp4MlaV2CacheLayoutPolicy.install_manager_capabilities(manager) - return kv_cache_config + self._fp4_mla_hp_pool_size = HP_BLOCK_SIZE + max_rewind_len + self._fp4_mla_view_cache: dict[tuple, torch.Tensor] = {} - @staticmethod - def install_manager_capabilities(manager) -> None: - """Attach attention-layout hooks to a non-policy lifecycle manager.""" - if isinstance(manager, Fp4MlaV2CacheLayoutPolicy): - return - method_names = ( - "_bf16_mla_local_layer_indices", - "_fp4_mla_local_layer_indices", - "_fp4_mla_compact_layer_idx", - "_bf16_mla_bytes_per_token", - "_get_buffer_roles_for_layer", - "_storage_head_dim", - "get_buffers", - "get_kv_cache_dtype", - "get_kv_cache_num_blocks", - "_v_scale_bytes_per_page", - "_v_packed_bytes_per_page", - "_hp_bytes_per_page", - "get_layer_bytes_per_token", - "_extra_buffers_per_layer", - "_prepare_page_table_tensor", - "_validate_fp4_mla_layer_groups", - "get_fp4_mla_page_table_spec", - "_role_encoded_page_capacity", - "_role_view", - "get_fp4_mla_cache_buffers", - "_iter_fp4_mla_physical_pool_views", - "_all_layer_role_view", - "_role_pool_base_view", - "get_mla_v_scale_pool", - "get_mla_v_scale_pool_base", - "get_mla_v_scale_page_offset", - "get_mla_v_packed_pool", - "get_mla_v_packed_pool_base", - "get_mla_v_packed_page_offset", - "get_fp4_mla_hp_pool", - "get_disagg_role_mapper_kinds", - "get_disagg_transfer_roles", - "get_disagg_global_layer_ids", - "_get_runtime_cache_size_layer_components", - "_get_generation_request_capacity", - ) - for method_name in method_names: - method = getattr(Fp4MlaV2CacheLayoutPolicy, method_name) - setattr(manager, method_name, MethodType(method, manager)) + super().__init__(*args, **kwargs) - @staticmethod - def finalize_manager(manager) -> None: - """Validate allocated pools and initialize derived FP4 state.""" - Fp4MlaV2CacheLayoutPolicy._validate_fp4_mla_layer_groups(manager) + self._validate_fp4_mla_layer_groups() # Partial pages leave unused V-scale tiles untouched while the # fixed-width PV path can load the complete scale page. Match V1's # deterministic initialization before any warmup or graph capture. - with torch.cuda.stream(manager._stream): - Fp4MlaV2CacheLayoutPolicy.get_mla_v_scale_pool_base(manager).zero_() - manager._stream.synchronize() - - def _bf16_mla_local_layer_indices(self) -> list[int]: - fallback_layers = getattr(self, "_bf16_mla_global_layer_ids", frozenset()) - is_linear_layer = getattr(self, "_is_local_mamba_layer", None) - return [ - local_layer - for local_layer in range(self.num_local_layers) - if self.pp_layers[local_layer] in fallback_layers - and (not callable(is_linear_layer) or not is_linear_layer(local_layer)) - ] - - def _fp4_mla_local_layer_indices(self) -> list[int]: - fallback_layers = set(self._bf16_mla_local_layer_indices()) - is_linear_layer = getattr(self, "_is_local_mamba_layer", None) - return [ - local_layer - for local_layer in range(self.num_local_layers) - if local_layer not in fallback_layers - and (not callable(is_linear_layer) or not is_linear_layer(local_layer)) - ] - - def _fp4_mla_compact_layer_idx(self, local_layer: int) -> int: - try: - return self._fp4_mla_local_to_compact[local_layer] - except KeyError as error: - raise ValueError( - f"Local layer {local_layer} is not an FP4 MLA attention layer." - ) from error - - def _bf16_mla_bytes_per_token(self, local_layer_idx: int) -> int: - return ( - self.num_kv_heads_per_layer[local_layer_idx] - * self.head_dim_per_layer[local_layer_idx] - * torch.empty((), dtype=torch.bfloat16).element_size() - ) - - def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: - if local_layer_idx in self._bf16_mla_local_layer_indices(): - return [Role.KEY] - return KVCacheManagerV2._get_buffer_roles_for_layer(self, local_layer_idx) - - def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: - """Return the layer's primary paged-cache view. - - BF16 fallback layers use a buffer role and physical page geometry that - differ from the manager-wide NVFP4 layout. Construct their view from - the role descriptor directly so FMHA dispatch observes the same BF16 - dtype and logical shape as an all-BF16 cache manager. - """ - local_layer = self.layer_offsets[layer_idx] - is_linear_layer = getattr(self, "_is_local_mamba_layer", None) - if callable(is_linear_layer) and is_linear_layer(local_layer): - return None - if local_layer not in self._bf16_mla_local_layer_indices(): - return KVCacheManagerV2.get_buffers(self, layer_idx, kv_layout) - if kv_layout not in ("NHD", "HND"): - raise ValueError(f"Unsupported kv_layout: {kv_layout}") - - page_shape = ( - ( - self.kv_factor, - self.tokens_per_block, - self.num_kv_heads_per_layer[local_layer], - self.head_dim_per_layer[local_layer], - ) - if kv_layout == "NHD" - else ( - self.kv_factor, - self.num_kv_heads_per_layer[local_layer], - self.tokens_per_block, - self.head_dim_per_layer[local_layer], - ) - ) - return self._role_view( - LayerId(local_layer), - Role.KEY, - torch.bfloat16, - page_shape, - ) - - def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: - if layer_idx is None: - return self.dtype - local_layer = self.layer_offsets[layer_idx] - if local_layer in self._bf16_mla_local_layer_indices(): - return DataType.BF16 - return self.dtype - - def get_kv_cache_num_blocks(self, layer_idx: int) -> int: - local_layer = self.layer_offsets[layer_idx] - if local_layer in self._bf16_mla_local_layer_indices(): - manager_layer = LayerId(local_layer) - else: - manager_layer = self._cache_manager_layer_ids[ - self._fp4_mla_compact_layer_idx(local_layer) - ] - return self._role_encoded_page_capacity(manager_layer, Role.KEY) + with torch.cuda.stream(self._stream): + self.get_mla_v_scale_pool_base().zero_() + self._stream.synchronize() @property def blocks_in_primary_pool(self) -> int: @@ -344,12 +159,6 @@ def _hp_bytes_per_page(self, local_layer_idx: int) -> int: ) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole): - if local_layer_idx in self._bf16_mla_local_layer_indices(): - if data_role in (Role.KEY, Role.ALL): - return self._bf16_mla_bytes_per_token(local_layer_idx) - raise ValueError(f"Invalid BF16 MLA V2 data role: {data_role}") - if local_layer_idx not in self._fp4_mla_local_layer_indices(): - return KVCacheManagerV2.get_layer_bytes_per_token(self, local_layer_idx, data_role) storage_head_dim = self._storage_head_dim(local_layer_idx) role_sizes = { Role.KEY: math.ceil(storage_head_dim / 2), @@ -373,7 +182,7 @@ def _extra_buffers_per_layer( self, *, tokens_per_block: int ) -> Optional[dict[int, List[BufferConfig]]]: result = {} - for local_layer in self._fp4_mla_local_layer_indices(): + for local_layer in range(self.num_local_layers): buffers = [ BufferConfig( role=Role.MLA_V_SCALE, @@ -392,39 +201,9 @@ def _extra_buffers_per_layer( def _build_cache_config(self, config): cache_layers = list(config.layers) - bf16_local_layers = self._bf16_mla_local_layer_indices() - fp4_local_layers = self._fp4_mla_local_layer_indices() - if not fp4_local_layers: - raise ValueError("FP4 MLA V2 layout requires at least one local attention layer.") - self._bf16_mla_manager_layer_ids = [ - LayerId(local_layer) for local_layer in bf16_local_layers - ] - if bf16_local_layers: - bf16_lifecycle_window = ( - self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1 - ) - for local_layer in bf16_local_layers: - cache_layers[local_layer] = AttentionLayerConfig( - layer_id=LayerId(local_layer), - buffers=[ - BufferConfig( - role=Role.KEY, - size=self._bf16_mla_bytes_per_token(local_layer) - * self.tokens_per_block, - ) - ], - # A distinct, non-evicting lifecycle gives BF16 pages their - # own attention-op pool without changing the model window. - sliding_window_size=bf16_lifecycle_window, - num_sink_tokens=None, - ) - self._fp4_mla_local_to_compact = { - local_layer: compact_layer for compact_layer, local_layer in enumerate(fp4_local_layers) - } - self._fp4_mla_compact_to_local = fp4_local_layers - self._cache_manager_layer_ids = [LayerId(i) for i in fp4_local_layers] + self._cache_manager_layer_ids = [LayerId(i) for i in range(self.num_local_layers)] self._hp_manager_layer_ids = [] - for local_layer in fp4_local_layers: + for local_layer in range(self.num_local_layers): layer_id = LayerId(len(cache_layers)) self._hp_manager_layer_ids.append(layer_id) cache_layers.append( @@ -449,12 +228,6 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: hp_pool_id = int(self.impl.get_layer_group_id(hp_layer)) if cache_pool_id == hp_pool_id: raise RuntimeError("FP4 MLA full-history and HP state must use distinct lifecycles.") - bf16_pool_id = None - if self._bf16_mla_manager_layer_ids: - bf16_layer = self._bf16_mla_manager_layer_ids[0] - bf16_pool_id = int(self.impl.get_layer_group_id(bf16_layer)) - if bf16_pool_id in (cache_pool_id, hp_pool_id): - raise RuntimeError("FP4 MLA, BF16 MLA, and HP state must use distinct lifecycles.") num_pools = len(self.impl.layer_grouping) pool_pointers = [[[0, 0], [0, 0]] for _ in range(num_pools)] @@ -469,18 +242,6 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: [int(self.impl.get_mem_pool_base_address(hp_layer, Role.MLA_HP_TAIL)), 0], [0, 0], ] - if bf16_pool_id is not None: - pool_pointers[bf16_pool_id] = [ - [ - int( - self.impl.get_mem_pool_base_address( - self._bf16_mla_manager_layer_ids[0], Role.KEY - ) - ), - 0, - ], - [0, 0], - ] self.kv_cache_pool_pointers = torch.tensor( pool_pointers, dtype=torch.int64, @@ -488,18 +249,10 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: pin_memory=prefer_pinned(), ) - mapping = [[0, 0] for _ in range(self.num_local_layers)] - for compact_layer, layer_id in enumerate(self._cache_manager_layer_ids): + mapping = [] + for layer_id in self._cache_manager_layer_ids: converter = self.impl.get_page_index_converter(layer_id, Role.KEY) - local_layer = self._fp4_mla_compact_to_local[compact_layer] - mapping[local_layer] = [cache_pool_id, int(converter.layer_offset)] - if bf16_pool_id is not None: - for layer_id in self._bf16_mla_manager_layer_ids: - converter = self.impl.get_page_index_converter(layer_id, Role.KEY) - mapping[int(layer_id)] = [ - bf16_pool_id, - int(converter.layer_offset), - ] + mapping.append([cache_pool_id, int(converter.layer_offset)]) self.kv_cache_pool_mapping = torch.tensor( mapping, dtype=torch.int32, @@ -516,12 +269,6 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: self.index_scales[hp_pool_id] = int( self.impl.get_page_index_converter(hp_layer, Role.MLA_HP_TAIL).scale ) - if bf16_pool_id is not None: - self.index_scales[bf16_pool_id] = int( - self.impl.get_page_index_converter( - self._bf16_mla_manager_layer_ids[0], Role.KEY - ).scale - ) self.kv_offset = torch.zeros_like(self.index_scales) self._index_scale_ints = self.index_scales.tolist() self.num_attention_op_pools = num_pools @@ -548,22 +295,10 @@ def _validate_fp4_mla_layer_groups(self) -> None: "FP4 MLA V2 requires one full-history and one HP sliding layer group; " f"got cache={sorted(cache_groups)}, hp={sorted(hp_groups)}." ) - bf16_groups = { - int(self.impl.get_layer_group_id(layer_id)) - for layer_id in self._bf16_mla_manager_layer_ids - } - if self._bf16_mla_manager_layer_ids and ( - len(bf16_groups) != 1 or bf16_groups == cache_groups or bf16_groups == hp_groups - ): - raise RuntimeError( - "BF16 MLA fallback layers require one lifecycle distinct from " - f"FP4 and HP; got bf16={sorted(bf16_groups)}." - ) cache_roles = [Role.KEY, Role.KEY_BLOCK_SCALE, Role.MLA_V_SCALE] if self.mla_v_head_dim is not None: cache_roles.append(Role.MLA_V_PACKED) - for compact_layer, layer_id in enumerate(self._cache_manager_layer_ids): - local_layer = self._fp4_mla_compact_to_local[compact_layer] + for local_layer, layer_id in enumerate(self._cache_manager_layer_ids): converters = [ self.impl.get_page_index_converter(layer_id, role) for role in cache_roles ] @@ -581,14 +316,12 @@ def _validate_fp4_mla_layer_groups(self) -> None: f"one encoded page geometry for local layer {local_layer}; " f"got {sorted(geometries)}." ) - num_fp4_layers = len(self._hp_manager_layer_ids) - for compact_layer, layer_id in enumerate(self._hp_manager_layer_ids): - local_layer = self._fp4_mla_compact_to_local[compact_layer] + for local_layer, layer_id in enumerate(self._hp_manager_layer_ids): converter = self.impl.get_page_index_converter(layer_id, Role.MLA_HP_TAIL) if ( - int(converter.scale) != num_fp4_layers + int(converter.scale) != self.num_local_layers or int(converter.expansion) != 1 - or int(converter.layer_offset) != compact_layer + or int(converter.layer_offset) != local_layer ): raise RuntimeError( "FP4 MLA V2 HP roles require one coalesced page per local " @@ -596,14 +329,9 @@ def _validate_fp4_mla_layer_groups(self) -> None: ) def get_fp4_mla_page_table_spec(self, layer_idx: Optional[int] = None) -> Fp4MlaPageTableSpec: - local_layer = ( - self._fp4_mla_compact_to_local[0] - if layer_idx is None - else self.layer_offsets[layer_idx] - ) - compact_layer = self._fp4_mla_compact_layer_idx(local_layer) - cache_layer = self._cache_manager_layer_ids[compact_layer] - hp_layer = self._hp_manager_layer_ids[compact_layer] + local_layer = 0 if layer_idx is None else self.layer_offsets[layer_idx] + cache_layer = self._cache_manager_layer_ids[local_layer] + hp_layer = self._hp_manager_layer_ids[local_layer] return Fp4MlaPageTableSpec( cache_pool_id=int(self.impl.get_layer_group_id(cache_layer)), # copy_batch_block_offsets already applies the V2 converter scale. @@ -670,8 +398,7 @@ def get_fp4_mla_cache_buffers( if kv_layout != "NHD": raise ValueError("FP4 MLA V2 cache buffers support only NHD layout.") local_layer = self.layer_offsets[layer_idx] - compact_layer = self._fp4_mla_compact_layer_idx(local_layer) - manager_layer = self._cache_manager_layer_ids[compact_layer] + manager_layer = self._cache_manager_layer_ids[local_layer] storage_head_dim = self._storage_head_dim(local_layer) kv_cache = self._role_view( manager_layer, @@ -687,55 +414,6 @@ def get_fp4_mla_cache_buffers( ) return kv_cache, sf_cache - def _iter_fp4_mla_physical_pool_views(self) -> Iterable[torch.Tensor]: - """Yield page-major views spanning each physical MLA pool once.""" - local_layer = self._fp4_mla_compact_to_local[0] - cache_layer = self._cache_manager_layer_ids[0] - storage_head_dim = self._storage_head_dim(local_layer) - yield self._role_pool_base_view( - cache_layer, - Role.KEY, - torch.uint8, - (self.kv_factor, self.tokens_per_block, 1, storage_head_dim // 2), - ) - yield self._role_pool_base_view( - cache_layer, - Role.KEY_BLOCK_SCALE, - torch.uint8, - (self.tokens_per_block, storage_head_dim // FP4_BLOCK_SIZE), - ).view(torch.float8_e4m3fn) - if self._bf16_mla_manager_layer_ids: - bf16_layer = self._bf16_mla_manager_layer_ids[0] - bf16_local_layer = int(bf16_layer) - yield self._role_pool_base_view( - bf16_layer, - Role.KEY, - torch.bfloat16, - ( - self.kv_factor, - self.tokens_per_block, - self.num_kv_heads_per_layer[bf16_local_layer], - self.head_dim_per_layer[bf16_local_layer], - ), - ) - yield self.get_mla_v_scale_pool_base().view(torch.float8_e4m3fn) - if self.mla_v_head_dim is not None: - yield self._role_pool_base_view( - cache_layer, - Role.MLA_V_PACKED, - torch.uint8, - (self.mla_v_head_dim, self.tokens_per_block // 2), - ) - yield self._role_pool_base_view( - self._hp_manager_layer_ids[0], - Role.MLA_HP_TAIL, - torch.bfloat16, - ( - 1, - self._fp4_mla_hp_pool_size * self.head_dim_per_layer[local_layer], - ), - ) - def _all_layer_role_view( self, layer_ids: list[LayerId], @@ -825,16 +503,14 @@ def get_mla_v_scale_pool_base(self) -> torch.Tensor: ) def get_mla_v_scale_page_offset(self, local_layer: int) -> int: - if not 0 <= local_layer < len(self._cache_manager_layer_ids): - raise IndexError(f"Invalid compact FP4 MLA layer {local_layer}.") layer_id = self._cache_manager_layer_ids[local_layer] return int(self.impl.get_page_index_converter(layer_id, Role.MLA_V_SCALE).layer_offset) def get_mla_v_packed_pool(self, local_layer: int) -> Optional[torch.Tensor]: if self.mla_v_head_dim is None: return None - if not 0 <= local_layer < len(self._cache_manager_layer_ids): - raise IndexError(f"Invalid compact FP4 MLA layer {local_layer}.") + if not 0 <= local_layer < self.num_local_layers: + raise IndexError(f"Invalid FP4 MLA local layer {local_layer}.") pool = self._role_view( self._cache_manager_layer_ids[local_layer], Role.MLA_V_PACKED, @@ -855,8 +531,6 @@ def get_mla_v_packed_pool_base(self) -> Optional[torch.Tensor]: return pool.reshape(pool.shape[0] * pool.shape[1], pool.shape[2]) def get_mla_v_packed_page_offset(self, local_layer: int) -> int: - if not 0 <= local_layer < len(self._cache_manager_layer_ids): - raise IndexError(f"Invalid compact FP4 MLA layer {local_layer}.") layer_id = self._cache_manager_layer_ids[local_layer] return int(self.impl.get_page_index_converter(layer_id, Role.MLA_V_PACKED).layer_offset) @@ -865,64 +539,20 @@ def get_fp4_mla_hp_pool(self) -> torch.Tensor: self._hp_manager_layer_ids, Role.MLA_HP_TAIL, torch.bfloat16, - ( - 1, - self._fp4_mla_hp_pool_size - * self.head_dim_per_layer[self._fp4_mla_compact_to_local[0]], - ), + (1, self._fp4_mla_hp_pool_size * self.head_dim_per_layer[0]), ).permute(1, 0, 2, 3) - def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: - return { - Role.ALL: MapperKind.NHD, - Role.MLA_HP_TAIL: MapperKind.REPLICATED, - } - - def get_disagg_transfer_roles(self) -> Optional[frozenset[DataRole]]: - return frozenset((Role.KEY, Role.KEY_BLOCK_SCALE, Role.MLA_HP_TAIL)) - - def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: - local_layer_ids = list(self.impl.layer_grouping[layer_group_id]) - cache_local = { - int(layer_id): self._fp4_mla_compact_to_local[compact_layer] - for compact_layer, layer_id in enumerate(self._cache_manager_layer_ids) - } - cache_local.update( - {int(layer_id): int(layer_id) for layer_id in self._bf16_mla_manager_layer_ids} - ) - hp_local = { - int(layer_id): self._fp4_mla_compact_to_local[compact_layer] - for compact_layer, layer_id in enumerate(self._hp_manager_layer_ids) - } - result = [] - for layer_id in local_layer_ids: - internal_layer = int(layer_id) - if internal_layer in cache_local: - result.append(2 * int(self.pp_layers[cache_local[internal_layer]])) - elif internal_layer in hp_local: - result.append(2 * int(self.pp_layers[hp_local[internal_layer]]) + 1) - else: - raise ValueError(f"Unknown FP4 MLA V2 internal layer {internal_layer}.") - return result - def _get_runtime_cache_size_layer_components(self): - fp4_local_layers = self._fp4_mla_local_layer_indices() - bf16_local_layers = self._bf16_mla_local_layer_indices() sizes = [ self.get_layer_bytes_per_token(local_layer, Role.ALL) - for local_layer in fp4_local_layers + for local_layer in range(self.num_local_layers) ] - windows: list[Optional[int]] = [None] * len(fp4_local_layers) - sizes.extend( - self.get_layer_bytes_per_token(local_layer, Role.ALL) - for local_layer in bf16_local_layers - ) - windows.extend([None] * len(bf16_local_layers)) + windows: list[Optional[int]] = [None] * self.num_local_layers sizes.extend( self.get_layer_bytes_per_token(local_layer, Role.MLA_HP_TAIL) - for local_layer in fp4_local_layers + for local_layer in range(self.num_local_layers) ) - windows.extend([self._fp4_mla_hp_pool_size] * len(fp4_local_layers)) + windows.extend([self._fp4_mla_hp_pool_size] * self.num_local_layers) return sizes, windows def _get_generation_request_capacity(self) -> int: @@ -958,36 +588,11 @@ def get_cache_size_per_token(model_config, mapping, num_layers=None, **kwargs): local_layers = KVCacheManager._resolve_num_attention_layers( model_config, mapping, num_layers ) - bf16_layer_ids = get_kimi_k3_bf16_kv_layer_ids(config) - if bf16_layer_ids: - local_global_layers = set(mapping.pp_layers(int(config.num_hidden_layers))) - num_bf16_layers = len(bf16_layer_ids & local_global_layers) - if num_bf16_layers > local_layers: - raise ValueError( - "Kimi K3 BF16 KV-cache fallback layer count exceeds the local MLA layer count." - ) - else: - num_bf16_layers = 0 - num_fp4_layers = local_layers - num_bf16_layers spec_config = kwargs.get("spec_config") rewind = int(spec_config.tokens_per_gen_step - 1) if spec_config else 0 - hp_bytes = (HP_BLOCK_SIZE + rewind) * logical_head_dim * 2 * num_fp4_layers + hp_bytes = (HP_BLOCK_SIZE + rewind) * logical_head_dim * 2 * local_layers max_batch_size = int(kwargs.get("max_batch_size") or 0) - cache_bytes = ( - per_layer * num_fp4_layers - + logical_head_dim - * torch.empty((), dtype=torch.bfloat16).element_size() - * num_bf16_layers - ) - return cache_bytes, hp_bytes * max_batch_size * mapping.pp_size - - -class Fp4MlaKVCacheManagerV2(Fp4MlaV2CacheLayoutPolicy, KVCacheManagerV2): - """Compatibility manager applying the FP4 MLA V2 layout policy.""" + return per_layer * local_layers, hp_bytes * max_batch_size * mapping.pp_size -__all__ = [ - "Fp4MlaKVCacheManagerV2", - "Fp4MlaPageTableSpec", - "Fp4MlaV2CacheLayoutPolicy", -] +__all__ = ["Fp4MlaKVCacheManagerV2", "Fp4MlaPageTableSpec"] diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index cb6466258562..941c2460ed0e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -188,6 +188,26 @@ def get_kv_cache_manager_cls( config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) + quant_config = model_config.quant_config + if (is_mla(config) and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache()): + if is_disagg: + raise NotImplementedError( + "FP4 MLA disaggregated serving requires the follow-up " + "Python NIXL integration.") + if is_hybrid_linear(config): + raise NotImplementedError( + "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " + "support hybrid linear-attention models.") + if sparse_attn_config is not None: + raise NotImplementedError( + "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " + f"support sparse attention algorithm {sparse_attn_algorithm!r}." + ) + from ..attention.backends.fp4_mla.cache_manager import \ + Fp4MlaKVCacheManagerV2 + + return Fp4MlaKVCacheManagerV2 use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True if is_hybrid_linear(config): # Degenerate case: model is flagged as hybrid but the config has zero @@ -886,6 +906,13 @@ def _validate_or_fallback_kv_cache_manager_v2( f"Gemma4 hybrid attention requires KVCacheManagerV2, " f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") + quant_config = model_config.quant_config + if (is_mla(config) and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache()): + raise NotImplementedError( + "FP4 MLA requires Fp4MlaKVCacheManagerV2, which is " + f"not yet supported with {incompat_str}. Disable these " + "features to run FP4 MLA.") if is_hybrid_linear(config): raise NotImplementedError( "Hybrid Mamba cache managers do not support " diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c8386d9de9ce..72e56fa84797 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6105,6 +6105,14 @@ def _forward_scheduled(self, scheduled_requests: ScheduledRequests, LlmRequest]]): assert not self._is_packed_runner, ( "a packed-batch runner cannot execute scheduled requests") + if not self._disable_overlap_scheduler: + # Do not refill reusable host staging while the previous + # iteration's asynchronous H2D copies still consume it. This + # event precedes the model forward, so its synchronization retains + # GPU forward overlap while protecting every forward entry, + # including speculative/draft paths. + self.wait_for_input_copy() + kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) if self._runner is not None and not self._is_encoder_decoder_model(): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index a5ddf0b68d43..0f34152d13a7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -59,6 +59,8 @@ _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS_STR = "/".join( f"SM{sm_version}" for sm_version in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) +FLASH_MLA_TOKENS_PER_BLOCK = 64 +FP4_MLA_TOKENS_PER_BLOCK = 128 class _ExecutorMemoryMonitor: @@ -197,6 +199,68 @@ def _set_model_engines_cache_reuse(model_engines, cache_reuse: bool): engine.attn_runtime_features.cache_reuse = cache_reuse +def _has_fp4_kv_cache(model_config, kv_cache_config) -> bool: + kv_cache_quant_algo = getattr(getattr(model_config, "quant_config", None), + "kv_cache_quant_algo", None) + fp4_quant_values = { + QuantAlgo.NVFP4, + getattr(QuantAlgo.NVFP4, "value", None), + "NVFP4", + } + kv_cache_dtype = getattr(kv_cache_config, "dtype", None) + return ((isinstance(kv_cache_dtype, str) + and kv_cache_dtype.lower() == "nvfp4") + or (isinstance(kv_cache_quant_algo, str) + and kv_cache_quant_algo.upper() == "NVFP4") + or kv_cache_quant_algo in fp4_quant_values) + + +def _uses_fp4_mla_attention(model_config, kv_cache_config) -> bool: + return (_has_fp4_kv_cache(model_config, kv_cache_config) + and getattr(model_config, "attn_backend", None) == "TRTLLM") + + +def _select_mla_tokens_per_block(config, model_config, kv_cache_config, + tokens_per_block: int) -> int: + if not is_mla(config): + return tokens_per_block + + if _uses_fp4_mla_attention(model_config, kv_cache_config): + tokens_per_block = FP4_MLA_TOKENS_PER_BLOCK + logger.info( + f"Change tokens_per_block to: {tokens_per_block} for using FP4 MLA attention" + ) + kv_cache_config.tokens_per_block = tokens_per_block + return tokens_per_block + + if model_config.enable_flash_mla: + tokens_per_block = FLASH_MLA_TOKENS_PER_BLOCK + logger.info( + f"Change tokens_per_block to: {tokens_per_block} for using FlashMLA" + ) + kv_cache_config.tokens_per_block = tokens_per_block + + return tokens_per_block + + +def _configure_fp4_mla_speculative_kv_cache(spec_config, config, model_config, + kv_cache_config, model) -> None: + if (spec_config is None or not spec_config.spec_dec_mode.use_one_engine() + or not is_mla(config) + or not _uses_fp4_mla_attention(model_config, kv_cache_config)): + return + + # FP4 MLA side pools live on attention metadata and are not switched by + # draft_kv_cache_context. Keep target and one-model draft layers in one + # manager so their global layer IDs address distinct HP/V-scale state. + spec_config._allow_separate_draft_kv_cache = False + if hasattr(model, 'use_separate_draft_kv_cache'): + model.use_separate_draft_kv_cache = False + spec_worker = getattr(model, 'spec_worker', None) + if spec_worker is not None: + spec_worker.use_separate_draft_kv_cache = False + + def _get_mapping(_mapping: Mapping) -> Mapping: if _mapping is None: mapping = Mapping(world_size=tensorrt_llm.mpi_world_size(), @@ -634,6 +698,13 @@ def allocation_scope(current_stage: ExecutorMemoryType): resolve_cache_transceiver_config(cache_transceiver_config) config = model_engine.model.model_config.pretrained_config + _configure_fp4_mla_speculative_kv_cache( + spec_config, + config, + model_engine.model.model_config, + kv_cache_config, + model_engine.model, + ) max_num_seq_slots = getattr( model_engine, "max_num_seq_slots", None) or compute_max_num_sequences( mapping, @@ -641,24 +712,10 @@ def allocation_scope(current_stage: ExecutorMemoryType): llm_args.disable_overlap_scheduler, enable_overlap_headroom=getattr(model_engine, "_enable_overlap_headroom", False)) + tokens_per_block = _select_mla_tokens_per_block( + config, model_engine.model.model_config, kv_cache_config, + tokens_per_block) if is_mla(config): - if model_engine.model.model_config.enable_flash_mla: - tokens_per_block = 64 - # Propagate the override back to kv_cache_config so any consumer - # that later reads llm_args.kv_cache_config.tokens_per_block sees - # the effective value. KvCacheConnectorScheduler subclasses - # (LMCache, Dynamo KVBM) are instantiated further down via - # scheduler_cls(llm_args) and size their block pools from - # llm_args.kv_cache_config.tokens_per_block. Without this the - # connector's block size desynced from the KVCacheManager's - # actual tokens_per_block (user-set or default 32 vs. FlashMLA's - # forced 64), producing a frozen cache_block_ids view to the - # connector and silently-corrupted decode KV (#13320). - kv_cache_config.tokens_per_block = tokens_per_block - logger.info( - f"Change tokens_per_block to: {tokens_per_block} for using FlashMLA" - ) - sm_version = get_sm_version() if (kv_cache_config.enable_block_reuse and sm_version not in _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS): diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index c755dda60358..59553f235cc8 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -398,6 +398,10 @@ def __init__( self.mapping = mapping self.dtype = dtype self.kv_cache_type = kv_cache_type + if dtype == DataType.NVFP4 and kv_cache_type == CacheTypeCpp.SELFKONLY: + raise ValueError( + "NVFP4 SELFKONLY cache storage requires " + "Fp4MlaKVCacheManagerV2; KVCacheManager V1 is not supported.") # Consumed by the disaggregation page-table builder to expose the DSA # indexer K cache pool as a REPLICATED pool view. self.enable_indexer_k_cache = enable_indexer_k_cache @@ -1576,6 +1580,12 @@ def get_cache_size_per_token(model_config: ModelConfigPython, # get head dim mla = hasattr(config, "kv_lora_rank") and config.kv_lora_rank is not None + quant_config = model_config.quant_config + if (mla and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache()): + raise ValueError( + "FP4 MLA cache sizing requires Fp4MlaKVCacheManagerV2; " + "KVCacheManager V1 is not supported.") if mla: head_dim = config.kv_lora_rank + config.qk_rope_head_dim kv_factor = 1 @@ -1592,7 +1602,6 @@ def get_cache_size_per_token(model_config: ModelConfigPython, # K and V mem_per_token = kv_factor * num_attention_layers * head_dim # The data type bytes. - quant_config = model_config.quant_config if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( ): mem_per_token *= 1 @@ -1877,17 +1886,24 @@ def get_buffers(self, result = self.impl.get_primary_pool_data(layer_offset) pool = self.get_pool_for_layer(layer_offset) + layer_dtype = pool.dtype if pool else self.dtype layer_head_dim = pool.head_dim if pool else self.head_dim assert kv_layout in ["NHD", "HND"], f"Unsupported kv_layout: {kv_layout}" + + element_per_container = 1 + if layer_dtype == DataType.NVFP4: + element_per_container = 2 + effective_head_dim = layer_head_dim // element_per_container + if kv_layout == "NHD": return result.reshape( result.shape[0], self.kv_factor, self.tokens_per_block, self.num_kv_heads_per_layer[layer_offset], - layer_head_dim, + effective_head_dim, ) else: return result.reshape( @@ -1895,7 +1911,7 @@ def get_buffers(self, self.kv_factor, self.num_kv_heads_per_layer[layer_offset], self.tokens_per_block, - layer_head_dim, + effective_head_dim, ) def get_indexer_k_cache_pool_data(self, layer_idx: int) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index f2cb1e00d237..10caf2c61667 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -945,6 +945,12 @@ def change_attn_metadata(self, num_accepted_tokens: torch.Tensor, attn_metadata.kv_lens_cuda[num_contexts:batch_size].clamp_( min=runtime_draft_len) attn_metadata.on_update_kv_lens() + attn_metadata.update_for_spec_dec() + elif getattr(attn_metadata, "kv_lens_cuda_runtime", None) is not None: + attn_metadata.kv_lens_cuda_runtime[num_contexts:batch_size] -= ( + runtime_draft_len + 1 - + num_accepted_tokens[num_contexts:batch_size]) + attn_metadata.update_for_spec_dec() if attn_metadata.kv_cache_params is not None and not attn_metadata.is_cuda_graph: for i in range(num_contexts, batch_size): From 342559e1482c247a7159b652d120ab1dd427be44 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:40:44 -0700 Subject: [PATCH 03/21] [None][test] add FP4 MLA core coverage Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../defs/accuracy/test_llm_api_pytorch.py | 17 + .../unittest/_torch/attention/backend_case.py | 48 +- .../unittest/_torch/attention/test_fp4_mla.py | 541 +----------------- 3 files changed, 60 insertions(+), 546 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8d5e4ae62676..e3c050662653 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -298,6 +298,23 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, acceptance_length, ) + @skip_no_rubin + @pytest.mark.skip_less_device_memory(60000) + def test_nvfp4_mla_gsm8k(self): + kv_cache_config = KvCacheConfig(dtype="nvfp4", + free_gpu_memory_fraction=0.75) + with LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only", + attn_backend="TRTLLM", + kv_cache_config=kv_cache_config, + max_num_tokens=8192, + max_batch_size=1350, + ) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("enable_chunked_prefill", [False, True]) def test_bfloat16_flashinfer(self, enable_chunked_prefill): diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index b23358200999..e95caab3d2da 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -28,7 +28,11 @@ PredefinedAttentionMask, RopeParams, ) -from tensorrt_llm._torch.attention.backends.utils import create_attention, get_attention_backend +from tensorrt_llm._torch.attention.backends.utils import ( + append_mla_latent_cache, + create_attention, + get_attention_backend, +) from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -461,7 +465,8 @@ def _run_mla_gen_backend( """Run one backend's absorbed-MLA generation; return [nnz_q, heads*kv_lora]. No ``mla_rope_generation`` call (see ``generate_mla_gen_inputs``): ``fused_q`` - is passed as-is and the backend's ``forward`` appends the new latent + MQA. + is passed as-is. The harness prepares TRTLLM's cache and scheduler inputs + directly, while the other backends retain their normal Python cache update. """ AttentionCls = get_attention_backend(backend) H = case.num_heads @@ -490,20 +495,41 @@ def _run_mla_gen_backend( expected_latents = _split_packed_tokens(latent_cache, case.seq_lens) def _forward(metadata, q, q_pe): + forward_args = AttentionForwardArgs( + latent_cache=latent_cache, + q_pe=q_pe, + attention_input_type=AttentionInputType.generation_only, + ) + if backend == "TRTLLM": + num_ctx = metadata.num_contexts + num_gen = metadata.num_generations + gen_q_lens = metadata.seq_lens_cuda[num_ctx : num_ctx + num_gen].to(torch.int32) + gen_kv_lens = metadata.kv_lens_cuda_runtime[num_ctx : num_ctx + num_gen].to(torch.int32) + cu_q_seqlens = torch.zeros(num_gen + 1, dtype=torch.int32, device=q.device) + cu_kv_seqlens = torch.zeros(num_gen + 1, dtype=torch.int32, device=q.device) + cu_q_seqlens[1:] = torch.cumsum(gen_q_lens, dim=0).to(torch.int32) * H + cu_kv_seqlens[1:] = torch.cumsum(gen_kv_lens, dim=0).to(torch.int32) + forward_args.cu_q_seqlens = cu_q_seqlens + forward_args.cu_kv_seqlens = cu_kv_seqlens + forward_args.fmha_scheduler_counter = torch.zeros( + 1, dtype=torch.uint32, device=q.device + ) + append_mla_latent_cache( + metadata.kv_cache_manager, + attn.get_local_layer_idx(metadata), + metadata.request_ids, + metadata.seq_lens.tolist(), + metadata.kv_cache_params.num_cached_tokens_per_seq, + latent_cache, + kv_layout=metadata.kv_layout, + seq_start=num_ctx, + ) out = attn.forward( q, None, None, metadata, - forward_args=AttentionForwardArgs( - latent_cache=latent_cache, - q_pe=q_pe, - attention_input_type=AttentionInputType.generation_only, - # The harness feeds a pre-RoPE'd fused_q, so skip the RoPE step; - # the TRTLLM backend still appends the new latent and inits its - # scheduler buffers. Vanilla/FlashInfer ignore this flag. - skip_mla_rope_generation=True, - ), + forward_args=forward_args, ) return out[0] if isinstance(out, tuple) else out diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index 521ad2554009..e90a69e3b037 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -10,12 +10,10 @@ import tensorrt_llm import tensorrt_llm._torch.attention.backends.fp4_mla as fp4_mla_backend -from tensorrt_llm._torch.attention.backends.fmha.fp4_mla import Fp4MlaFmha from tensorrt_llm._torch.attention.backends.fp4_mla import ( FP4_BLOCK_SIZE, FP4_MLA_ATTENTION_BACKEND_ENV, FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, - FP4_MLA_CUTEDSL_MUFU16_ENV, FP4_MLA_K_RESIDUAL_DIM, FP4_MLA_KV_GLOBAL_SCALE, FP4_MLA_P_GLOBAL_SCALE, @@ -26,27 +24,15 @@ _cutedsl_backend_available, _fp4_mla_attention_backend, _get_fp4_mla_global_scale, - load_fp4_mla_chunked_kv_cache, run_fp4_mla_attention_decode, scatter_fp4_mla_kv_cache, ) from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 -from tensorrt_llm._torch.attention.backends.fp4_mla.fp4_mla_context import ( - _build_fp8_mla_context_metadata, -) -from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - AttentionInputType, -) -from tensorrt_llm._torch.kimi_k3_cache_policy import KIMI_K3_BF16_KV_LAYERS_ENV -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - BASE_GENERATION_TOKEN_COUNT, - KVCacheManagerV2, - Role, -) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexConverter _DataType = tensorrt_llm.bindings.DataType _CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -71,60 +57,6 @@ def _is_cutedsl_unavailable() -> bool: ) -def test_fp4_mla_request_validation_uses_sparse_runtime_params() -> None: - metadata = SimpleNamespace( - num_sparse_topk=0, - kv_cache_manager=SimpleNamespace(dtype=_DataType.NVFP4, kv_factor=1), - high_precision_kv_pool=object(), - fp4_mla_v_scale_pool=object(), - beam_width=1, - ) - forward_args = AttentionForwardArgs(attention_input_type=AttentionInputType.generation_only) - - Fp4MlaFmha._validate_request(None, None, None, metadata, forward_args) - - forward_args.sparse_runtime_params.sparse_attn_indices = torch.tensor([0]) - with pytest.raises(NotImplementedError, match="does not support sparse attention"): - Fp4MlaFmha._validate_request(None, None, None, metadata, forward_args) - - -def test_fp8_mla_context_partition_metadata_separates_q_and_kv_lengths() -> None: - prompt_lens_cuda = torch.tensor([32, 48, 1], dtype=torch.int32) - prompt_lens_cpu = prompt_lens_cuda.clone() - kv_lens_cuda = torch.tensor([128, 64], dtype=torch.int32) - kv_lens_cpu = kv_lens_cuda.clone() - meta = SimpleNamespace( - num_contexts=2, - num_ctx_tokens=80, - prompt_lens_cuda_runtime=prompt_lens_cuda, - prompt_lens_cpu_runtime=prompt_lens_cpu, - host_request_types_runtime=torch.tensor([0, 0, 1], dtype=torch.int32), - positions=torch.arange(80, dtype=torch.int32), - _fp4_mla_fp8_context_state=object(), - ) - scratch = SimpleNamespace( - cache_manager_view=object(), - block_offsets=object(), - block_ids_per_seq=object(), - host_total_kv_lens=torch.tensor([192, 0], dtype=torch.int64), - ) - - fp8_meta = _build_fp8_mla_context_metadata( - meta, - scratch, - kv_lens_cuda=kv_lens_cuda, - kv_lens_cpu=kv_lens_cpu, - ) - - assert fp8_meta.prompt_lens_cuda_runtime.data_ptr() == prompt_lens_cuda.data_ptr() - assert fp8_meta.prompt_lens_cpu_runtime.data_ptr() == prompt_lens_cpu.data_ptr() - assert fp8_meta.kv_lens_cuda_runtime is kv_lens_cuda - assert fp8_meta.kv_lens_runtime is kv_lens_cpu - assert fp8_meta.host_total_kv_lens is scratch.host_total_kv_lens - assert fp8_meta.helix_position_offsets.shape == (80,) - assert fp8_meta._fp4_mla_fp8_context_state is None - - def _reset_triton_allocator() -> None: import triton @@ -198,7 +130,7 @@ def populate_generation_lengths(*args, **kwargs) -> None: def test_fp4_mla_v2_encoded_page_capacity_is_layer_invariant(layer_offset: int) -> None: manager = object.__new__(Fp4MlaKVCacheManagerV2) manager.impl = SimpleNamespace( - get_page_index_converter=lambda *_: SimpleNamespace( + get_page_index_converter=lambda *_: PageIndexConverter( scale=4, expansion=1, layer_offset=layer_offset, @@ -209,11 +141,8 @@ def test_fp4_mla_v2_encoded_page_capacity_is_layer_invariant(layer_offset: int) assert manager._role_encoded_page_capacity(0, Role.KEY) == 17 -def test_fp4_mla_v2_cache_size_accounts_for_hp_intercept( - monkeypatch: pytest.MonkeyPatch, -) -> None: +def test_fp4_mla_v2_cache_size_accounts_for_hp_intercept(monkeypatch) -> None: monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") - monkeypatch.delenv(KIMI_K3_BF16_KV_LAYERS_ENV, raising=False) model_config = SimpleNamespace( pretrained_config=SimpleNamespace( kv_lora_rank=512, @@ -234,35 +163,6 @@ def test_fp4_mla_v2_cache_size_accounts_for_hp_intercept( assert intercept == 3 * 2 * (HP_BLOCK_SIZE + 3) * 576 * 2 * 2 -def test_fp4_mla_v2_cache_size_accounts_for_bf16_fallback( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") - monkeypatch.setenv(KIMI_K3_BF16_KV_LAYERS_ENV, "3") - model_config = SimpleNamespace( - pretrained_config=SimpleNamespace( - model_type="kimi_linear", - num_hidden_layers=4, - kv_lora_rank=512, - qk_rope_head_dim=64, - linear_attn_config={"full_attn_layers": [2, 4]}, - ) - ) - - slope, intercept = Fp4MlaKVCacheManagerV2.get_cache_size_per_token( - model_config, - Mapping(world_size=1, tp_size=1, pp_size=1, rank=0), - num_layers=2, - tokens_per_block=FP4_MLA_TOKENS_PER_BLOCK, - max_batch_size=3, - spec_config=SimpleNamespace(tokens_per_gen_step=4), - ) - - fp4_bytes_per_token = 320 + 40 + 32 - assert slope == fp4_bytes_per_token + 576 * torch.bfloat16.itemsize - assert intercept == 3 * (HP_BLOCK_SIZE + 3) * 576 * 2 - - def test_fp4_mla_v2_runtime_sizing_accounts_for_pipeline_slots() -> None: manager = object.__new__(Fp4MlaKVCacheManagerV2) manager.max_batch_size = 3 @@ -270,18 +170,12 @@ def test_fp4_mla_v2_runtime_sizing_accounts_for_pipeline_slots() -> None: manager.max_num_tokens = 100 manager.tokens_per_block = FP4_MLA_TOKENS_PER_BLOCK manager.enable_swa_scratch_reuse = False - manager._has_cp_helix = False - manager._generation_kv_capacity_headroom = BASE_GENERATION_TOKEN_COUNT manager._get_runtime_cache_size_layer_components = lambda: ([10, 8], [None, 19]) quota = manager._get_quota_from_max_tokens(manager.max_num_tokens) - # KVCacheManagerV2 charges every resident generation request the pages that - # the live window interval can straddle, which is two 128-token pages for a - # 19-token HP ring plus the base generation headroom. hp_page_bytes = FP4_MLA_TOKENS_PER_BLOCK * 8 - hp_pages_per_request = 2 - expected_quota = manager.max_num_tokens * (10 + 8) + 6 * hp_pages_per_request * hp_page_bytes + expected_quota = manager.max_num_tokens * (10 + 8) + 6 * hp_page_bytes assert quota == expected_quota assert manager._get_max_tokens_from_quota(quota) == manager.max_num_tokens @@ -322,201 +216,6 @@ def capture_base_init(manager, *args, **kwargs) -> None: assert captured["manager"].mla_v_head_dim == expected_v_head_dim -@pytest.mark.parametrize( - "fused_v_transpose", - [False, True], - ids=["mufu16-packed-v", "fused-v-canonical-cache"], -) -def test_fp4_mla_disagg_import_rebuilds_variant_sidecars( - monkeypatch, - fused_v_transpose: bool, -) -> None: - page_ids = torch.tensor([17, 29], dtype=torch.int32) - page_valid_tokens = torch.tensor([FP4_MLA_TOKENS_PER_BLOCK, 1], dtype=torch.int32) - v_scale_pool = torch.empty((2, 32, 1), dtype=torch.uint8) - kv_caches = { - 7: torch.empty(1, dtype=torch.uint8), - 11: torch.empty(1, dtype=torch.uint8), - } - sf_caches = { - 7: torch.empty(1, dtype=torch.uint8), - 11: torch.empty(1, dtype=torch.uint8), - } - v_packed_pools = { - 0: torch.empty(1, dtype=torch.uint8), - 1: torch.empty(1, dtype=torch.uint8), - } - staged = [] - scale_rebuilds = [] - packed_rebuilds = [] - - def stage_page_metadata(block_ids, *, prompt_len, page_size, device): - staged.append((block_ids, prompt_len, page_size, device)) - return page_ids, page_valid_tokens - - def rebuild_v_scales( - sf_cache, - observed_v_scale_pool, - observed_page_ids, - observed_page_valid_tokens, - **kwargs, - ) -> None: - scale_rebuilds.append( - ( - sf_cache, - observed_v_scale_pool, - observed_page_ids, - observed_page_valid_tokens, - kwargs, - ) - ) - - def rebuild_packed_v( - v_packed, - kv_cache, - observed_page_ids, - **kwargs, - ) -> None: - packed_rebuilds.append((v_packed, kv_cache, observed_page_ids, kwargs)) - - def get_v_packed_pool(local_layer: int) -> torch.Tensor: - if fused_v_transpose: - raise AssertionError("fused-V import must not request a packed-V pool") - return v_packed_pools[local_layer] - - manager = SimpleNamespace( - dtype=_DataType.NVFP4, - kv_factor=1, - mla_v_scale_head_dim=512, - mla_v_head_dim=None if fused_v_transpose else 512, - tokens_per_block=FP4_MLA_TOKENS_PER_BLOCK, - pp_layers=[7, 11], - num_local_layers=2, - get_fp4_mla_page_table_spec=lambda *_: SimpleNamespace(), - get_batch_cache_indices=lambda request_ids, layer_idx=None: [[17, 29, 31]], - get_mla_v_scale_pool=lambda: v_scale_pool, - get_fp4_mla_cache_buffers=lambda layer_idx: ( - kv_caches[layer_idx], - sf_caches[layer_idx], - ), - get_mla_v_packed_pool=get_v_packed_pool, - ) - monkeypatch.setattr(fp4_mla_backend, "_fp4_mla_attention_backend", lambda: "cutedsl") - monkeypatch.setattr( - fp4_mla_backend, - "_fp4_mla_cutedsl_fused_v_transpose_enabled", - lambda: fused_v_transpose, - ) - monkeypatch.setattr( - fp4_mla_backend, - "_stage_fp4_mla_import_page_metadata", - stage_page_metadata, - ) - monkeypatch.setattr( - fp4_mla_backend, - "_rebuild_fp4_mla_v_scales_from_k_scales", - rebuild_v_scales, - ) - monkeypatch.setattr( - fp4_mla_backend, - "_repack_cutedsl_v_packed_cache", - rebuild_packed_v, - ) - - assert fp4_mla_backend.rebuild_fp4_mla_disagg_imported_cache( - manager, - request_id=41, - prompt_len=FP4_MLA_TOKENS_PER_BLOCK + 1, - ) - - assert staged == [ - ( - [17, 29], - FP4_MLA_TOKENS_PER_BLOCK + 1, - FP4_MLA_TOKENS_PER_BLOCK, - v_scale_pool.device, - ) - ] - assert len(scale_rebuilds) == 2 - for local_layer, layer_idx in enumerate(manager.pp_layers): - sf_cache, observed_pool, observed_ids, observed_valid, kwargs = scale_rebuilds[local_layer] - assert sf_cache is sf_caches[layer_idx] - assert observed_pool is v_scale_pool - assert observed_ids is page_ids - assert observed_valid is page_valid_tokens - assert kwargs == { - "local_layer": local_layer, - "v_head_dim": 512, - "page_size": FP4_MLA_TOKENS_PER_BLOCK, - } - assert len(packed_rebuilds) == (0 if fused_v_transpose else 2) - for local_layer, layer_idx in enumerate(manager.pp_layers[: len(packed_rebuilds)]): - v_packed, kv_cache, observed_ids, kwargs = packed_rebuilds[local_layer] - assert v_packed is v_packed_pools[local_layer] - assert kv_cache is kv_caches[layer_idx] - assert observed_ids is page_ids - assert kwargs == { - "v_head_dim": 512, - "page_size": FP4_MLA_TOKENS_PER_BLOCK, - "block_v": fp4_mla_backend.FP4_MLA_SCALE_ROW_GROUP, - } - - -def test_fp4_mla_disagg_import_skips_hybrid_linear_attention_layers(monkeypatch) -> None: - page_ids = torch.tensor([17, 29], dtype=torch.int32) - page_valid_tokens = torch.tensor([FP4_MLA_TOKENS_PER_BLOCK, 1], dtype=torch.int32) - v_scale_pool = torch.empty((2, 32, 1), dtype=torch.uint8) - requested_page_layers = [] - rebuilt_cache_layers = [] - rebuilt_compact_layers = [] - - def get_batch_cache_indices(request_ids, *, layer_idx): - requested_page_layers.append((request_ids, layer_idx)) - return [[17, 29]] - - def get_fp4_mla_cache_buffers(layer_idx): - rebuilt_cache_layers.append(layer_idx) - return torch.empty(1, dtype=torch.uint8), torch.empty(1, dtype=torch.uint8) - - def rebuild_v_scales(_sf_cache, _v_scale_pool, _page_ids, _page_valid_tokens, **kwargs): - rebuilt_compact_layers.append(kwargs["local_layer"]) - - manager = SimpleNamespace( - dtype=_DataType.NVFP4, - kv_factor=1, - mla_v_scale_head_dim=512, - tokens_per_block=FP4_MLA_TOKENS_PER_BLOCK, - pp_layers=[3, 7, 11], - num_local_layers=3, - _fp4_mla_compact_to_local=[1, 2], - get_fp4_mla_page_table_spec=lambda *_: SimpleNamespace(), - get_batch_cache_indices=get_batch_cache_indices, - get_mla_v_scale_pool=lambda: v_scale_pool, - get_fp4_mla_cache_buffers=get_fp4_mla_cache_buffers, - ) - monkeypatch.setattr(fp4_mla_backend, "_fp4_mla_attention_backend", lambda: "triton") - monkeypatch.setattr( - fp4_mla_backend, - "_stage_fp4_mla_import_page_metadata", - lambda *args, **kwargs: (page_ids, page_valid_tokens), - ) - monkeypatch.setattr( - fp4_mla_backend, - "_rebuild_fp4_mla_v_scales_from_k_scales", - rebuild_v_scales, - ) - - assert fp4_mla_backend.rebuild_fp4_mla_disagg_imported_cache( - manager, - request_id=41, - prompt_len=FP4_MLA_TOKENS_PER_BLOCK + 1, - ) - - assert requested_page_layers == [([41], 7)] - assert rebuilt_cache_layers == [7, 11] - assert rebuilt_compact_layers == [0, 1] - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_fp4_mla_manager_v2_registers_native_cache_and_hp_roles(monkeypatch) -> None: monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") @@ -564,9 +263,6 @@ def test_fp4_mla_manager_v2_registers_native_cache_and_hp_roles(monkeypatch) -> assert v_scale_pool.dtype == torch.float8_e4m3fn assert manager.get_mla_v_scale_pool_base().dtype == torch.uint8 assert hp_pool.shape[1:] == (2, 1, HP_BLOCK_SIZE * 576) - assert manager.get_disagg_transfer_roles() == frozenset( - (Role.KEY, Role.KEY_BLOCK_SCALE, Role.MLA_HP_TAIL) - ) finally: manager.shutdown() @@ -836,8 +532,6 @@ def _materialize_reference_cache_tokens( batch_indices: torch.Tensor, positions: torch.Tensor, head_dim: int, - *, - include_k_residual: bool = False, ) -> torch.Tensor: kv_cache, sf_cache = metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) sf_cache = sf_cache.view(torch.float8_e4m3fn) @@ -863,112 +557,10 @@ def _materialize_reference_cache_tokens( sf_per_token=storage_head_dim // FP4_BLOCK_SIZE, global_scale=static_global_scale, ) - token = dequantized_pages[physical_page][page_position, :head_dim].clone() - if include_k_residual: - residual_end = head_dim + FP4_MLA_K_RESIDUAL_DIM - token[-FP4_MLA_K_RESIDUAL_DIM:] += dequantized_pages[physical_page][ - page_position, - head_dim:residual_end, - ] - tokens.append(token) + tokens.append(dequantized_pages[physical_page][page_position, :head_dim]) return torch.stack(tokens, dim=0) -@pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), - reason="requires Rubin SM107", -) -def test_fp4_mla_chunked_cache_gather_dequantizes_k_residual(monkeypatch) -> None: - _reset_triton_allocator() - monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") - seq_lens = [300, 180] - kv_lora_rank = 512 - qk_rope_head_dim = 64 - head_dim = kv_lora_rank + qk_rope_head_dim - kv_cache_manager = _create_fp4_mla_v2_manager( - max_tokens=640, - max_seq_len=max(seq_lens), - max_batch_size=len(seq_lens), - ) - try: - kv_cache_manager.add_dummy_requests(list(range(len(seq_lens))), seq_lens) - metadata = _build_multi_seq_metadata( - kv_cache_manager, - seq_lens=seq_lens, - page_size=FP4_MLA_TOKENS_PER_BLOCK, - ) - latent = ( - torch.randn(sum(seq_lens), head_dim, dtype=torch.bfloat16, device="cuda") * 0.25 - ).clamp_(-1.0, 1.0) - scatter_fp4_mla_kv_cache( - metadata, - latent, - layer_idx=0, - token_offset=0, - phase="context", - local_layer=0, - v_head_dim=kv_lora_rank, - ) - - chunk_lens = [96, 80] - chunk_offsets = torch.tensor([128, 64], dtype=torch.int64, device="cuda") - cu_chunk_lens = torch.tensor([0, 96, 176], dtype=torch.int64, device="cuda") - gathered_compressed_kv, gathered_k_pe = load_fp4_mla_chunked_kv_cache( - metadata, - layer_idx=0, - num_ctx_cached_tokens=sum(chunk_lens), - cu_chunked_seq_len=cu_chunk_lens, - chunked_global_offset=chunk_offsets, - chunked_max_seq_len=max(chunk_lens), - out_dtype=torch.bfloat16, - kv_lora_rank=kv_lora_rank, - qk_rope_head_dim=qk_rope_head_dim, - ) - batch_indices = torch.tensor( - [0] * chunk_lens[0] + [1] * chunk_lens[1], - dtype=torch.int32, - device="cuda", - ) - positions = torch.cat( - [ - torch.arange( - offset, - offset + length, - dtype=torch.int32, - device="cuda", - ) - for offset, length in zip(chunk_offsets.tolist(), chunk_lens) - ] - ) - reference = _materialize_reference_cache_tokens( - metadata, - layer_idx=0, - batch_indices=batch_indices, - positions=positions, - head_dim=head_dim, - include_k_residual=True, - ).to(torch.bfloat16) - - torch.testing.assert_close( - gathered_compressed_kv, - reference[:, :kv_lora_rank], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - gathered_k_pe, - reference[:, kv_lora_rank:], - rtol=0, - atol=0, - ) - finally: - torch.cuda.synchronize() - _reset_triton_allocator() - kv_cache_manager.shutdown() - torch.cuda.synchronize() - torch.cuda.empty_cache() - - def _build_fp4_mla_attention_decode_case( *, seq_lens, @@ -1268,7 +860,6 @@ def _assert_fp4_mla_attention_decode_accuracy( ) -> None: _reset_triton_allocator() monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, backend) - monkeypatch.setenv(FP4_MLA_CUTEDSL_MUFU16_ENV, "1") monkeypatch.setenv( FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, str(int(fused_v_transpose)), @@ -1509,126 +1100,6 @@ def test_fp4_mla_context_tail_uses_draft_slack_ring( torch.cuda.empty_cache() -@pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), - reason="requires Rubin SM107", -) -def test_fp4_mla_context_scatter_fuses_qk_rope(monkeypatch) -> None: - _reset_triton_allocator() - monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") - seq_len = 17 - kv_lora_rank = 512 - qk_nope_head_dim = 128 - qk_rope_head_dim = 64 - num_heads = 4 - head_dim = kv_lora_rank + qk_rope_head_dim - kv_cache_manager = _create_fp4_mla_v2_manager( - max_tokens=FP4_MLA_TOKENS_PER_BLOCK, - max_seq_len=FP4_MLA_TOKENS_PER_BLOCK, - max_batch_size=1, - ) - try: - kv_cache_manager.add_dummy_requests([0], [seq_len]) - metadata = _build_multi_seq_metadata( - kv_cache_manager, - seq_lens=[seq_len], - page_size=FP4_MLA_TOKENS_PER_BLOCK, - ) - latent = torch.randn( - seq_len, - head_dim, - dtype=torch.bfloat16, - device="cuda", - ) - q = torch.randn( - seq_len, - num_heads * (qk_nope_head_dim + qk_rope_head_dim), - dtype=torch.bfloat16, - device="cuda", - ) - original_latent = latent.clone() - original_q = q.clone().view( - seq_len, - num_heads, - qk_nope_head_dim + qk_rope_head_dim, - ) - rotary_cos_sin = torch.zeros( - (FP4_MLA_TOKENS_PER_BLOCK, qk_rope_head_dim, 2), - dtype=torch.float32, - device="cuda", - ) - rotary_cos_sin[..., 1] = 1.0 - - scatter_fp4_mla_kv_cache( - metadata, - latent, - layer_idx=0, - token_offset=0, - phase="context", - local_layer=0, - v_head_dim=kv_lora_rank, - rotary_cos_sin=rotary_cos_sin, - q_context=q, - q_nope_head_dim=qk_nope_head_dim, - ) - torch.cuda.synchronize() - - expected_k_pe = original_latent[:, kv_lora_rank:].reshape( - seq_len, - qk_rope_head_dim // 2, - 2, - ) - expected_k_pe = torch.stack( - (-expected_k_pe[..., 1], expected_k_pe[..., 0]), - dim=-1, - ).flatten(1) - q_view = q.view( - seq_len, - num_heads, - qk_nope_head_dim + qk_rope_head_dim, - ) - expected_q_pe = original_q[..., qk_nope_head_dim:].reshape( - seq_len, - num_heads, - qk_rope_head_dim // 2, - 2, - ) - expected_q_pe = torch.stack( - (-expected_q_pe[..., 1], expected_q_pe[..., 0]), - dim=-1, - ).flatten(-2) - torch.testing.assert_close( - latent[:, :kv_lora_rank], - original_latent[:, :kv_lora_rank], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - latent[:, kv_lora_rank:], - expected_k_pe, - rtol=0, - atol=0, - ) - torch.testing.assert_close( - q_view[..., :qk_nope_head_dim], - original_q[..., :qk_nope_head_dim], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - q_view[..., qk_nope_head_dim:], - expected_q_pe, - rtol=0, - atol=0, - ) - finally: - torch.cuda.synchronize() - _reset_triton_allocator() - kv_cache_manager.shutdown() - torch.cuda.synchronize() - torch.cuda.empty_cache() - - _V_REPACK_PAGE_SIZE = 128 _V_REPACK_HEAD_DIM = 512 _V_REPACK_PACKED_DIM = _V_REPACK_HEAD_DIM // 2 From 777a3ff584be98c64d32047a2a0967ee99438fb6 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 31 Aug 2026 05:17:29 -0700 Subject: [PATCH 04/21] Migrate FP4 MLA kernels to public CuTeDSL APIs Signed-off-by: root --- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 104 +++++++-------- ...p4_mla_cutedsl_mufu16_fused_v_transpose.py | 126 +++++++++--------- 2 files changed, 111 insertions(+), 119 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 785f06656c68..caf71204a31a 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -18,19 +18,6 @@ import cutlass.cute as cute import cutlass.pipeline as pipeline import torch -from ctm.Operations import ptx -from ctm.Operations.ptx import ( - AtomicOpKind, - CvtaSpace, - MBarrierArriveScope, - MBarrierArriveSem, - MBarrierSpace, - MemScopeKind, - SharedSpace, - cvta_to, - mbarrier_arrive, -) -from ctm.Operations.ptx import cp_async as _cp_async from cutlass._mlir.dialects import llvm from cutlass.base_dsl.dsl import BaseDSL from cutlass.cute.arch.nvvm_wrappers import inline_ptx as cute_inline_ptx @@ -38,9 +25,9 @@ from cutlass.experimental import cuda as cuda_tma from cutlass.experimental import primitives as prims -nvvm_add_packed_f32x2 = partial(ptx.add_packed_f32x2, rnd="rn") -nvvm_mul_packed_f32x2 = partial(ptx.mul_packed_f32x2, rnd="rn") -nvvm_fma_packed_f32x2 = partial(ptx.fma_packed_f32x2, rnd="rn") +nvvm_add_packed_f32x2 = partial(prims.add_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_mul_packed_f32x2 = partial(prims.mul_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 _CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" @@ -59,7 +46,7 @@ def _as_shared_cta(address, *, loc=None, ip=None): address_ir = address.ir_value() if hasattr(address, "ir_value") else address address_space = llvm.PointerType(address_ir.type).address_space if address_space == ctm.AddressSpace.dsmem: - return cvta_to(address, CvtaSpace.SHARED, loc=loc, ip=ip) + return prims.cvta_to(address, prims.CvtaSpace.SHARED, loc=loc, ip=ip) return address @@ -70,12 +57,10 @@ def _mapa_shared_cluster(address, rank, *, loc=None, ip=None): @ctm.dsl_user_op def _mbarrier_arrive_release_cta_shared_cluster(mbar, count=1, *, loc=None, ip=None) -> None: - mbarrier_arrive( + prims.mbarrier_arrive( mbar, - count, - sem=MBarrierArriveSem.RELEASE, - scope=MBarrierArriveScope.CTA, - space=MBarrierSpace.SHARED_CLUSTER, + count=count, + scope=prims.MemScope.CTA, loc=loc, ip=ip, ) @@ -1511,7 +1496,7 @@ def _tma_gather4_cluster( leader_barrier = _as_shared_cta(_mapa_shared_cluster(barrier, ctm.Int32(0))) barrier_ptr = leader_barrier.data_ptr() multicast_mask_u16 = ctm.Uint16(multicast_mask) - _cp_async._predicated_inline_ptx( + cute_inline_ptx( "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" " [{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], " @@ -2294,7 +2279,11 @@ def _softmax_exp2_pair_packed_f16x2( (score0, score1), (softmax_scale_log2, softmax_scale_log2), (p_bias, p_bias) ) shifted_h2 = _pack_f32_pair_to_f16x2(shifted[0], shifted[1]) - return ptx.ex2_f16x2(shifted_h2) + return cute_inline_ptx( + "ex2.approx.f16x2 {$w0}, {$r0};", + write_only_types=[ctm.Int32], + read_only_args=[shifted_h2], + ) @cute.jit @@ -3591,28 +3580,41 @@ def _smem_p4_p4_materialize_group_col( @cute.jit def _float_to_ordered_u32_for_atomic_max(value: ctm.Float32) -> ctm.Uint32: - bits = ptx.mov_b32(value, target_type=ctm.Int32) + bits = prims.mov_b32(value, target_type=ctm.Int32) sign_mask = bits >> ctm.Int32(31) | ctm.Int32(2147483648) encoded = bits ^ sign_mask - return ptx.mov_b32(encoded, target_type=ctm.Uint32) + return prims.mov_b32(encoded, target_type=ctm.Uint32) @cute.jit def _ordered_u32_to_float_after_atomic_max(value: ctm.Uint32) -> ctm.Float32: - encoded = ptx.mov_b32(value, target_type=ctm.Int32) + encoded = prims.mov_b32(value, target_type=ctm.Int32) sign_mask = ~(encoded >> ctm.Int32(31)) | ctm.Int32(2147483648) bits = encoded ^ sign_mask - return ptx.mov_b32(bits, target_type=ctm.Float32) + return prims.mov_b32(bits, target_type=ctm.Float32) @cute.jit def _smem_atomic_max_ordered_u32(pointer, value: ctm.Uint32) -> None: - ptx.atom( - AtomicOpKind.MAX, + prims.atomicrmw( + prims.AtomicOp.MAX, pointer, value, - syncscope=MemScopeKind.CTA, - space=SharedSpace.shared_cta, + syncscope=prims.MemScope.CTA, + space=prims.SharedSpace.shared_cta, + ) + + +@cute.jit +def _tcgen05_ld_red_32x32b_x16_max_f32(tmem) -> tuple: + tmem_addr = tmem.toint(ctm.Int32) + return cute_inline_ptx( + "tcgen05.ld.red.sync.aligned.32x32b.x16.max.f32 " + "{{$w0}, {$w1}, {$w2}, {$w3}, {$w4}, {$w5}, {$w6}, {$w7}, " + "{$w8}, {$w9}, {$w10}, {$w11}, {$w12}, {$w13}, {$w14}, {$w15}}, " + "{$w16}, [{$r0}];", + write_only_types=[ctm.Int32] * 17, + read_only_args=[tmem_addr], ) @@ -3698,13 +3700,7 @@ def _load_p4_n256_score_half_from_tmem( pending_score_groups.append(group_scores) pending_group_stats.append(None) else: - regs = ptx.tcgen05_ld_red( - ptx.Tcgen05LdStShape.SHAPE_32X32B, - tmem, - num=SF_VEC_SIZE, - red_op="max", - type_="f32", - ) + regs = _tcgen05_ld_red_32x32b_x16_max_f32(tmem) pending_score_groups.append(regs) pending_group_stats.append(regs[SF_VEC_SIZE]) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) @@ -5432,14 +5428,13 @@ def _runtime_t336_producer_tile( ) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - 1, - sem=MBarrierArriveSem.RELAXED, - scope=MBarrierArriveScope.CLUSTER, - space=MBarrierSpace.SHARED_CLUSTER, - pred=common_source_ready_pred, - ) + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -5691,14 +5686,13 @@ def _runtime_t336_producer_tile_v23( warp_rebase = cute.arch.vote_any_sync(lane_rebase) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - 1, - sem=MBarrierArriveSem.RELAXED, - scope=MBarrierArriveScope.CLUSTER, - space=MBarrierSpace.SHARED_CLUSTER, - pred=common_source_ready_pred, - ) + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py index df51457a43c2..1301574752f5 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py @@ -19,18 +19,6 @@ import cutlass.cute as cute import cutlass.pipeline as pipeline import torch -from ctm.Operations import ptx -from ctm.Operations.ptx import ( - AtomicOpKind, - CvtaSpace, - MBarrierArriveScope, - MBarrierArriveSem, - MBarrierSpace, - MemScopeKind, - SharedSpace, - cvta_to, -) -from ctm.Operations.ptx import cp_async as _cp_async from cutlass._mlir.dialects import llvm from cutlass.base_dsl.array import EvictPriority from cutlass.base_dsl.dsl import BaseDSL @@ -39,9 +27,9 @@ from cutlass.experimental import cuda as cuda_tma from cutlass.experimental import primitives as prims -nvvm_add_packed_f32x2 = partial(ptx.add_packed_f32x2, rnd="rn") -nvvm_mul_packed_f32x2 = partial(ptx.mul_packed_f32x2, rnd="rn") -nvvm_fma_packed_f32x2 = partial(ptx.fma_packed_f32x2, rnd="rn") +nvvm_add_packed_f32x2 = partial(prims.add_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_mul_packed_f32x2 = partial(prims.mul_packed_f32x2, rnd=prims.FPRoundingMode.RN) +nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 _CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" @@ -59,7 +47,7 @@ def _as_tma_completion_mbar(mbar, *, loc=None, ip=None): mbar_ir = mbar.ir_value() if hasattr(mbar, "ir_value") else mbar if llvm.PointerType(mbar_ir.type).address_space == ctm.AddressSpace.dsmem: - return cvta_to(mbar, CvtaSpace.SHARED, loc=loc, ip=ip) + return prims.cvta_to(mbar, prims.CvtaSpace.SHARED, loc=loc, ip=ip) return mbar @@ -73,12 +61,10 @@ def _mapa_shared_cluster(mbar, rank, *, loc=None, ip=None): def _mbarrier_arrive_shared_cluster(mbar, count=1, *, loc=None, ip=None) -> None: # Keep the validated release.cta form. The cluster-scoped NVVM form fails # to lower on the target toolchain (CUDA error 715). - ptx.mbarrier_arrive( + prims.mbarrier_arrive( mbar, - count, - sem=MBarrierArriveSem.RELEASE, - scope=MBarrierArriveScope.CTA, - space=MBarrierSpace.SHARED_CLUSTER, + count=count, + scope=prims.MemScope.CTA, loc=loc, ip=ip, ) @@ -1469,7 +1455,7 @@ def _load_qk_qonly_kblock_stage( if cta_rank == ctm.Int32(0): while not prims.mbarrier_try_wait_parity(q_tma_mbar, q_tma_phase, time_limit=10000000): pass - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) @cute.jit @@ -1515,10 +1501,13 @@ def _tma_gather4_cluster( """Gather four KSF atoms and multicast the compact image to both CTAs.""" smem_ptr = smem_dst.data_ptr() tma_ptr = tma_desc.data_ptr() if hasattr(tma_desc, "data_ptr") else tma_desc - leader_barrier = cvta_to(_mapa_shared_cluster(barrier, ctm.Int32(0)), CvtaSpace.SHARED) + leader_barrier = prims.cvta_to( + _mapa_shared_cluster(barrier, ctm.Int32(0)), + prims.CvtaSpace.SHARED, + ) barrier_ptr = leader_barrier.data_ptr() multicast_mask_u16 = ctm.Uint16(multicast_mask) - _cp_async._predicated_inline_ptx( + cute_inline_ptx( "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4" ".mbarrier::complete_tx::bytes.multicast::cluster.cta_group::2" " [{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], " @@ -2365,7 +2354,7 @@ def _load_v_tile_stage( if should_wait_v_tma: while not prims.mbarrier_try_wait_parity(v_tma_mbar, v_tma_phase, time_limit=10000000): pass - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) @cute.jit @@ -2398,7 +2387,7 @@ def _issue_raw_k_pages_to_v_staging( raw_stage_offset = ctm.Int32( logical_page_idx * raw_page_bytes + n_tile_idx * raw_n_tile_bytes ) - _cp_async.cp_async_bulk_tensor_shared_cta_global( + prims.cp_async_bulk_tensor_shared_cta_global( sRawV.subview(raw_stage_offset), tma_k_v_raw_ptr, (packed_dim_begin, ctm.Int32(0), physical_page), @@ -2437,7 +2426,7 @@ def _transpose_raw_k_staging_to_v_stage( ) while not prims.mbarrier_try_wait_parity(raw_v_tma_mbar, raw_phase, time_limit=10000000): pass - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) raw_ptr = sRawV.data_ptr() v_word_ptr = sV.data_ptr() @@ -2555,7 +2544,7 @@ def _transpose_raw_k_staging_to_v_stage( ).bitcast(cutlass.Uint8) raw_fragment.store(packed_fragment) cute.copy(stsm_atom, raw_fragment, pv_lane) - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) prims.barrier( barrier_id=SMEM_P4_RAW_V_READY_BAR_ID, number_of_threads=SMEM_P4_RAW_V_READY_BAR_THREADS, @@ -2619,7 +2608,7 @@ def _load_vsf_tile_stage_only( @cute.jit def _fence_async_shared_cta() -> None: - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) @cute.jit @@ -2746,7 +2735,11 @@ def _softmax_exp2_pair_packed_f16x2( (score0, score1), (softmax_scale_log2, softmax_scale_log2), (p_bias, p_bias) ) shifted_h2 = _pack_f32_pair_to_f16x2(shifted[0], shifted[1]) - return ptx.ex2_f16x2(shifted_h2) + return cute_inline_ptx( + "ex2.approx.f16x2 {$w0}, {$r0};", + write_only_types=[ctm.Int32], + read_only_args=[shifted_h2], + ) @cute.jit @@ -4029,28 +4022,28 @@ def _smem_p4_p4_materialize_group_col( @cute.jit def _float_to_ordered_u32_for_atomic_max(value: ctm.Float32) -> ctm.Uint32: - bits = ptx.mov_b32(value, target_type=ctm.Int32) + bits = prims.mov_b32(value, target_type=ctm.Int32) sign_mask = bits >> ctm.Int32(31) | ctm.Int32(2147483648) encoded = bits ^ sign_mask - return ptx.mov_b32(encoded, target_type=ctm.Uint32) + return prims.mov_b32(encoded, target_type=ctm.Uint32) @cute.jit def _ordered_u32_to_float_after_atomic_max(value: ctm.Uint32) -> ctm.Float32: - encoded = ptx.mov_b32(value, target_type=ctm.Int32) + encoded = prims.mov_b32(value, target_type=ctm.Int32) sign_mask = ~(encoded >> ctm.Int32(31)) | ctm.Int32(2147483648) bits = encoded ^ sign_mask - return ptx.mov_b32(bits, target_type=ctm.Float32) + return prims.mov_b32(bits, target_type=ctm.Float32) @cute.jit def _smem_atomic_max_ordered_u32(pointer, value: ctm.Uint32) -> None: - ptx.atom( - AtomicOpKind.MAX, + prims.atomicrmw( + prims.AtomicOp.MAX, pointer, value, - syncscope=MemScopeKind.CTA, - space=SharedSpace.shared_cta, + syncscope=prims.MemScope.CTA, + space=prims.SharedSpace.shared_cta, ) @@ -4099,6 +4092,19 @@ def _prepare_p4_n256_score_half_tmem_addresses( return tuple(score_tmem_addresses) +@cute.jit +def _tcgen05_ld_red_x16_max_f32(tmem) -> tuple: + tmem_addr = tmem.toint(ctm.Int32) + return cute_inline_ptx( + "tcgen05.ld.red.sync.aligned.32x32b.x16.max.f32 " + "{{$w0}, {$w1}, {$w2}, {$w3}, {$w4}, {$w5}, {$w6}, {$w7}, " + "{$w8}, {$w9}, {$w10}, {$w11}, {$w12}, {$w13}, {$w14}, {$w15}}, " + "{$w16}, [{$r0}];", + write_only_types=[ctm.Int32] * 17, + read_only_args=[tmem_addr], + ) + + @cute.jit def _load_p4_n256_score_half_from_tmem( sAtomicRunningRowMax, @@ -4136,13 +4142,7 @@ def _load_p4_n256_score_half_from_tmem( pending_score_groups.append(group_scores) pending_group_stats.append(None) else: - regs = ptx.tcgen05_ld_red( - ptx.Tcgen05LdStShape.SHAPE_32X32B, - tmem, - num=SF_VEC_SIZE, - red_op="max", - type_="f32", - ) + regs = _tcgen05_ld_red_x16_max_f32(tmem) pending_score_groups.append(regs) pending_group_stats.append(regs[SF_VEC_SIZE]) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) @@ -4607,7 +4607,7 @@ def _finalize_p4_n256_score_half_psf_source( ) sPSF[psf_source_word_offset] = pv_sf_word cute.nvgpu.cfence() - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) cute.nvgpu.cfence() return (denominator_sf_word, prequant_row_sum) @@ -4719,7 +4719,7 @@ def _runtime_finalize_final_p4_n256_score_half_psf_source( ) sPSF[psf_source_word_offset] = pv_sf_word cute.nvgpu.cfence() - prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=SharedSpace.shared_cta) + prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) cute.nvgpu.cfence() return (denominator_sf_word, prequant_row_sum) @@ -5865,14 +5865,13 @@ def _runtime_t336_producer_tile( ) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - ptx.mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - 1, - sem=MBarrierArriveSem.RELAXED, - scope=MBarrierArriveScope.CLUSTER, - space=MBarrierSpace.SHARED_CLUSTER, - pred=common_source_ready_pred, - ) + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -6139,14 +6138,13 @@ def _runtime_t336_producer_tile_v23( warp_rebase = cute.arch.vote_any_sync(lane_rebase) if has_previous_tile: common_source_ready_pred = prims.elect_sync() & ~warp_rebase - ptx.mbarrier_arrive( - leader_pair_p_source_ready_dsmem_mbar, - 1, - sem=MBarrierArriveSem.RELAXED, - scope=MBarrierArriveScope.CLUSTER, - space=MBarrierSpace.SHARED_CLUSTER, - pred=common_source_ready_pred, - ) + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) if warp_rebase: pv_done_li_idx = stream_li_idx - ctm.Int32(1) pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -6841,7 +6839,7 @@ def _run_mla_decode_body( ) prims.fence_proxy( kind=prims.Proxy.ASYNC_SHARED, - space=SharedSpace.shared_cta, + space=prims.SharedSpace.shared_cta, ) _load_vsf_tile_stage_only( sPageIdPlan, From 1f4232f447aef1be0766dc80aaf577bf4aa7d194 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Mon, 31 Aug 2026 19:40:40 -0700 Subject: [PATCH 05/21] [None][fix] resolve FP4 MLA pre-commit failures Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/trtllm.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index b246312d7ed3..93edcc209bbf 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1275,7 +1275,6 @@ def prepare_encoder_decoder_from_precomputed_lengths( host_request_types=self.host_request_types[:num_seqs], ) - def prepare_encoder_only(self) -> None: """Fast path for encoder-only forward (eager + CUDA graph capture).""" extra_attrs = get_model_extra_attrs() From 76bd06cb0efd527807cf584d8d682f35b2c88abc Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:34:50 -0700 Subject: [PATCH 06/21] [None][fix] gate FP4 MLA fallback during FMHA registration Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/fmha/fallback.py | 5 ++--- tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index bee07b32a455..2f91b685fe6d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -62,6 +62,8 @@ class FallbackFmha(Fmha): @classmethod def _is_available(cls, attn: "TrtllmAttention") -> bool: + if attn.is_mla_enable and attn.has_fp4_kv_cache: + return False sparse_algorithm = getattr(attn.sparse_params, "algorithm", None) if sparse_algorithm in ("deepseek_v4", "dsa"): if getattr(attn, "kv_cache_dtype", None) == "fp8_ds_mla": @@ -99,9 +101,6 @@ def _is_supported( return False if q is not None and q.dtype == torch.float8_e4m3fn: return False - attn = self.attn - if attn.is_mla_enable and attn.has_fp4_kv_cache: - return False if forward_args.attention_mask == CustomAttentionMask.CUSTOM: return False if not forward_args.update_kv_cache and not metadata.is_cross: diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index 2a9a3ea9d73a..7c436abcd92d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -51,7 +51,7 @@ class Fp4MlaFmha(PhasedFmha): """TRTLLM FMHA library for FP4 MLA context and no-dequant decode.""" @classmethod - def is_available(cls, attn: "TrtllmAttention") -> bool: + def _is_available(cls, attn: "TrtllmAttention") -> bool: return attn.is_mla_enable and attn.has_fp4_kv_cache def forward( From 85a02e493eef6f246933372135319c2814f8de70 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:22:53 -0700 Subject: [PATCH 07/21] [None][fix] guard optional FP4 MLA metadata hooks Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/trtllm.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 93edcc209bbf..cc398ec36363 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -817,7 +817,8 @@ def on_update_kv_lens(self): self._invalidate_mla_scheduler_buffers() if getattr(self, '_fp4_mla_device_page_table', False): self._fp4_mla_device_page_table_valid = False - self._update_fp4_mla_append_metadata() + if getattr(self, 'high_precision_kv_pool', None) is not None: + self._update_fp4_mla_append_metadata() def _update_fp4_mla_append_metadata(self) -> None: if self.high_precision_kv_pool is not None and self.num_tokens > 0: @@ -832,7 +833,7 @@ def update_for_spec_dec(self) -> None: if self.enable_flash_mla: self._flash_mla_metadata_valid = False self._invalidate_mla_scheduler_buffers() - if self.high_precision_kv_pool is None: + if getattr(self, 'high_precision_kv_pool', None) is None: return num_seqs = self.num_seqs @@ -856,7 +857,7 @@ def restore_from_spec_dec(self) -> None: # positions/kv lens -> corrupted KV writes and illegal memory access # in the RoPE table lookup). super().restore_from_spec_dec() - if self.high_precision_kv_pool is None: + if getattr(self, 'high_precision_kv_pool', None) is None: return num_seqs = self.num_seqs self.kv_lens_cuda_runtime = self.kv_lens_cuda[:num_seqs] From b10e9c602ae02716aa1078415cf3e4c92a1dfffc Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 1 Sep 2026 04:14:34 -0700 Subject: [PATCH 08/21] [None][fix] preserve executor contracts with FP4 MLA routing Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 4 +- .../_torch/pyexecutor/model_engine.py | 2 +- .../_torch/pyexecutor/py_executor_creator.py | 52 +++++++++---------- 3 files changed, 28 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 941c2460ed0e..c2955d67bcc4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -188,7 +188,7 @@ def get_kv_cache_manager_cls( config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) - quant_config = model_config.quant_config + quant_config = getattr(model_config, "quant_config", None) if (is_mla(config) and quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache()): if is_disagg: @@ -906,7 +906,7 @@ def _validate_or_fallback_kv_cache_manager_v2( f"Gemma4 hybrid attention requires KVCacheManagerV2, " f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") - quant_config = model_config.quant_config + quant_config = getattr(model_config, "quant_config", None) if (is_mla(config) and quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache()): raise NotImplementedError( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 72e56fa84797..43ea9ea83296 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6105,7 +6105,7 @@ def _forward_scheduled(self, scheduled_requests: ScheduledRequests, LlmRequest]]): assert not self._is_packed_runner, ( "a packed-batch runner cannot execute scheduled requests") - if not self._disable_overlap_scheduler: + if not getattr(self, "_disable_overlap_scheduler", True): # Do not refill reusable host staging while the previous # iteration's asynchronous H2D copies still consume it. This # event precedes the model forward, so its synchronization retains diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 0f34152d13a7..a1e2c4c2485d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -59,7 +59,6 @@ _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS_STR = "/".join( f"SM{sm_version}" for sm_version in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) -FLASH_MLA_TOKENS_PER_BLOCK = 64 FP4_MLA_TOKENS_PER_BLOCK = 128 @@ -220,29 +219,6 @@ def _uses_fp4_mla_attention(model_config, kv_cache_config) -> bool: and getattr(model_config, "attn_backend", None) == "TRTLLM") -def _select_mla_tokens_per_block(config, model_config, kv_cache_config, - tokens_per_block: int) -> int: - if not is_mla(config): - return tokens_per_block - - if _uses_fp4_mla_attention(model_config, kv_cache_config): - tokens_per_block = FP4_MLA_TOKENS_PER_BLOCK - logger.info( - f"Change tokens_per_block to: {tokens_per_block} for using FP4 MLA attention" - ) - kv_cache_config.tokens_per_block = tokens_per_block - return tokens_per_block - - if model_config.enable_flash_mla: - tokens_per_block = FLASH_MLA_TOKENS_PER_BLOCK - logger.info( - f"Change tokens_per_block to: {tokens_per_block} for using FlashMLA" - ) - kv_cache_config.tokens_per_block = tokens_per_block - - return tokens_per_block - - def _configure_fp4_mla_speculative_kv_cache(spec_config, config, model_config, kv_cache_config, model) -> None: if (spec_config is None or not spec_config.spec_dec_mode.use_one_engine() @@ -712,10 +688,32 @@ def allocation_scope(current_stage: ExecutorMemoryType): llm_args.disable_overlap_scheduler, enable_overlap_headroom=getattr(model_engine, "_enable_overlap_headroom", False)) - tokens_per_block = _select_mla_tokens_per_block( - config, model_engine.model.model_config, kv_cache_config, - tokens_per_block) if is_mla(config): + if _uses_fp4_mla_attention(model_engine.model.model_config, + kv_cache_config): + tokens_per_block = FP4_MLA_TOKENS_PER_BLOCK + kv_cache_config.tokens_per_block = tokens_per_block + logger.info( + f"Change tokens_per_block to: {tokens_per_block} for using FP4 MLA attention" + ) + else: + if model_engine.model.model_config.enable_flash_mla: + tokens_per_block = 64 + # Propagate the override back to kv_cache_config so any consumer + # that later reads llm_args.kv_cache_config.tokens_per_block sees + # the effective value. KvCacheConnectorScheduler subclasses + # (LMCache, Dynamo KVBM) are instantiated further down via + # scheduler_cls(llm_args) and size their block pools from + # llm_args.kv_cache_config.tokens_per_block. Without this the + # connector's block size desynced from the KVCacheManager's + # actual tokens_per_block (user-set or default 32 vs. FlashMLA's + # forced 64), producing a frozen cache_block_ids view to the + # connector and silently-corrupted decode KV (#13320). + kv_cache_config.tokens_per_block = tokens_per_block + logger.info( + f"Change tokens_per_block to: {tokens_per_block} for using FlashMLA" + ) + sm_version = get_sm_version() if (kv_cache_config.enable_block_reuse and sm_version not in _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS): From 4c37a417b1aea0c07bd3fedb835ac9ebf08c46b8 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:27:15 -0700 Subject: [PATCH 09/21] [None][fix] preserve attention metadata compatibility for FP4 MLA Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/trtllm.py | 10 ++++++++++ tests/unittest/_torch/attention/test_fp4_mla.py | 8 +++++--- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index cc398ec36363..64f4f6c21706 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -150,6 +150,9 @@ def effective_beam_width(self) -> int: _max_seq_len_storage: Optional[int] = field(default=None, init=True, repr=False) + _page_size_override: Optional[int] = field(init=False, + default=None, + repr=False) # Encoder CUDA graph compatibility: overrides host-side max_context_q_len # so FMHA kernel launch params are stable across graph capture/replay even @@ -365,9 +368,16 @@ def page_size(self) -> int: """ Number of tokens per cache page. """ + page_size_override = getattr(self, '_page_size_override', None) + if page_size_override is not None: + return page_size_override assert self.kv_cache_manager is not None, "page_size requires a KV cache manager" return self.kv_cache_manager.tokens_per_block + @page_size.setter + def page_size(self, value: int) -> None: + self._page_size_override = value + @property def paged_kv_indices(self) -> torch.Tensor: """ diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index e90a69e3b037..7134df332060 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -32,7 +32,6 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexConverter _DataType = tensorrt_llm.bindings.DataType _CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -130,7 +129,7 @@ def populate_generation_lengths(*args, **kwargs) -> None: def test_fp4_mla_v2_encoded_page_capacity_is_layer_invariant(layer_offset: int) -> None: manager = object.__new__(Fp4MlaKVCacheManagerV2) manager.impl = SimpleNamespace( - get_page_index_converter=lambda *_: PageIndexConverter( + get_page_index_converter=lambda *_: SimpleNamespace( scale=4, expansion=1, layer_offset=layer_offset, @@ -170,12 +169,15 @@ def test_fp4_mla_v2_runtime_sizing_accounts_for_pipeline_slots() -> None: manager.max_num_tokens = 100 manager.tokens_per_block = FP4_MLA_TOKENS_PER_BLOCK manager.enable_swa_scratch_reuse = False + manager._has_cp_helix = False + manager._generation_kv_capacity_headroom = 1 manager._get_runtime_cache_size_layer_components = lambda: ([10, 8], [None, 19]) quota = manager._get_quota_from_max_tokens(manager.max_num_tokens) hp_page_bytes = FP4_MLA_TOKENS_PER_BLOCK * 8 - expected_quota = manager.max_num_tokens * (10 + 8) + 6 * hp_page_bytes + # The retained HP window can straddle two physical pages per request. + expected_quota = manager.max_num_tokens * (10 + 8) + 6 * 2 * hp_page_bytes assert quota == expected_quota assert manager._get_max_tokens_from_quota(quota) == manager.max_num_tokens From c81735329642dfe881c255c40a7c34179c513af5 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Mon, 7 Sep 2026 23:21:07 -0700 Subject: [PATCH 10/21] [None][fix] support large KV lengths in FP4 MLA kernels Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 216 ++++++++++++----- ...p4_mla_cutedsl_mufu16_fused_v_transpose.py | 226 +++++++++++++----- 2 files changed, 330 insertions(+), 112 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index caf71204a31a..622654e48b24 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -30,6 +30,7 @@ nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 +INT32_MAX = (1 << 31) - 1 _CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" _PYIR_STDOUT_LINES = frozenset( { @@ -292,9 +293,14 @@ def _initial_float_from_argv(option: str, default: float) -> float: # normalization below cancels this factor without changing the attention math. FP4_MLA_E4M3_MAX_FINITE = 448.0 FP4_MLA_P_GLOBAL_SCALE = FP4_MLA_E4M3_MAX_FINITE * 6.0 -SMEM_P4_RUNTIME_MAX_KV = 160 * 1024 -SMEM_P4_RUNTIME_MAX_PAGES = SMEM_P4_RUNTIME_MAX_KV // TRTLLM_PAGE_SIZE -SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_RUNTIME_MAX_PAGES +SMEM_P4_PAGE_PLAN_PROFILE_KV = 160 * 1024 +SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT = 4 +SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT = SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT * TRTLLM_PAGE_SIZE +# Keep the device-side ceil division ``valid_k + KV_TILE - 1`` in Int32 range. +SMEM_P4_RUNTIME_MAX_KV = ( + (INT32_MAX - (KV_TILE - 1)) // SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT +) * SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT +SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_PAGE_PLAN_PROFILE_KV // TRTLLM_PAGE_SIZE SMEM_P4_PAGE_ID_PLAN_BYTES = SMEM_P4_PAGE_ID_PLAN_INTS * ctm.Int32.bytes SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS = 2 SMEM_P4_RUNTIME_SCALE_PAIR_BYTES = SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS * ctm.Float32.bytes @@ -962,6 +968,26 @@ def _ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +def _select_runtime_kv_profile(max_kv_len: int) -> int: + if type(max_kv_len) is not int: + raise TypeError(f"max_kv_len must be an int, got {type(max_kv_len).__name__}") + if max_kv_len <= 0: + raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") + if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: + raise ValueError( + "max_kv_len exceeds the Int32-safe runtime-KV limit: " + f"{max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" + ) + if max_kv_len <= SMEM_P4_PAGE_PLAN_PROFILE_KV: + return SMEM_P4_PAGE_PLAN_PROFILE_KV + # Eager execution can report a different batch maximum on every step. + # Shifted power-of-two buckets bound the number of compiled variants and + # retain the established 1 Mi-token plus four-page-reserve profile. + profile_payload = max_kv_len - SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT + profile_kv = (1 << (profile_payload - 1).bit_length()) + (SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT) + return min(profile_kv, SMEM_P4_RUNTIME_MAX_KV) + + def _validate_output_dtype(output_dtype: torch.dtype) -> None: if output_dtype not in {torch.float16, torch.bfloat16}: raise TypeError(f"output dtype must be torch.float16 or torch.bfloat16, got {output_dtype}") @@ -1038,6 +1064,7 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, + use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, ) -> None: @@ -1060,9 +1087,7 @@ def fused_fp4_mla_decode_ctm( ) page_table_tensor = cute.make_tensor( cute.recast_ptr(page_table_ptr, dtype=cutlass.Int32), - cute.make_layout( - (SMEM_P4_RUNTIME_MAX_PAGES * (batch_size // query_len_per_seq),), stride=(1,) - ), + cute.make_layout(((k // page_size) * (batch_size // query_len_per_seq),), stride=(1,)), ) page_indptr_tensor = cute.make_tensor( cute.recast_ptr(page_indptr_ptr, dtype=cutlass.Int32), @@ -1363,6 +1388,7 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, + use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, @@ -1478,6 +1504,31 @@ def _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx: ctm.Int32) -> t return (physical_page0, physical_page1) +@cute.jit +def _load_global_page_native_tile_pair( + mPageTable_pl: cute.Tensor, + kv_tile_idx: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, +) -> tuple: + tile_page_idx = ctm.Int32(kv_tile_idx * SMEM_P4_PAGES_PER_KV_TILE) + # Keep these as two scalar loads: a CSR row may start at an odd Int32, and + # the second logical page must clamp to the first for a one-page tail. + physical_page0 = _lookup_physical_page( + mPageTable_pl, tile_page_idx, ctm.Int32(0), page_begin, page_count + ) + physical_page1 = _lookup_physical_page( + mPageTable_pl, + tile_page_idx + ctm.Int32(1), + ctm.Int32(0), + page_begin, + page_count, + ) + physical_page0 = cute.arch.make_warp_uniform(physical_page0) + physical_page1 = cute.arch.make_warp_uniform(physical_page1) + return (physical_page0, physical_page1) + + @cute.jit def _tma_gather4_cluster( smem_dst, @@ -2017,6 +2068,7 @@ def _load_v_tile_stage( v_tma_phase: ctm.Int32, stage: ctm.Int32 = 0, manage_mbarrier: ctm.Constexpr = True, + use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, ) -> None: del tidx @@ -2030,15 +2082,15 @@ def _load_v_tile_stage( if should_wait_v_tma: if prims.elect_sync(): prims.mbarrier_arrive_expect_tx(v_tma_mbar, v_tma_bytes) - # Resolve the page pair at the point of use so V-ring wrap cannot reuse a - # stale loop-carried ID. The plan was populated once before the producer - # warps start, so reloading it here avoids a second CSR/global lookup and - # its serial uniform-address chain on every tile. - del mPageTable_pl, bidz, page_begin, page_count - del tile_physical_page0, tile_physical_page1 - resolved_physical_page0, resolved_physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) + del bidz, tile_physical_page0, tile_physical_page1 + if cutlass.const_expr(use_smem_page_plan): + resolved_physical_page0, resolved_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) + else: + resolved_physical_page0, resolved_physical_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, page_begin, page_count + ) if cutlass.const_expr(use_consecutive_page_pair): resolved_physical_page1 = resolved_physical_page0 + ctm.Int32(1) if prims.elect_sync(): @@ -5856,6 +5908,7 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, @@ -6151,9 +6204,10 @@ def _run_mla_decode_body( _stage_smem_p4_runtime_scale_pair( mQGlobalScale, mKvGlobalScale, sRuntimeScalePair, softmax_scale_log2 ) - _stage_smem_p4_page_id_plan( - mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count - ) + if cutlass.const_expr(use_smem_page_plan): + _stage_smem_p4_page_id_plan( + mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count + ) cute.arch.cluster_arrive() cute.arch.cluster_wait() runtime_scale_pair = sRuntimeScalePair.data_ptr().load(count=2, alignment=8) @@ -6203,7 +6257,14 @@ def _run_mla_decode_body( qk0_handle = qk_smem_producer.acquire_and_advance() qk0_tma_mbar = ctm.Array(qk0_handle.barrier, shape=1) if is_leader_cta: - qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, ctm.Int32(0)) + if cutlass.const_expr(use_smem_page_plan): + qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(0) + ) + else: + qk0_page0, qk0_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, ctm.Int32(0), csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6237,9 +6298,14 @@ def _run_mla_decode_body( ) qk_prefix_tma_mbar = ctm.Array(qk_prefix_handle.barrier, shape=1) if is_leader_cta: - qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) + if cutlass.const_expr(use_smem_page_plan): + qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) + else: + qk_prefix_page0, qk_prefix_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6276,9 +6342,14 @@ def _run_mla_decode_body( ) qk_tma_mbar = ctm.Array(qk_handle.barrier, shape=1) if is_leader_cta: - qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) + if cutlass.const_expr(use_smem_page_plan): + qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) + else: + qk_steady_page0, qk_steady_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6310,9 +6381,14 @@ def _run_mla_decode_body( ) qk15_tma_mbar = ctm.Array(qk15_handle.barrier, shape=1) if is_leader_cta: - qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(15) - ) + if cutlass.const_expr(use_smem_page_plan): + qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(15) + ) + else: + qk15_page0, qk15_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, ctm.Int32(15), csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6342,9 +6418,12 @@ def _run_mla_decode_body( cute.arch.cp_async_bulk_wait_group(0, read=True) q_smem_producer.tail() if warp_idx == TMA_V_WARP_ID: - v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(0) - ) + v_physical_page0 = ctm.Int32(0) + v_physical_page1 = ctm.Int32(0) + if cutlass.const_expr(use_smem_page_plan): + v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(0) + ) for kv_tile_idx in cutlass.range(0, kv_tiles, 1, unroll=1): v_handle = v_smem_producer.acquire_and_advance() v_tma_mbar = ctm.Array(v_handle.barrier, shape=1) @@ -6373,20 +6452,27 @@ def _run_mla_decode_body( v_tma_phase=v_stage, stage=v_stage, manage_mbarrier=False, + use_smem_page_plan=use_smem_page_plan, use_consecutive_page_pair=use_consecutive_page_pair, ) - if kv_tile_idx + ctm.Int32(1) < kv_tiles: - v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx + ctm.Int32(1) - ) + if cutlass.const_expr(use_smem_page_plan): + if kv_tile_idx + ctm.Int32(1) < kv_tiles: + v_physical_page0, v_physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ctm.Int32(1) + ) v_smem_producer.tail() if warp_idx == ROWMETA_WARP_ID: if is_leader_cta: qk_prefix_end = min(kv_tiles, ctm.Int32(3)) for qk_prefix_li_idx in cutlass.range(1, qk_prefix_end, 1, unroll=1): - rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) + if cutlass.const_expr(use_smem_page_plan): + rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) + else: + rowmeta_page0, rowmeta_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count + ) _issue_runtime_t336_qk_prefix_rank1( mPageTable_pl, tma_k_page_ptr, @@ -6867,6 +6953,7 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, @@ -6902,6 +6989,7 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, + use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, @@ -7155,13 +7243,25 @@ def _compile_fused( vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, ) -> Callable: - if kv != SMEM_P4_RUNTIME_MAX_KV: - raise ValueError(f"runtime-KV compile requires fixed K={SMEM_P4_RUNTIME_MAX_KV}, got {kv}") + if type(kv) is not int: + raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") + is_page_plan_profile = kv == SMEM_P4_PAGE_PLAN_PROFILE_KV + is_bucketed_runtime_profile = ( + SMEM_P4_PAGE_PLAN_PROFILE_KV < kv <= SMEM_P4_RUNTIME_MAX_KV + and _select_runtime_kv_profile(kv) == kv + ) + if not (is_page_plan_profile or is_bucketed_runtime_profile): + raise ValueError( + "runtime-KV compile requires the fixed page-plan profile " + f"{SMEM_P4_PAGE_PLAN_PROFILE_KV} or a larger geometric profile " + f"up to {SMEM_P4_RUNTIME_MAX_KV}, got {kv}" + ) use_ksf_gather4 = ( not use_consecutive_page_pair and ksf_page_stride_bytes == page_size * TRTLLM_K_SF_GROUPS ) cache_key = ( n, + kv, page_size, use_mixed_imlp, output_dtype, @@ -7210,6 +7310,7 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, + use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, options="--opt-level 2 --ptxas-options '--uumn'", @@ -7257,12 +7358,7 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"page-native decode requires page_size={TRTLLM_PAGE_SIZE}, got {page_size}" ) - if max_kv_len <= 0: - raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") - if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: - raise ValueError( - f"max_kv_len exceeds the fixed runtime-KV profile: {max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" - ) + physical_k = _select_runtime_kv_profile(max_kv_len) for name, value in ( ("assume_valid_k_prefix_tiles", assume_valid_k_prefix_tiles), ( @@ -7288,7 +7384,6 @@ def run_trtllm_fp4_mla_decode_page_native( del assume_valid_k_prefix_tiles, partition_runtime_valid_k required_consecutive_tiles = _ceil_div(max_kv_len, KV_TILE) use_consecutive_page_pair = assume_consecutive_page_prefix_tiles >= required_consecutive_tiles - physical_k = SMEM_P4_RUNTIME_MAX_KV if q_internal.dtype != torch.uint8 or q_internal.dim() != 3: raise TypeError( "q_internal must be a uint8 [M, Q640/2, L] tensor, " @@ -7348,10 +7443,15 @@ def run_trtllm_fp4_mla_decode_page_native( "src_page_ids must be a 1D int32 physical-page list, " f"got dtype={src_page_ids.dtype} shape={tuple(src_page_ids.shape)}" ) - if src_page_ids.numel() == 0 or src_page_ids.stride(0) != 1: + src_page_id_count = src_page_ids.numel() + if src_page_id_count == 0 or src_page_ids.stride(0) != 1: raise ValueError( "src_page_ids must be non-empty and contiguous, " - f"got numel={src_page_ids.numel()} stride={src_page_ids.stride()}" + f"got numel={src_page_id_count} stride={src_page_ids.stride()}" + ) + if src_page_id_count > INT32_MAX: + raise ValueError( + f"src_page_ids exceeds the Int32 element limit: {src_page_id_count} > {INT32_MAX}" ) if ( paged_kv_indptr_decode.dtype != torch.int32 @@ -7389,9 +7489,14 @@ def run_trtllm_fp4_mla_decode_page_native( ) num_cache_pages = cache_layout.num_pages page_table_capacity = num_sequences * (physical_k // page_size) - if src_page_ids.numel() > page_table_capacity: + if page_table_capacity > INT32_MAX: + raise ValueError( + "bucketed CSR capacity exceeds the Int32 layout limit: " + f"{page_table_capacity} > {INT32_MAX}" + ) + if src_page_id_count > page_table_capacity: raise ValueError( - f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_ids.numel()} > {page_table_capacity}" + f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_id_count} > {page_table_capacity}" ) expected_v_rows = num_cache_pages * (TRTLLM_V_HEAD_DIM // v_pack_block) * v_pack_block if ( @@ -7414,12 +7519,11 @@ def run_trtllm_fp4_mla_decode_page_native( if not isinstance(v_page_offset, int): raise TypeError(f"v_page_offset must be an int, got {type(v_page_offset).__name__}") num_v_cache_pages = v_packed.shape[0] // v_rows_per_page - int32_max = torch.iinfo(torch.int32).max - if v_page_offset < 0 or v_page_offset > int32_max: - raise ValueError(f"v_page_offset must be in [0, {int32_max}], got {v_page_offset}") - if num_v_cache_pages > int32_max: + if v_page_offset < 0 or v_page_offset > INT32_MAX: + raise ValueError(f"v_page_offset must be in [0, {INT32_MAX}], got {v_page_offset}") + if num_v_cache_pages > INT32_MAX: raise ValueError( - f"v_packed exceeds the Int32 physical-page limit: {num_v_cache_pages} > {int32_max}" + f"v_packed exceeds the Int32 physical-page limit: {num_v_cache_pages} > {INT32_MAX}" ) if v_page_offset + num_cache_pages > num_v_cache_pages: raise ValueError( diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py index 1301574752f5..6d582c78ce0a 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py @@ -32,6 +32,7 @@ nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 +INT32_MAX = (1 << 31) - 1 _CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" _PYIR_STDOUT_LINES = frozenset( { @@ -294,9 +295,14 @@ def _initial_float_from_argv(option: str, default: float) -> float: # normalization below cancels this factor without changing the attention math. FP4_MLA_E4M3_MAX_FINITE = 448.0 FP4_MLA_P_GLOBAL_SCALE = FP4_MLA_E4M3_MAX_FINITE * 6.0 -SMEM_P4_RUNTIME_MAX_KV = 160 * 1024 -SMEM_P4_RUNTIME_MAX_PAGES = SMEM_P4_RUNTIME_MAX_KV // TRTLLM_PAGE_SIZE -SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_RUNTIME_MAX_PAGES +SMEM_P4_PAGE_PLAN_PROFILE_KV = 160 * 1024 +SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT = 4 +SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT = SMEM_P4_RUNTIME_PROFILE_PAGE_ALIGNMENT * TRTLLM_PAGE_SIZE +# Keep the device-side ceil division ``valid_k + KV_TILE - 1`` in Int32 range. +SMEM_P4_RUNTIME_MAX_KV = ( + (INT32_MAX - (KV_TILE - 1)) // SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT +) * SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT +SMEM_P4_PAGE_ID_PLAN_INTS = SMEM_P4_PAGE_PLAN_PROFILE_KV // TRTLLM_PAGE_SIZE SMEM_P4_PAGE_ID_PLAN_BYTES = SMEM_P4_PAGE_ID_PLAN_INTS * ctm.Int32.bytes SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS = 2 SMEM_P4_RUNTIME_SCALE_PAIR_BYTES = SMEM_P4_RUNTIME_SCALE_PAIR_FLOATS * ctm.Float32.bytes @@ -963,6 +969,26 @@ def _ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +def _select_runtime_kv_profile(max_kv_len: int) -> int: + if type(max_kv_len) is not int: + raise TypeError(f"max_kv_len must be an int, got {type(max_kv_len).__name__}") + if max_kv_len <= 0: + raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") + if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: + raise ValueError( + "max_kv_len exceeds the Int32-safe runtime-KV limit: " + f"{max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" + ) + if max_kv_len <= SMEM_P4_PAGE_PLAN_PROFILE_KV: + return SMEM_P4_PAGE_PLAN_PROFILE_KV + # Eager execution can report a different batch maximum on every step. + # Shifted power-of-two buckets bound the number of compiled variants and + # retain the established 1 Mi-token plus four-page-reserve profile. + profile_payload = max_kv_len - SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT + profile_kv = (1 << (profile_payload - 1).bit_length()) + (SMEM_P4_RUNTIME_PROFILE_KV_ALIGNMENT) + return min(profile_kv, SMEM_P4_RUNTIME_MAX_KV) + + def _validate_output_dtype(output_dtype: torch.dtype) -> None: if output_dtype not in {torch.float16, torch.bfloat16}: raise TypeError(f"output dtype must be torch.float16 or torch.bfloat16, got {output_dtype}") @@ -1032,6 +1058,7 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, + use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, ) -> None: @@ -1060,10 +1087,9 @@ def fused_fp4_mla_decode_ctm( stride=(1, TRTLLM_K_STORAGE_DIM // 2, kv_page_stride_bytes), ), ) - page_count = cute.assume(k // page_size, 1) page_table_tensor = cute.make_tensor( cute.recast_ptr(page_table_ptr, dtype=cutlass.Int32), - cute.make_layout((SMEM_P4_RUNTIME_MAX_PAGES * (l // query_len_per_seq),), stride=(1,)), + cute.make_layout(((k // page_size) * (l // query_len_per_seq),), stride=(1,)), ) page_indptr_tensor = cute.make_tensor( cute.recast_ptr(page_indptr_ptr, dtype=cutlass.Int32), @@ -1371,6 +1397,7 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, + use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, @@ -1473,6 +1500,31 @@ def _lookup_physical_page( return ctm.Int32(mPageTable_pl[page_table_idx]) +@cute.jit +def _load_global_page_native_tile_pair( + mPageTable_pl: cute.Tensor, + kv_tile_idx: ctm.Int32, + page_begin: ctm.Int32, + page_count: ctm.Int32, +) -> tuple: + tile_page_idx = ctm.Int32(kv_tile_idx * SMEM_P4_PAGES_PER_KV_TILE) + # Keep these as two scalar loads: a CSR row may start at an odd Int32, and + # the second logical page must clamp to the first for a one-page tail. + physical_page0 = _lookup_physical_page( + mPageTable_pl, tile_page_idx, ctm.Int32(0), page_begin, page_count + ) + physical_page1 = _lookup_physical_page( + mPageTable_pl, + tile_page_idx + ctm.Int32(1), + ctm.Int32(0), + page_begin, + page_count, + ) + physical_page0 = cute.arch.make_warp_uniform(physical_page0) + physical_page1 = cute.arch.make_warp_uniform(physical_page1) + return (physical_page0, physical_page1) + + @cute.jit def _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx: ctm.Int32) -> tuple: physical_page0 = ctm.Int32(0) @@ -2223,7 +2275,6 @@ def _issue_runtime_t336_qk_prefix_rank1( @cute.jit def _load_v_tile_stage( mPageTable_pl: cute.Tensor, - sPageIdPlan, tma_v_tile_ptr, tma_v_pair_ptr, tma_vsf_ptr, @@ -2258,15 +2309,11 @@ def _load_v_tile_stage( if should_wait_v_tma: if prims.elect_sync(): prims.mbarrier_arrive_expect_tx(v_tma_mbar, v_tma_bytes) - # Resolve the page pair at the point of use so V-ring wrap cannot reuse a - # stale loop-carried ID. The plan was populated once before the producer - # warps start, so reloading it here avoids a second CSR/global lookup and - # its serial uniform-address chain on every tile. + # Resolve the supplied pair at the point of use so V-ring wrap cannot + # reuse stale loop-carried IDs. del mPageTable_pl, bidz, page_begin, page_count - del tile_physical_page0, tile_physical_page1 - resolved_physical_page0, resolved_physical_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) + resolved_physical_page0 = tile_physical_page0 + resolved_physical_page1 = tile_physical_page1 if cutlass.const_expr(use_consecutive_page_pair): resolved_physical_page1 = resolved_physical_page0 + ctm.Int32(1) if prims.elect_sync(): @@ -2362,15 +2409,20 @@ def _issue_raw_k_pages_to_v_staging( tma_k_v_raw_ptr, raw_v_tma_mbar, sRawV, - sV, sPageIdPlan, - tidx: ctm.Int32, cta_rank: ctm.Int32, kv_tile_idx: ctm.Int32, - v_stage: ctm.Int32, + direct_physical_page0: ctm.Int32, + direct_physical_page1: ctm.Int32, + use_smem_page_plan: ctm.Constexpr, ) -> None: """Stage the V-bearing K prefix in a transpose-friendly layout.""" - physical_page0, physical_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx) + physical_page0 = direct_physical_page0 + physical_page1 = direct_physical_page1 + if cutlass.const_expr(use_smem_page_plan): + physical_page0, physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) raw_packed_dims_per_n_tile: ctm.Constexpr = SMEM_P4_V_N_PER_CTA // 2 raw_n_tile_bytes: ctm.Constexpr = TRTLLM_PAGE_SIZE * raw_packed_dims_per_n_tile raw_page_bytes: ctm.Constexpr = raw_n_tile_bytes * SMEM_P4_N_OUT_TILES @@ -2560,12 +2612,20 @@ def _load_vsf_tile_stage_only( sVSF, cta_rank: ctm.Int32, kv_tile_idx: ctm.Int32, + direct_physical_page0: ctm.Int32, + direct_physical_page1: ctm.Int32, v_page_offset: ctm.Int32, stage: ctm.Int32, + use_smem_page_plan: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr = False, ) -> None: sSFB_stage = sVSF.subview(SMEM_P4_V_SFB_TMA_STAGE_BYTES * stage) - physical_page0, physical_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, kv_tile_idx) + physical_page0 = direct_physical_page0 + physical_page1 = direct_physical_page1 + if cutlass.const_expr(use_smem_page_plan): + physical_page0, physical_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) if cutlass.const_expr(use_consecutive_page_pair): physical_page1 = physical_page0 + ctm.Int32(1) if cta_rank == ctm.Int32(0): @@ -6309,6 +6369,7 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, @@ -6626,9 +6687,10 @@ def _run_mla_decode_body( _stage_smem_p4_runtime_scale_pair( mQGlobalScale, mKvGlobalScale, sRuntimeScalePair, softmax_scale_log2 ) - _stage_smem_p4_page_id_plan( - mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count - ) + if cutlass.const_expr(use_smem_page_plan): + _stage_smem_p4_page_id_plan( + mPageTable_pl, mPageIndptr_s, sPageIdPlan, tidx, page_batch, planned_page_count + ) cute.arch.cluster_arrive() cute.arch.cluster_wait() runtime_scale_pair = sRuntimeScalePair.data_ptr().load(count=2, alignment=8) @@ -6678,7 +6740,14 @@ def _run_mla_decode_body( qk0_handle = qk_smem_producer.acquire_and_advance() qk0_tma_mbar = ctm.Array(qk0_handle.barrier, shape=1) if is_leader_cta: - qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair(sPageIdPlan, ctm.Int32(0)) + if cutlass.const_expr(use_smem_page_plan): + qk0_page0, qk0_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(0) + ) + else: + qk0_page0, qk0_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, ctm.Int32(0), csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6712,9 +6781,14 @@ def _run_mla_decode_body( ) qk_prefix_tma_mbar = ctm.Array(qk_prefix_handle.barrier, shape=1) if is_leader_cta: - qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) + if cutlass.const_expr(use_smem_page_plan): + qk_prefix_page0, qk_prefix_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) + else: + qk_prefix_page0, qk_prefix_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6750,9 +6824,14 @@ def _run_mla_decode_body( expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] ) qk_tma_mbar = ctm.Array(qk_handle.barrier, shape=1) - qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, kv_tile_idx - ) + if cutlass.const_expr(use_smem_page_plan): + qk_steady_page0, qk_steady_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) + else: + qk_steady_page0, qk_steady_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile_dual_steady( mPageTable_pl, tma_k_page_ptr, @@ -6785,9 +6864,14 @@ def _run_mla_decode_body( ) qk15_tma_mbar = ctm.Array(qk15_handle.barrier, shape=1) if is_leader_cta: - qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, ctm.Int32(15) - ) + if cutlass.const_expr(use_smem_page_plan): + qk15_page0, qk15_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, ctm.Int32(15) + ) + else: + qk15_page0, qk15_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, ctm.Int32(15), csr_page_begin, csr_page_count + ) _load_runtime_t336_qk_tile( mPageTable_pl, tma_k_page_ptr, @@ -6822,16 +6906,22 @@ def _run_mla_decode_body( v_tma_mbar = ctm.Array(v_handle.barrier, shape=1) v_stage = kv_tile_idx % ctm.Int32(SMEM_P4_V_PIPELINE_STAGES) raw_v_stage_mbar = raw_v_tma_mbar.subview(v_stage) + v_physical_page0 = ctm.Int32(0) + v_physical_page1 = ctm.Int32(0) + if cutlass.const_expr(not use_smem_page_plan): + v_physical_page0, v_physical_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count + ) _issue_raw_k_pages_to_v_staging( tma_k_v_raw_ptr, raw_v_stage_mbar, sRawV, - sV, sPageIdPlan, - tidx, cta_rank, kv_tile_idx, - v_stage, + v_physical_page0, + v_physical_page1, + use_smem_page_plan, ) prims.barrier( barrier_id=SMEM_P4_RAW_V_READY_BAR_ID, @@ -6849,8 +6939,11 @@ def _run_mla_decode_body( sVSF, cta_rank, kv_tile_idx, + v_physical_page0, + v_physical_page1, v_page_offset, v_stage, + use_smem_page_plan, use_consecutive_page_pair=use_consecutive_page_pair, ) v_smem_producer.tail() @@ -6858,9 +6951,14 @@ def _run_mla_decode_body( if is_leader_cta: qk_prefix_end = min(kv_tiles, ctm.Int32(3)) for qk_prefix_li_idx in cutlass.range(1, qk_prefix_end, 1, unroll=1): - rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( - sPageIdPlan, qk_prefix_li_idx - ) + if cutlass.const_expr(use_smem_page_plan): + rowmeta_page0, rowmeta_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, qk_prefix_li_idx + ) + else: + rowmeta_page0, rowmeta_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, qk_prefix_li_idx, csr_page_begin, csr_page_count + ) _issue_runtime_t336_qk_prefix_rank1( mPageTable_pl, tma_k_page_ptr, @@ -7355,6 +7453,7 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, @@ -7391,6 +7490,7 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, + use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, @@ -7590,12 +7690,22 @@ def _compile_fused( vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, ) -> Callable: - if kv != SMEM_P4_RUNTIME_MAX_KV: - raise ValueError(f"runtime-KV compile requires fixed K={SMEM_P4_RUNTIME_MAX_KV}, got {kv}") + if type(kv) is not int: + raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") + is_bucketed_runtime_profile = ( + SMEM_P4_PAGE_PLAN_PROFILE_KV < kv <= SMEM_P4_RUNTIME_MAX_KV + and _select_runtime_kv_profile(kv) == kv + ) + if kv != SMEM_P4_PAGE_PLAN_PROFILE_KV and not is_bucketed_runtime_profile: + raise ValueError( + "runtime-KV compile requires the fixed page-plan profile or a geometric runtime " + f"profile in ({SMEM_P4_PAGE_PLAN_PROFILE_KV}, {SMEM_P4_RUNTIME_MAX_KV}], got {kv}" + ) use_ksf_gather4 = ( not use_consecutive_page_pair and ksf_page_stride_bytes == page_size * TRTLLM_K_SF_GROUPS ) cache_key = ( + kv, n, page_size, use_mixed_imlp, @@ -7645,6 +7755,7 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, + use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, options="--opt-level 2 --ptxas-options '--uumn'", @@ -7692,12 +7803,7 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"page-native decode requires page_size={TRTLLM_PAGE_SIZE}, got {page_size}" ) - if max_kv_len <= 0: - raise ValueError(f"max_kv_len must be positive, got {max_kv_len}") - if max_kv_len > SMEM_P4_RUNTIME_MAX_KV: - raise ValueError( - f"max_kv_len exceeds the fixed runtime-KV profile: {max_kv_len} > {SMEM_P4_RUNTIME_MAX_KV}" - ) + physical_k = _select_runtime_kv_profile(max_kv_len) for name, value in ( ("assume_valid_k_prefix_tiles", assume_valid_k_prefix_tiles), ( @@ -7723,7 +7829,6 @@ def run_trtllm_fp4_mla_decode_page_native( del assume_valid_k_prefix_tiles, partition_runtime_valid_k required_consecutive_tiles = _ceil_div(max_kv_len, KV_TILE) use_consecutive_page_pair = assume_consecutive_page_prefix_tiles >= required_consecutive_tiles - physical_k = SMEM_P4_RUNTIME_MAX_KV if q_internal.dtype != torch.uint8 or q_internal.dim() != 3: raise TypeError( f"q_internal must be a uint8 [M, Q640/2, L] tensor, got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" @@ -7780,9 +7885,14 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"src_page_ids must be a 1D int32 physical-page list, got dtype={src_page_ids.dtype} shape={tuple(src_page_ids.shape)}" ) - if src_page_ids.numel() == 0 or src_page_ids.stride(0) != 1: + src_page_id_count = src_page_ids.numel() + if src_page_id_count == 0 or src_page_ids.stride(0) != 1: raise ValueError( - f"src_page_ids must be non-empty and contiguous, got numel={src_page_ids.numel()} stride={src_page_ids.stride()}" + f"src_page_ids must be non-empty and contiguous, got numel={src_page_id_count} stride={src_page_ids.stride()}" + ) + if src_page_id_count > INT32_MAX: + raise ValueError( + f"src_page_ids exceeds the Int32 element limit: {src_page_id_count} > {INT32_MAX}" ) if ( paged_kv_indptr_decode.dtype != torch.int32 @@ -7810,21 +7920,25 @@ def run_trtllm_fp4_mla_decode_page_native( ) num_cache_pages = cache_layout.num_pages page_table_capacity = num_sequences * (physical_k // page_size) - if src_page_ids.numel() > page_table_capacity: + if page_table_capacity > INT32_MAX: + raise ValueError( + "bucketed CSR capacity exceeds the Int32 layout limit: " + f"{page_table_capacity} > {INT32_MAX}" + ) + if src_page_id_count > page_table_capacity: raise ValueError( - f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_ids.numel()} > {page_table_capacity}" + f"src_page_ids exceeds the bucketed CSR capacity, got {src_page_id_count} > {page_table_capacity}" ) if not isinstance(v_page_offset, int): raise TypeError(f"v_page_offset must be an int, got {type(v_page_offset).__name__}") if v_sf.dim() == 0: raise ValueError("v_sf must expose a physical-page dimension") num_v_cache_pages = int(v_sf.shape[0]) - int32_max = torch.iinfo(torch.int32).max - if v_page_offset < 0 or v_page_offset > int32_max: - raise ValueError(f"v_page_offset must be in [0, {int32_max}], got {v_page_offset}") - if num_v_cache_pages > int32_max: + if v_page_offset < 0 or v_page_offset > INT32_MAX: + raise ValueError(f"v_page_offset must be in [0, {INT32_MAX}], got {v_page_offset}") + if num_v_cache_pages > INT32_MAX: raise ValueError( - f"v_sf exceeds the Int32 physical-page limit: {num_v_cache_pages} > {int32_max}" + f"v_sf exceeds the Int32 physical-page limit: {num_v_cache_pages} > {INT32_MAX}" ) if v_page_offset + num_cache_pages > num_v_cache_pages: raise ValueError( From adebf19dee2308743aaa16db1973a70e05696094 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:50:14 -0700 Subject: [PATCH 11/21] fix: address FP4 MLA core review feedback Align FP4 routing and V2 selection with resolved quantization, validate unsupported profiles before allocation, and configure draft-cache sharing before model construction. Drain auxiliary work on context failures, remove duplicate cache append and unrelated executor/V1 changes, scope MTP metadata updates, and consolidate kernel helpers. Strengthen existing routing and GSM8K checks and register the Rubin accuracy case in the QA test list. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 9 ++ .../attention/backends/fp4_mla/__init__.py | 8 +- .../backends/fp4_mla/cache_manager.py | 35 +++++-- .../backends/fp4_mla/fp4_mla_context.py | 21 +++-- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 7 +- ...p4_mla_cutedsl_mufu16_fused_v_transpose.py | 7 +- .../fp4_mla/fp4_mla_cutedsl_v_repack.py | 15 +-- .../backends/fp4_mla/fp4_mla_triton.py | 62 +++---------- .../_torch/attention/backends/trtllm.py | 48 +--------- tensorrt_llm/_torch/attention/mla.py | 30 ++++-- tensorrt_llm/_torch/pyexecutor/_util.py | 6 ++ .../_torch/pyexecutor/config_utils.py | 21 ++++- .../_torch/pyexecutor/model_engine.py | 8 -- .../_torch/pyexecutor/model_loader.py | 51 ++++++++++- .../_torch/pyexecutor/py_executor_creator.py | 91 +++++-------------- .../_torch/pyexecutor/resource_manager.py | 11 +-- tensorrt_llm/_torch/speculative/mtp.py | 9 +- .../defs/accuracy/test_llm_api_pytorch.py | 38 +++++++- .../test_lists/qa/llm_function_core.txt | 3 + .../unittest/_torch/attention/backend_case.py | 48 +++------- .../unittest/_torch/attention/test_fp4_mla.py | 80 +++++++--------- tests/unittest/_torch/test_model_config.py | 54 ++++++++++- 22 files changed, 330 insertions(+), 332 deletions(-) diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index fdfe6e628710..461567fd7b85 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -428,6 +428,15 @@ The FMHA package is split by role: - `fmha/combined.py` composes different context and generation implementations for non-MLA mixed batches. - `fmha/fp4_mla.py` implements FP4 MLA context and no-dequant decode. + The core implementation requires dense TRTLLM MLA, BF16 absorption weights, + fused RoPE with duplicated rotary tables, and KV Cache Manager V2. It uses + FP8 context attention with an FP4 cache update and FP4 generation attention. + Chunked prefill and context parallelism are rejected before KV allocation. + Cached-context attention is not implemented, so executor block reuse remains + disabled even though the manager supports full-block reuse of its pools. + Disaggregated serving is reserved for the follow-up integration. On SM107, + this dense TRTLLM path keeps NVFP4 KV quantization; unsupported profiles retain + the existing FP8 fallback. - `fmha/triton_custom_mask.py` implements the Triton custom-mask context phase. Custom-mask data applies to context requests; for mixed batches, `TrtllmAttention` can pair it with a later causal-generation provider through diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py index c60cd9512ff7..9924a4a7e832 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -1192,9 +1192,9 @@ def _get_fp4_mla_hp_pool_layout( ) -> tuple[int, int]: """Return the manager-owned HP ring size and per-token head dimension.""" manager = getattr(metadata, "kv_cache_manager", None) - if manager is None or not hasattr(manager, "_fp4_mla_hp_pool_size"): + if manager is None or not hasattr(manager, "fp4_mla_hp_pool_size"): raise ValueError("FP4 MLA requires a V2 manager-owned HP ring.") - hp_pool_size = int(manager._fp4_mla_hp_pool_size) + hp_pool_size = manager.fp4_mla_hp_pool_size if ( hp_pool_size < HP_BLOCK_SIZE or pool.ndim != 4 @@ -2585,7 +2585,7 @@ def _update_triton_v_packed_cache( return None if page_ids.numel() == 0: return None - from .fp4_mla_triton import fp4_mla_repack_v_cache + from .fp4_mla_triton import fp4_mla_repack_v_cache_triton def _tma_alloc(size: int, alignment: int, stream): return torch.empty(size, device=kv_cache.device, dtype=torch.int8) @@ -2599,7 +2599,7 @@ def _tma_alloc(size: int, alignment: int, stream): dtype=torch.uint8, device=kv_cache.device, ) - fp4_mla_repack_v_cache( + fp4_mla_repack_v_cache_triton( v_packed, kv_cache, page_ids, diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index 4ab4e0458cee..e5e8f59ed92f 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -4,7 +4,7 @@ import math from dataclasses import dataclass, replace -from typing import List, Optional +from typing import TYPE_CHECKING, List, Optional import torch @@ -12,14 +12,20 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, BufferConfig, DataRole, + KVCacheManagerConfig, LayerId, PageIndexMode, ) +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + from . import ( _FP4_MLA_CUTEDSL_BACKEND, _FP4_MLA_K_RESIDUAL_BACKENDS, @@ -136,6 +142,11 @@ def __init__(self, *args, **kwargs) -> None: self.get_mla_v_scale_pool_base().zero_() self._stream.synchronize() + @property + def fp4_mla_hp_pool_size(self) -> int: + """Number of BF16 tokens in each layer's rewind-capable HP ring.""" + return self._fp4_mla_hp_pool_size + @property def blocks_in_primary_pool(self) -> int: return self._role_encoded_page_capacity(self._cache_manager_layer_ids[0], Role.KEY) @@ -158,7 +169,7 @@ def _hp_bytes_per_page(self, local_layer_idx: int) -> int: * torch.empty((), dtype=torch.bfloat16).element_size() ) - def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole): + def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) -> int: storage_head_dim = self._storage_head_dim(local_layer_idx) role_sizes = { Role.KEY: math.ceil(storage_head_dim / 2), @@ -199,7 +210,7 @@ def _extra_buffers_per_layer( result[local_layer] = buffers return result - def _build_cache_config(self, config): + def _build_cache_config(self, config: KVCacheManagerConfig) -> KVCacheManagerConfig: cache_layers = list(config.layers) self._cache_manager_layer_ids = [LayerId(i) for i in range(self.num_local_layers)] self._hp_manager_layer_ids = [] @@ -542,7 +553,7 @@ def get_fp4_mla_hp_pool(self) -> torch.Tensor: (1, self._fp4_mla_hp_pool_size * self.head_dim_per_layer[0]), ).permute(1, 0, 2, 3) - def _get_runtime_cache_size_layer_components(self): + def _get_runtime_cache_size_layer_components(self) -> tuple[list[int], list[Optional[int]]]: sizes = [ self.get_layer_bytes_per_token(local_layer, Role.ALL) for local_layer in range(self.num_local_layers) @@ -565,7 +576,16 @@ def get_cache_bytes_per_token(self) -> int: return sum(sizes) @staticmethod - def get_cache_size_per_token(model_config, mapping, num_layers=None, **kwargs): + def get_cache_size_per_token( + model_config: "ModelConfig", + mapping: Mapping, + num_layers: Optional[int] = None, + *, + tokens_per_block: int, + spec_config: Optional["DecodingBaseConfig"] = None, + max_batch_size: Optional[int] = None, + **kwargs: object, + ) -> tuple[int, int]: config = model_config.pretrained_config logical_head_dim = int(config.kv_lora_rank + config.qk_rope_head_dim) backend = _fp4_mla_attention_backend() @@ -576,7 +596,6 @@ def get_cache_size_per_token(model_config, mapping, num_layers=None, **kwargs): ) residual = FP4_MLA_K_RESIDUAL_DIM if backend in _FP4_MLA_K_RESIDUAL_BACKENDS else 0 storage_head_dim = logical_head_dim + residual - tokens_per_block = int(kwargs["tokens_per_block"]) v_scale_page = get_fp4_mla_v_scale_pool_size(config.kv_lora_rank, tokens_per_block) per_layer = ( math.ceil(storage_head_dim / 2) @@ -588,11 +607,9 @@ def get_cache_size_per_token(model_config, mapping, num_layers=None, **kwargs): local_layers = KVCacheManager._resolve_num_attention_layers( model_config, mapping, num_layers ) - spec_config = kwargs.get("spec_config") rewind = int(spec_config.tokens_per_gen_step - 1) if spec_config else 0 hp_bytes = (HP_BLOCK_SIZE + rewind) * logical_head_dim * 2 * local_layers - max_batch_size = int(kwargs.get("max_batch_size") or 0) - return per_layer * local_layers, hp_bytes * max_batch_size * mapping.pp_size + return per_layer * local_layers, hp_bytes * (max_batch_size or 0) * mapping.pp_size __all__ = ["Fp4MlaKVCacheManagerV2", "Fp4MlaPageTableSpec"] diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py index 6759b9026996..0cfd34cb55ec 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py @@ -67,12 +67,21 @@ def _execute_fp8_context_with_cache_update( start_event.record(current_stream) # Launch the auxiliary work first so the main stream cannot drain the # attention queue before the cache update has been submitted. - with torch.cuda.stream(aux_stream): - aux_stream.wait_event(start_event) - cache_update_fn() - done_event.record(aux_stream) - attention_fn() - current_stream.wait_event(done_event) + joined = False + try: + with torch.cuda.stream(aux_stream): + aux_stream.wait_event(start_event) + cache_update_fn() + done_event.record(aux_stream) + attention_fn() + current_stream.wait_event(done_event) + joined = True + finally: + if not joined: + # Either callback may enqueue work before raising, and done_event + # may not have been recorded. Drain before scratch or KV pages can + # be released/reused; the successful path stays asynchronous. + aux_stream.synchronize() @dataclass(frozen=True, slots=True) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 622654e48b24..f33349506a9d 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -498,15 +498,10 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_PV_SFB_S2T_STRIDE_BYTES = 128 SMEM_P4_PV_SFB_CP_GROUP_ONE = False SMEM_P4_QK_COMPLETION_MBARS = SMEM_P4_TMEM_SCORE_PIPELINE_STAGES -_EXPLICIT_TORCH_STREAM: torch.cuda.Stream | None = None def _current_cu_stream() -> cuda.CUstream: - global _EXPLICIT_TORCH_STREAM - if os.environ.get("DKG_MLA_EXPLICIT_STREAM") == "1": - if _EXPLICIT_TORCH_STREAM is None: - _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() - torch.cuda.set_stream(_EXPLICIT_TORCH_STREAM) + """Use the caller's stream for launch ordering and CUDA graph capture.""" return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py index 6d582c78ce0a..c7a9d73120ad 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py @@ -500,15 +500,10 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_PV_SFB_S2T_STRIDE_BYTES = 128 SMEM_P4_PV_SFB_CP_GROUP_ONE = False SMEM_P4_QK_COMPLETION_MBARS = SMEM_P4_TMEM_SCORE_PIPELINE_STAGES -_EXPLICIT_TORCH_STREAM: torch.cuda.Stream | None = None def _current_cu_stream() -> cuda.CUstream: - global _EXPLICIT_TORCH_STREAM - if os.environ.get("DKG_MLA_EXPLICIT_STREAM") == "1": - if _EXPLICIT_TORCH_STREAM is None: - _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() - torch.cuda.set_stream(_EXPLICIT_TORCH_STREAM) + """Use the caller's stream for launch ordering and CUDA graph capture.""" return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py index 26f24caedaab..78418ab6ebfd 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py @@ -77,21 +77,8 @@ def _compile_cutedsl(*args, **kwargs): stdout_filter.finish() -_EXPLICIT_TORCH_STREAM: torch.cuda.Stream | None = None - - def _current_cu_stream() -> cuda.CUstream: - """Return the active PyTorch stream as a CUDA driver stream. - - cuda2ctl capture cannot reliably query the default null stream. Keep the - default behavior for normal pytest/AModel runs, but allow SMART wrappers to - request an explicit stream through the environment. - """ - global _EXPLICIT_TORCH_STREAM - if os.environ.get("DKG_MLA_EXPLICIT_STREAM") == "1": - if _EXPLICIT_TORCH_STREAM is None: - _EXPLICIT_TORCH_STREAM = torch.cuda.Stream() - torch.cuda.set_stream(_EXPLICIT_TORCH_STREAM) + """Use the caller's stream for launch ordering and CUDA graph capture.""" return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py index 98bc954cc56a..a9f2b83f78c4 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py @@ -31,28 +31,13 @@ import triton import triton.language as tl -_LOG2_E = tl.constexpr(1.4426950408889634) - +from .fp4_mla_kernels import _fp4_mla_swizzled_sf_offset -@triton.jit -def _fp4_mla_swizzled_sf_offset( - row_idx, - col_idx, - SF_PER_TOKEN: tl.constexpr, -): - padded_cols = ((SF_PER_TOKEN + 3) // 4) * 4 - col_in_group = col_idx % 4 - col_group = col_idx // 4 - row_in_group0 = row_idx % 32 - row_in_group1 = (row_idx % 128) // 32 - row_group = row_idx // 128 - return ( - col_in_group - + col_group * (4 * 128) - + row_in_group0 * 16 - + row_in_group1 * 4 - + row_group * (128 * padded_cols) - ) +_LOG2_E = tl.constexpr(1.4426950408889634) +# Triton 3.6.0/sm_100 fails NVWSInsertTmemAref::hasOneUse() for the +# chained residual-Q dot_scaled path. Enable only after verifying a fixed +# compiler; gate descriptor construction together with its consumer below. +_ENABLE_RESIDUAL_Q_FAST_PATH = tl.constexpr(False) @triton.jit @@ -71,30 +56,6 @@ def _fp4_mla_swizzled_sf_offset_row_block( return col_part + row_part + row_group * (128 * padded_cols) -@triton.jit -def _fp4_e2m1_quantize(x): - abs_x = tl.abs(x) - magnitude = tl.where( - abs_x < 0.25, - 0, - tl.where( - abs_x < 0.75, - 1, - tl.where( - abs_x < 1.25, - 2, - tl.where( - abs_x < 1.75, - 3, - tl.where(abs_x < 2.5, 4, tl.where(abs_x < 3.5, 5, tl.where(abs_x < 5.0, 6, 7))), - ), - ), - ), - ) - sign = tl.where(x < 0.0, 8, 0) - return (magnitude | sign).to(tl.uint8) - - @triton.jit def _fp4_pack_low_nibbles(even_packed, odd_packed): """PTX helper: pack the low nibbles of two bytes into one byte (low + high<<4).""" @@ -327,7 +288,7 @@ def _fp4_mla_attention_v_repack_pages_dyn_kernel( out_desc.store([row_base.to(tl.int32), 0], v_vals) -def fp4_mla_repack_v_cache( +def fp4_mla_repack_v_cache_triton( v_packed: Any, kv_cache: Any, page_ids: Optional[Any] = None, @@ -491,7 +452,12 @@ def _fp4_mla_qk_scores_tile( strides=[kv_s0, kv_s2, kv_s4], block_shape=[1, BLOCK_T, BLOCK_K // 2], ) - if USE_TMA_DATA_LOAD and Q_RESIDUAL_D == 64 and TAIL_BLOCK_K == 128: + if ( + _ENABLE_RESIDUAL_Q_FAST_PATH + and USE_TMA_DATA_LOAD + and Q_RESIDUAL_D == 64 + and TAIL_BLOCK_K == 128 + ): q_tail_desc = tl.make_tensor_descriptor( q_fp4_ptr, shape=[q_num_rows, Q_STORAGE_HEAD_D // 2], @@ -585,7 +551,7 @@ def _fp4_mla_qk_scores_tile( # calls into the same accumulator (one for even Q lane, one for odd), # which lowers to a TMEM alloc with multiple uses and trips # NVWSInsertTmemAref::hasOneUse() on Triton 3.6.0 / sm_100. - if False and Q_RESIDUAL_D == 64 and TAIL_BLOCK_K == 128: + if _ENABLE_RESIDUAL_Q_FAST_PATH and Q_RESIDUAL_D == 64 and TAIL_BLOCK_K == 128: residual_packed_offsets = tl.arange(0, 32) residual_scale_offsets = tl.arange(0, 4) packed_k_cols = non_residual_groups * (FP4_BLOCK // 2) + residual_packed_offsets diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 64f4f6c21706..ba51b5ec1e1a 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -660,7 +660,7 @@ def _post_init_with_buffers(self, buffers) -> None: num_local_layers = self.kv_cache_manager.num_local_layers head_dim = self.kv_cache_manager.head_dim kv_factor = self.kv_cache_manager.kv_factor - hp_ring_size = int(self.kv_cache_manager._fp4_mla_hp_pool_size) + hp_ring_size = self.kv_cache_manager.fp4_mla_hp_pool_size if hp_ring_size < HP_BLOCK_SIZE: raise RuntimeError( "FP4 MLA high-precision ring must retain at least one " @@ -2586,49 +2586,6 @@ def forward( assert metadata.kv_cache_manager is None assert metadata.num_contexts == metadata.num_seqs - # Testing only: ``mla_rope_generation`` normally rotates q_pe, appends the - # new latent to the paged cache, and fills the trtllm-gen scheduler - # buffers (cumulative q/kv seqlens + the FMHA scheduler counter). When the - # harness sets ``skip_mla_rope_generation`` it feeds a pre-RoPE'd fused_q, - # so we skip only the RoPE and do the append + scheduler init here: the - # generation FMHA only reads the cache, and the fallback path needs the - # scheduler buffers (the flashinfer trtllm-gen decode kernel ignores them). - if (self.is_mla_enable and forward_args.skip_mla_rope_generation - and forward_args.attention_input_type - == AttentionInputType.generation_only): - num_ctx = metadata.num_contexts - n_gen = metadata.num_generations - # Use the GPU-resident length tensors (no host->device copy) so this - # stays CUDA-graph-capturable. - gen_q_lens = metadata.seq_lens_cuda[num_ctx:num_ctx + n_gen].to( - torch.int32) - gen_kv_lens = metadata.kv_lens_cuda_runtime[num_ctx:num_ctx + - n_gen].to(torch.int32) - cu_q = torch.zeros(n_gen + 1, dtype=torch.int32, device=q.device) - cu_kv = torch.zeros(n_gen + 1, dtype=torch.int32, device=q.device) - cu_q[1:] = torch.cumsum(gen_q_lens, dim=0).to( - torch.int32) * self.num_heads - cu_kv[1:] = torch.cumsum(gen_kv_lens, dim=0).to(torch.int32) - forward_args.cu_q_seqlens = cu_q - forward_args.cu_kv_seqlens = cu_kv - if forward_args.fmha_scheduler_counter is None: - forward_args.fmha_scheduler_counter = torch.zeros( - 1, dtype=torch.uint32, device=q.device) - else: - forward_args.fmha_scheduler_counter.zero_() - assert forward_args.latent_cache is not None - from .utils import append_mla_latent_cache - append_mla_latent_cache( - metadata.kv_cache_manager, - self.get_local_layer_idx(metadata), - metadata.request_ids, - metadata.seq_lens.tolist(), - metadata.kv_cache_params.num_cached_tokens_per_seq, - forward_args.latent_cache, - kv_layout=metadata.kv_layout, - seq_start=num_ctx, - ) - fmha = self._fmha_manager.select(self, q, k, v, metadata, forward_args) if fmha is None: @@ -3026,6 +2983,9 @@ def _fp4_mla_rope_generation( if not fuse_q_quant: raise RuntimeError( "FP4 MLA generation requires fused Q quantization.") + if not self.rope_params.duplicate_data: + raise RuntimeError( + "FP4 MLA requires RopeParams.duplicate_data=True.") if not self.can_fuse_fp4_mla_q_quant(fused_q, q_pe, latent_cache, metadata): raise RuntimeError( diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index 5ed998391d17..64d9c8de7b1f 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -697,6 +697,20 @@ def create_weights(self): and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales() ) mla_weight_dtype = torch.float8_e4m3fn if has_fp8_block_scales else self.dtype + if isinstance(self.mqa, TrtllmAttention) and self.mqa.has_fp4_kv_cache: + if ( + mla_weight_dtype != torch.bfloat16 + or self.mapping.cp_size != 1 + or self.apply_rotary_emb + or self.kv_cache_dtype == "fp8_ds_mla" + ): + raise ValueError( + "FP4 MLA requires BF16 absorption weights and fused RoPE " + "on the TRTLLM backend without context parallelism or " + "fp8_ds_mla cache packing." + ) + if not self.mqa.rope_params.duplicate_data: + raise ValueError("FP4 MLA requires RopeParams.duplicate_data=True.") self.k_b_proj_trans = nn.Parameter( torch.empty( (self.num_heads_tp, self.kv_lora_rank, self.qk_nope_head_dim), @@ -1580,15 +1594,10 @@ def forward_absorption_generation( ) fp4_mla = isinstance(self.mqa, TrtllmAttention) and self.mqa.has_fp4_kv_cache - if fp4_mla and ( - self.k_b_proj_trans.dtype != torch.bfloat16 - or latent_cache is None - or self.mapping.has_cp_helix() - ): + if fp4_mla and latent_cache is None: raise RuntimeError( - "FP4 MLA generation requires fused BF16 Q " - "quantization, RoPE, and cache update on the TRT-LLM " - "backend." + "FP4 MLA generation requires a latent cache for fused " + "Q quantization, RoPE, and cache update." ) fuse_fp4_q_quant = fp4_mla fp4_rope_kwargs = {"fuse_fp4_q_quant": True} if fuse_fp4_q_quant else {} @@ -1625,7 +1634,10 @@ def _mla_gen_rope(): **fp4_rope_kwargs, ) - rope_stream = None if fuse_fp4_q_quant else self.aux_stream if not use_fp8_mla else None + if fuse_fp4_q_quant or use_fp8_mla: + rope_stream = None + else: + rope_stream = self.aux_stream if self.k_b_proj_trans.dtype == torch.bfloat16: # [num_heads, num_tokens, self.qk_nope_head_dim] q_nope_t = q_nope.transpose(0, 1) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index c2955d67bcc4..2cbe479c6254 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -191,6 +191,10 @@ def get_kv_cache_manager_cls( quant_config = getattr(model_config, "quant_config", None) if (is_mla(config) and quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache()): + if kv_cache_config.use_kv_cache_manager_v2 is False: + raise ValueError("FP4 MLA requires use_kv_cache_manager_v2=True.") + if model_config.attn_backend != "TRTLLM": + raise ValueError("FP4 MLA requires the TRTLLM attention backend.") if is_disagg: raise NotImplementedError( "FP4 MLA disaggregated serving requires the follow-up " @@ -200,6 +204,8 @@ def get_kv_cache_manager_cls( "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " "support hybrid linear-attention models.") if sparse_attn_config is not None: + sparse_attn_algorithm = (sparse_attn_algorithm + or type(sparse_attn_config).__name__) raise NotImplementedError( "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " f"support sparse attention algorithm {sparse_attn_algorithm!r}." diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 2e10388fdd6d..268cf0f08c07 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -3,7 +3,7 @@ import dataclasses from collections.abc import Mapping as AbcMapping -from typing import List, Optional, Sequence +from typing import TYPE_CHECKING, List, Optional, Sequence import torch import transformers @@ -12,6 +12,9 @@ from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.logger import logger +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + def resolve_cache_transceiver_config( cache_transceiver_config: Optional[CacheTransceiverConfig]) -> None: @@ -367,6 +370,22 @@ def is_mla(config): return False +def supports_fp4_mla_attention(model_config: "ModelConfig") -> bool: + """Whether the model uses the dedicated dense TRTLLM FP4 MLA path.""" + return (is_mla(model_config.pretrained_config) + and model_config.attn_backend == "TRTLLM" + and model_config.sparse_attention_config is None + and not is_hybrid_linear(model_config.pretrained_config)) + + +def uses_fp4_mla_attention(model_config: "ModelConfig") -> bool: + """Use the resolved quantization, never the requested cache dtype.""" + quant_config = getattr(model_config, "quant_config", None) + return (quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache() + and supports_fp4_mla_attention(model_config)) + + def is_minimax_m3(sparse_attention_config): """True when the sparse attention config selects the MiniMax-M3 algorithm.""" return sparse_attention_config is not None and sparse_attention_config.algorithm == "minimax_m3" diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 43ea9ea83296..c8386d9de9ce 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6105,14 +6105,6 @@ def _forward_scheduled(self, scheduled_requests: ScheduledRequests, LlmRequest]]): assert not self._is_packed_runner, ( "a packed-batch runner cannot execute scheduled requests") - if not getattr(self, "_disable_overlap_scheduler", True): - # Do not refill reusable host staging while the previous - # iteration's asynchronous H2D copies still consume it. This - # event precedes the model forward, so its synchronization retains - # GPU forward overlap while protecting every forward entry, - # including speculative/draft paths. - self.wait_for_input_copy() - kv_cache_manager = resource_manager.get_resource_manager( self.kv_cache_manager_key) if self._runner is not None and not self._is_encoder_decoder_model(): diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 8642943c4ab5..295e8f30d89b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -54,7 +54,8 @@ maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope -from .config_utils import (is_hybrid_linear, resolve_auto_ssm_cache_dtype, +from .config_utils import (is_hybrid_linear, is_mla, resolve_auto_ssm_cache_dtype, + supports_fp4_mla_attention, uses_fp4_mla_attention, validate_kimi_kda_state_dtype) _KV_CACHE_MAP = { @@ -161,9 +162,9 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, effective_kv_cache_quant = (kv_cache_quant if pyt_kv_cache_dtype == "auto" else mapped_pyt_quant) - if (torch.cuda.is_available() and get_sm_version() == 107 - and effective_kv_cache_quant - in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value)): + if (effective_kv_cache_quant in (QuantAlgo.NVFP4, QuantAlgo.NVFP4.value) + and not supports_fp4_mla_attention(model_config) + and torch.cuda.is_available() and get_sm_version() == 107): logger.warning( "NVFP4 KV cache is not supported by trtllm-gen on SM107; " "using FP8 KV cache instead.") @@ -207,6 +208,32 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, layer_quant_config.kv_cache_quant_algo = mapped_pyt_quant +def validate_fp4_mla_config(model_config: ModelConfig, + llm_args: TorchLlmArgs) -> None: + """Validate FP4 MLA before model construction and KV-cache allocation.""" + if not (is_mla(model_config.pretrained_config) + and model_config.quant_config.quant_mode.has_fp4_kv_cache()): + return + if not supports_fp4_mla_attention(model_config): + raise ValueError( + "FP4 MLA requires the TRTLLM attention backend with dense MLA; " + "sparse and hybrid linear attention are not supported.") + if model_config.mapping.cp_size != 1: + raise ValueError("FP4 MLA does not support context parallelism.") + if llm_args.kv_cache_config.use_kv_cache_manager_v2 is False: + raise ValueError("FP4 MLA requires use_kv_cache_manager_v2=True.") + if llm_args.enable_chunked_prefill: + raise ValueError( + "FP4 MLA does not support chunked prefill or cached context; " + "set enable_chunked_prefill=False.") + + spec_config = model_config.spec_config + if spec_config is not None and spec_config.spec_dec_mode.use_one_engine(): + # HP/V-scale side pools are not switched by draft_kv_cache_context. + # Set this before models and workers consult the shared predicate. + spec_config._allow_separate_draft_kv_cache = False + + def validate_encoder_decoder_kv_cache_config(model_config: ModelConfig, kv_cache_config) -> None: """Validate encoder-decoder KV-cache requirements for the PyTorch runtime. @@ -803,6 +830,21 @@ def load_config_and_apply_defaults( # Resolve "auto" sentinel values after model defaults are applied. _resolve_transceiver_runtime_auto(llm_args, preference_cls, config.pretrained_config) + if (original_kv_cache_manager_setting == "auto" and + (llm_args.kv_cache_config.dtype == "nvfp4" or + (llm_args.kv_cache_config.dtype == "auto" + and config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4))): + # Resolve FP4 MLA before the generic model preference turns auto + # into False. Keep an explicit False distinguishable and reject it + # during FP4 validation instead of silently overriding the user. + fp4_mla_config = copy.copy(config) + fp4_mla_config.attn_backend = llm_args.attn_backend + fp4_mla_config.sparse_attention_config = llm_args.sparse_attention_config + if supports_fp4_mla_attention(fp4_mla_config): + validate_and_set_kv_cache_quant(fp4_mla_config, + llm_args.kv_cache_config.dtype) + if uses_fp4_mla_attention(fp4_mla_config): + llm_args.kv_cache_config.use_kv_cache_manager_v2 = True _resolve_kv_cache_manager_v2_auto(llm_args, preference_cls, config.pretrained_config) _validate_and_adjust_mamba_snapshot_config(config, llm_args) @@ -1922,6 +1964,7 @@ def _load_and_validate_config( self.llm_args.kv_cache_config) validate_and_set_kv_cache_quant(config, self.llm_args.kv_cache_config.dtype) + validate_fp4_mla_config(config, self.llm_args) validate_and_set_mamba_ssm_cache_dtype( config, self.llm_args.kv_cache_config.mamba_ssm_cache_dtype, self.llm_args.kv_cache_config.mamba_ssm_stochastic_rounding, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index a1e2c4c2485d..6e611831c942 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -43,7 +43,7 @@ validate_feature_combination) from .config_utils import (is_hybrid_linear, is_minimax_m3, resolve_cache_transceiver_config, - uses_vswa_kv_cache_layout) + uses_fp4_mla_attention, uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager, get_global_dwdp_manager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder @@ -198,45 +198,6 @@ def _set_model_engines_cache_reuse(model_engines, cache_reuse: bool): engine.attn_runtime_features.cache_reuse = cache_reuse -def _has_fp4_kv_cache(model_config, kv_cache_config) -> bool: - kv_cache_quant_algo = getattr(getattr(model_config, "quant_config", None), - "kv_cache_quant_algo", None) - fp4_quant_values = { - QuantAlgo.NVFP4, - getattr(QuantAlgo.NVFP4, "value", None), - "NVFP4", - } - kv_cache_dtype = getattr(kv_cache_config, "dtype", None) - return ((isinstance(kv_cache_dtype, str) - and kv_cache_dtype.lower() == "nvfp4") - or (isinstance(kv_cache_quant_algo, str) - and kv_cache_quant_algo.upper() == "NVFP4") - or kv_cache_quant_algo in fp4_quant_values) - - -def _uses_fp4_mla_attention(model_config, kv_cache_config) -> bool: - return (_has_fp4_kv_cache(model_config, kv_cache_config) - and getattr(model_config, "attn_backend", None) == "TRTLLM") - - -def _configure_fp4_mla_speculative_kv_cache(spec_config, config, model_config, - kv_cache_config, model) -> None: - if (spec_config is None or not spec_config.spec_dec_mode.use_one_engine() - or not is_mla(config) - or not _uses_fp4_mla_attention(model_config, kv_cache_config)): - return - - # FP4 MLA side pools live on attention metadata and are not switched by - # draft_kv_cache_context. Keep target and one-model draft layers in one - # manager so their global layer IDs address distinct HP/V-scale state. - spec_config._allow_separate_draft_kv_cache = False - if hasattr(model, 'use_separate_draft_kv_cache'): - model.use_separate_draft_kv_cache = False - spec_worker = getattr(model, 'spec_worker', None) - if spec_worker is not None: - spec_worker.use_separate_draft_kv_cache = False - - def _get_mapping(_mapping: Mapping) -> Mapping: if _mapping is None: mapping = Mapping(world_size=tensorrt_llm.mpi_world_size(), @@ -674,13 +635,6 @@ def allocation_scope(current_stage: ExecutorMemoryType): resolve_cache_transceiver_config(cache_transceiver_config) config = model_engine.model.model_config.pretrained_config - _configure_fp4_mla_speculative_kv_cache( - spec_config, - config, - model_engine.model.model_config, - kv_cache_config, - model_engine.model, - ) max_num_seq_slots = getattr( model_engine, "max_num_seq_slots", None) or compute_max_num_sequences( mapping, @@ -689,30 +643,35 @@ def allocation_scope(current_stage: ExecutorMemoryType): enable_overlap_headroom=getattr(model_engine, "_enable_overlap_headroom", False)) if is_mla(config): - if _uses_fp4_mla_attention(model_engine.model.model_config, - kv_cache_config): + if uses_fp4_mla_attention(model_engine.model.model_config): tokens_per_block = FP4_MLA_TOKENS_PER_BLOCK kv_cache_config.tokens_per_block = tokens_per_block logger.info( f"Change tokens_per_block to: {tokens_per_block} for using FP4 MLA attention" ) - else: - if model_engine.model.model_config.enable_flash_mla: - tokens_per_block = 64 - # Propagate the override back to kv_cache_config so any consumer - # that later reads llm_args.kv_cache_config.tokens_per_block sees - # the effective value. KvCacheConnectorScheduler subclasses - # (LMCache, Dynamo KVBM) are instantiated further down via - # scheduler_cls(llm_args) and size their block pools from - # llm_args.kv_cache_config.tokens_per_block. Without this the - # connector's block size desynced from the KVCacheManager's - # actual tokens_per_block (user-set or default 32 vs. FlashMLA's - # forced 64), producing a frozen cache_block_ids view to the - # connector and silently-corrupted decode KV (#13320). - kv_cache_config.tokens_per_block = tokens_per_block - logger.info( - f"Change tokens_per_block to: {tokens_per_block} for using FlashMLA" - ) + if kv_cache_config.enable_block_reuse: + logger.warning( + "FP4 MLA cached-context attention is not supported yet; " + "disabling KV cache block reuse.") + kv_cache_config.enable_block_reuse = False + _set_model_engines_cache_reuse( + [model_engine, draft_model_engine], False) + elif model_engine.model.model_config.enable_flash_mla: + tokens_per_block = 64 + # Propagate the override back to kv_cache_config so any consumer + # that later reads llm_args.kv_cache_config.tokens_per_block sees + # the effective value. KvCacheConnectorScheduler subclasses + # (LMCache, Dynamo KVBM) are instantiated further down via + # scheduler_cls(llm_args) and size their block pools from + # llm_args.kv_cache_config.tokens_per_block. Without this the + # connector's block size desynced from the KVCacheManager's + # actual tokens_per_block (user-set or default 32 vs. FlashMLA's + # forced 64), producing a frozen cache_block_ids view to the + # connector and silently-corrupted decode KV (#13320). + kv_cache_config.tokens_per_block = tokens_per_block + logger.info( + f"Change tokens_per_block to: {tokens_per_block} for using FlashMLA" + ) sm_version = get_sm_version() if (kv_cache_config.enable_block_reuse and sm_version diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 59553f235cc8..c63cd14d93ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1886,24 +1886,17 @@ def get_buffers(self, result = self.impl.get_primary_pool_data(layer_offset) pool = self.get_pool_for_layer(layer_offset) - layer_dtype = pool.dtype if pool else self.dtype layer_head_dim = pool.head_dim if pool else self.head_dim assert kv_layout in ["NHD", "HND"], f"Unsupported kv_layout: {kv_layout}" - - element_per_container = 1 - if layer_dtype == DataType.NVFP4: - element_per_container = 2 - effective_head_dim = layer_head_dim // element_per_container - if kv_layout == "NHD": return result.reshape( result.shape[0], self.kv_factor, self.tokens_per_block, self.num_kv_heads_per_layer[layer_offset], - effective_head_dim, + layer_head_dim, ) else: return result.reshape( @@ -1911,7 +1904,7 @@ def get_buffers(self, self.kv_factor, self.num_kv_heads_per_layer[layer_offset], self.tokens_per_block, - effective_head_dim, + layer_head_dim, ) def get_indexer_k_cache_pool_data(self, layer_idx: int) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 10caf2c61667..636d639790f6 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -945,12 +945,9 @@ def change_attn_metadata(self, num_accepted_tokens: torch.Tensor, attn_metadata.kv_lens_cuda[num_contexts:batch_size].clamp_( min=runtime_draft_len) attn_metadata.on_update_kv_lens() - attn_metadata.update_for_spec_dec() - elif getattr(attn_metadata, "kv_lens_cuda_runtime", None) is not None: - attn_metadata.kv_lens_cuda_runtime[num_contexts:batch_size] -= ( - runtime_draft_len + 1 - - num_accepted_tokens[num_contexts:batch_size]) - attn_metadata.update_for_spec_dec() + if getattr(attn_metadata, "high_precision_kv_pool", + None) is not None: + attn_metadata.update_for_spec_dec() if attn_metadata.kv_cache_params is not None and not attn_metadata.is_cuda_graph: for i in range(num_contexts, batch_size): diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index e3c050662653..0ce6cde395b2 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -300,20 +300,50 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, @skip_no_rubin @pytest.mark.skip_less_device_memory(60000) - def test_nvfp4_mla_gsm8k(self): + def test_nvfp4_mla_gsm8k(self, monkeypatch): + from tensorrt_llm._torch.attention.backends.fmha.fp4_mla import \ + Fp4MlaFmha + from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import \ + Fp4MlaKVCacheManagerV2 + from tensorrt_llm._torch.attention.backends.trtllm import \ + TrtllmAttention + from tensorrt_llm._torch.attention.mla import MLA + + # Inspect the actual worker's quantization, cache pools and FMHA choice. + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") kv_cache_config = KvCacheConfig(dtype="nvfp4", - free_gpu_memory_fraction=0.75) + free_gpu_memory_fraction=0.75, + use_kv_cache_manager_v2=True, + enable_block_reuse=False) with LLM( f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only", attn_backend="TRTLLM", kv_cache_config=kv_cache_config, + enable_chunked_prefill=False, max_num_tokens=8192, max_batch_size=1350, ) as llm: - assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 + executor = llm._executor.engine + quant_config = executor.model_engine.model.model_config.quant_config + assert quant_config.quant_algo == QuantAlgo.NVFP4 + assert quant_config.quant_mode.has_fp4_kv_cache() + manager = executor.kv_cache_manager + assert isinstance(manager, Fp4MlaKVCacheManagerV2) + assert manager.get_mla_v_scale_pool().numel() > 0 + assert manager.get_fp4_mla_hp_pool().numel() > 0 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + mla_layers = [ + module.mqa for module in executor.model_engine.model.modules() + if isinstance(module, MLA) + ] + assert mla_layers + for layer in mla_layers: + assert isinstance(layer, TrtllmAttention) + assert layer.has_fp4_kv_cache + assert any( + isinstance(fmha, Fp4MlaFmha) + for fmha in layer._fmha_manager._cache.values()) @pytest.mark.skip_less_device_memory(60000) @parametrize_with_ids("enable_chunked_prefill", [False, True]) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 27921e20429d..d360fb6bc84f 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -778,6 +778,9 @@ test_e2e.py::test_relaxed_acceptance_quickstart_advanced_deepseek_r1_8gpus[DeepS test_e2e.py::test_trtllm_benchmark_serving[gpt_oss/gpt-oss-20b] test_e2e.py::test_trtllm_multimodal_benchmark_serving +# Dense FP4 MLA accuracy (SM107 only) +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_mla_gsm8k + # fine-grained sync tests (SM107+ only) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_fine_grained_sync[enable_autotuner=False-moe_backend=TRTLLM-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_fine_grained_sync[enable_autotuner=False-moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index e95caab3d2da..b23358200999 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -28,11 +28,7 @@ PredefinedAttentionMask, RopeParams, ) -from tensorrt_llm._torch.attention.backends.utils import ( - append_mla_latent_cache, - create_attention, - get_attention_backend, -) +from tensorrt_llm._torch.attention.backends.utils import create_attention, get_attention_backend from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -465,8 +461,7 @@ def _run_mla_gen_backend( """Run one backend's absorbed-MLA generation; return [nnz_q, heads*kv_lora]. No ``mla_rope_generation`` call (see ``generate_mla_gen_inputs``): ``fused_q`` - is passed as-is. The harness prepares TRTLLM's cache and scheduler inputs - directly, while the other backends retain their normal Python cache update. + is passed as-is and the backend's ``forward`` appends the new latent + MQA. """ AttentionCls = get_attention_backend(backend) H = case.num_heads @@ -495,41 +490,20 @@ def _run_mla_gen_backend( expected_latents = _split_packed_tokens(latent_cache, case.seq_lens) def _forward(metadata, q, q_pe): - forward_args = AttentionForwardArgs( - latent_cache=latent_cache, - q_pe=q_pe, - attention_input_type=AttentionInputType.generation_only, - ) - if backend == "TRTLLM": - num_ctx = metadata.num_contexts - num_gen = metadata.num_generations - gen_q_lens = metadata.seq_lens_cuda[num_ctx : num_ctx + num_gen].to(torch.int32) - gen_kv_lens = metadata.kv_lens_cuda_runtime[num_ctx : num_ctx + num_gen].to(torch.int32) - cu_q_seqlens = torch.zeros(num_gen + 1, dtype=torch.int32, device=q.device) - cu_kv_seqlens = torch.zeros(num_gen + 1, dtype=torch.int32, device=q.device) - cu_q_seqlens[1:] = torch.cumsum(gen_q_lens, dim=0).to(torch.int32) * H - cu_kv_seqlens[1:] = torch.cumsum(gen_kv_lens, dim=0).to(torch.int32) - forward_args.cu_q_seqlens = cu_q_seqlens - forward_args.cu_kv_seqlens = cu_kv_seqlens - forward_args.fmha_scheduler_counter = torch.zeros( - 1, dtype=torch.uint32, device=q.device - ) - append_mla_latent_cache( - metadata.kv_cache_manager, - attn.get_local_layer_idx(metadata), - metadata.request_ids, - metadata.seq_lens.tolist(), - metadata.kv_cache_params.num_cached_tokens_per_seq, - latent_cache, - kv_layout=metadata.kv_layout, - seq_start=num_ctx, - ) out = attn.forward( q, None, None, metadata, - forward_args=forward_args, + forward_args=AttentionForwardArgs( + latent_cache=latent_cache, + q_pe=q_pe, + attention_input_type=AttentionInputType.generation_only, + # The harness feeds a pre-RoPE'd fused_q, so skip the RoPE step; + # the TRTLLM backend still appends the new latent and inits its + # scheduler buffers. Vanilla/FlashInfer ignore this flag. + skip_mla_rope_generation=True, + ), ) return out[0] if isinstance(out, tuple) else out diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index 7134df332060..3b3572e6aca6 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -29,15 +29,17 @@ ) from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, MTPDecodingConfig from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig -from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig from tensorrt_llm.mapping import Mapping _DataType = tensorrt_llm.bindings.DataType _CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType -def _swizzled_sf_offset(row_idx: int, col_idx: int, sf_per_token: int) -> int: +def _swizzled_sf_offset( + row_idx: int | torch.Tensor, col_idx: int | torch.Tensor, sf_per_token: int +) -> int | torch.Tensor: padded_cols = ((sf_per_token + 3) // 4) * 4 return ( col_idx % 4 @@ -284,37 +286,17 @@ def _dequant_fp4_swizzled( ) fp4_bytes = fp4_tensor.view(torch.uint8) sf_flat = sf_tensor.view(torch.float8_e4m3fn).reshape(-1) - out = torch.empty( - (fp4_bytes.shape[0], logical_dim), - dtype=torch.float32, - device=fp4_tensor.device, - ) - - for row_idx in range(fp4_bytes.shape[0]): - for sf_col in range(sf_per_token): - start = sf_col * FP4_BLOCK_SIZE - packed = fp4_bytes[row_idx, start // 2 : start // 2 + 8] - low = packed & 0x0F - high = (packed >> 4) & 0x0F - vals = torch.empty(FP4_BLOCK_SIZE, dtype=torch.float32, device=fp4_tensor.device) - low_sign = torch.where( - (low & 0x08) != 0, - -torch.ones_like(low, dtype=torch.float32), - torch.ones_like(low, dtype=torch.float32), - ) - high_sign = torch.where( - (high & 0x08) != 0, - -torch.ones_like(high, dtype=torch.float32), - torch.ones_like(high, dtype=torch.float32), - ) - vals[0::2] = fp4_values[(low & 0x07).long()] * low_sign - vals[1::2] = fp4_values[(high & 0x07).long()] * high_sign - sf_offset = _swizzled_sf_offset(row_idx, sf_col, sf_per_token) - out[row_idx, start : start + FP4_BLOCK_SIZE] = ( - vals * sf_flat[sf_offset].float() / global_scale - ) + # Unpack all rows at once, preserving low-nibble/high-nibble order. + packed = fp4_bytes[:, : (logical_dim + 1) // 2] + nibbles = torch.stack((packed & 0x0F, packed >> 4), dim=-1) + nibbles = nibbles.flatten(1)[:, :logical_dim] + values = fp4_values[(nibbles & 0x07).long()] + values = torch.where((nibbles & 0x08) != 0, -values, values) - return out + rows = torch.arange(fp4_bytes.shape[0], device=fp4_tensor.device)[:, None] + cols = torch.arange(logical_dim, device=fp4_tensor.device)[None, :] // FP4_BLOCK_SIZE + offsets = _swizzled_sf_offset(rows, cols, sf_per_token) + return values * sf_flat.float()[offsets] / global_scale def _duplicate_tail_groups(tensor: torch.Tensor, residual_dim: int) -> torch.Tensor: @@ -353,7 +335,7 @@ def _create_fp4_mla_v2_manager( max_tokens: int, max_seq_len: int, max_batch_size: int, - spec_config=None, + spec_config: DecodingBaseConfig | None = None, enable_block_reuse: bool = False, ) -> Fp4MlaKVCacheManagerV2: return Fp4MlaKVCacheManagerV2( @@ -378,7 +360,9 @@ def _create_fp4_mla_v2_manager( ) -def _build_multi_seq_metadata(kv_cache_manager, *, seq_lens, page_size): +def _build_multi_seq_metadata( + kv_cache_manager: Fp4MlaKVCacheManagerV2, *, seq_lens: list[int], page_size: int +) -> SimpleNamespace: device = torch.device("cuda") num_seqs = len(seq_lens) request_ids = list(range(num_seqs)) @@ -565,12 +549,12 @@ def _materialize_reference_cache_tokens( def _build_fp4_mla_attention_decode_case( *, - seq_lens, - num_heads, - seed, - query_len_per_seq=1, - enable_block_reuse=False, -): + seq_lens: list[int], + num_heads: int, + seed: int, + query_len_per_seq: int = 1, + enable_block_reuse: bool = False, +) -> tuple[Fp4MlaKVCacheManagerV2, SimpleNamespace, torch.Tensor, int, int]: torch.manual_seed(seed) device = torch.device("cuda") @@ -585,7 +569,9 @@ def _build_fp4_mla_attention_decode_case( context_seq_lens = [seq_len - query_len_per_seq for seq_len in seq_lens] if min(context_seq_lens) <= 0: raise ValueError("FP4 MLA decode cases require a non-empty context for every sequence.") - spec_config = MTPDecodingConfig(max_draft_len=3) if query_len_per_seq == 4 else None + spec_config = ( + MTPDecodingConfig(max_draft_len=query_len_per_seq - 1) if query_len_per_seq > 1 else None + ) kv_cache_manager = _create_fp4_mla_v2_manager( max_tokens=max_tokens, @@ -743,13 +729,13 @@ def _build_fp4_mla_attention_decode_case( def _fp4_mla_attention_decode_reference( - metadata, - q_nope, - q_pe, + metadata: SimpleNamespace, + q_nope: torch.Tensor, + q_pe: torch.Tensor, *, - sm_scale, - kv_lora_rank, - qk_rope_head_dim, + sm_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, ) -> torch.Tensor: head_dim = kv_lora_rank + qk_rope_head_dim storage = _materialize_reference_cache_storage(metadata, 0, head_dim) diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 19bfe98eebb8..f4c9b06c5d84 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import json import struct import types @@ -6,10 +9,13 @@ import torch from tensorrt_llm._torch.model_config import _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT, ModelConfig +from tensorrt_llm._torch.pyexecutor.config_utils import uses_fp4_mla_attention from tensorrt_llm._torch.pyexecutor.model_loader import ( validate_and_set_kv_cache_quant, validate_encoder_decoder_kv_cache_config, + validate_fp4_mla_config, ) +from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -144,13 +150,53 @@ def _mock_sm107(monkeypatch: pytest.MonkeyPatch) -> None: ) -def test_validate_and_set_kv_cache_quant_downgrades_nvfp4_on_sm107( +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "nvfp4"]) +@pytest.mark.parametrize("chunked_prefill", [False, True]) +@pytest.mark.parametrize( + "mla,backend,sparse,hybrid,expected", + [ + (False, "TRTLLM", False, False, QuantAlgo.FP8), + (True, "TRTLLM", False, False, QuantAlgo.NVFP4), + (True, "FLASHINFER", False, False, QuantAlgo.FP8), + (True, "TRTLLM", True, False, QuantAlgo.FP8), + (True, "TRTLLM", False, True, QuantAlgo.FP8), + ], +) +def test_validate_and_set_kv_cache_quant_resolves_nvfp4_on_sm107( monkeypatch: pytest.MonkeyPatch, + kv_cache_dtype: str, + chunked_prefill: bool, + mla: bool, + backend: str, + sparse: bool, + hybrid: bool, + expected: QuantAlgo, ) -> None: _mock_sm107(monkeypatch) - model_config = _make_model_config_with_kv_quant(QuantAlgo.FP8) - validate_and_set_kv_cache_quant(model_config, "nvfp4") - assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + model_config = ModelConfig( + pretrained_config=types.SimpleNamespace( + kv_lora_rank=512 if mla else None, + qk_rope_head_dim=64 if mla else None, + hybrid_override_pattern="M*" if hybrid else None, + ), + attn_backend=backend, + sparse_attention_config=DeepSeekSparseAttentionConfig() if sparse else None, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4), + quant_config_dict={"layer": QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4)}, + ) + validate_and_set_kv_cache_quant(model_config, kv_cache_dtype) + assert model_config.quant_config.kv_cache_quant_algo == expected + assert model_config.quant_config_dict["layer"].kv_cache_quant_algo == expected + assert uses_fp4_mla_attention(model_config) == (expected == QuantAlgo.NVFP4) + llm_args = types.SimpleNamespace( + enable_chunked_prefill=chunked_prefill, + kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2=True), + ) + if expected == QuantAlgo.NVFP4 and chunked_prefill: + with pytest.raises(ValueError, match="does not support chunked prefill"): + validate_fp4_mla_config(model_config, llm_args) + else: + validate_fp4_mla_config(model_config, llm_args) def _make_mixed_precision_model_config(): From 1f477dd3a5ee830ae42045bb57ce53791b1299df Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Thu, 10 Sep 2026 04:03:03 -0700 Subject: [PATCH 12/21] fix: accept integer quant modes in FP4 MLA routing Normalize integer quant modes while preserving QuantModeWrapper behavior in the FP4 KV-cache predicate. Share the predicate across routing, validation, and FP4 cache guards, and extend existing tests to cover integer, enum, and wrapper representations. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 15 +++++++-------- tensorrt_llm/_torch/pyexecutor/config_utils.py | 15 +++++++++++++-- tensorrt_llm/_torch/pyexecutor/model_loader.py | 5 +++-- .../_torch/pyexecutor/resource_manager.py | 5 ++--- .../defs/accuracy/test_llm_api_pytorch.py | 3 ++- tests/unittest/_torch/test_model_config.py | 10 ++++++++++ 6 files changed, 37 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 2cbe479c6254..eae939502fda 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -55,10 +55,11 @@ from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, extract_mamba_kv_cache_params, extract_qwen4_exp_ple_cache_params, - get_layer_attention_window, is_gemma4_hybrid, - is_hybrid_linear, is_kimi_linear, is_mla, - is_nemotron_hybrid, is_qwen3_hybrid, is_qwen4_exp, - resolve_vocab_size, uses_vswa_kv_cache_layout) + get_layer_attention_window, has_fp4_kv_cache, + is_gemma4_hybrid, is_hybrid_linear, is_kimi_linear, + is_mla, is_nemotron_hybrid, is_qwen3_hybrid, + is_qwen4_exp, resolve_vocab_size, + uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder @@ -189,8 +190,7 @@ def get_kv_cache_manager_cls( sparse_attn_config = model_config.sparse_attention_config sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) quant_config = getattr(model_config, "quant_config", None) - if (is_mla(config) and quant_config is not None - and quant_config.quant_mode.has_fp4_kv_cache()): + if is_mla(config) and has_fp4_kv_cache(quant_config): if kv_cache_config.use_kv_cache_manager_v2 is False: raise ValueError("FP4 MLA requires use_kv_cache_manager_v2=True.") if model_config.attn_backend != "TRTLLM": @@ -913,8 +913,7 @@ def _validate_or_fallback_kv_cache_manager_v2( f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") quant_config = getattr(model_config, "quant_config", None) - if (is_mla(config) and quant_config is not None - and quant_config.quant_mode.has_fp4_kv_cache()): + if is_mla(config) and has_fp4_kv_cache(quant_config): raise NotImplementedError( "FP4 MLA requires Fp4MlaKVCacheManagerV2, which is " f"not yet supported with {incompat_str}. Disable these " diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 268cf0f08c07..9ae4beed1cc5 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -11,9 +11,11 @@ from tensorrt_llm._utils import str_dtype_to_torch from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantMode if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.models.modeling_utils import QuantConfig def resolve_cache_transceiver_config( @@ -378,11 +380,20 @@ def supports_fp4_mla_attention(model_config: "ModelConfig") -> bool: and not is_hybrid_linear(model_config.pretrained_config)) +def has_fp4_kv_cache(quant_config: Optional["QuantConfig"]) -> bool: + """Check resolved KV quantization for integer modes and Python wrappers.""" + if quant_config is None: + return False + quant_mode = quant_config.quant_mode + if isinstance(quant_mode, int): + quant_mode = QuantMode(quant_mode) + return quant_mode.has_fp4_kv_cache() + + def uses_fp4_mla_attention(model_config: "ModelConfig") -> bool: """Use the resolved quantization, never the requested cache dtype.""" quant_config = getattr(model_config, "quant_config", None) - return (quant_config is not None - and quant_config.quant_mode.has_fp4_kv_cache() + return (has_fp4_kv_cache(quant_config) and supports_fp4_mla_attention(model_config)) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 295e8f30d89b..f9da30f170dc 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -54,7 +54,8 @@ maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope -from .config_utils import (is_hybrid_linear, is_mla, resolve_auto_ssm_cache_dtype, +from .config_utils import (has_fp4_kv_cache, is_hybrid_linear, is_mla, + resolve_auto_ssm_cache_dtype, supports_fp4_mla_attention, uses_fp4_mla_attention, validate_kimi_kda_state_dtype) @@ -212,7 +213,7 @@ def validate_fp4_mla_config(model_config: ModelConfig, llm_args: TorchLlmArgs) -> None: """Validate FP4 MLA before model construction and KV-cache allocation.""" if not (is_mla(model_config.pretrained_config) - and model_config.quant_config.quant_mode.has_fp4_kv_cache()): + and has_fp4_kv_cache(model_config.quant_config)): return if not supports_fp4_mla_attention(model_config): raise ValueError( diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index c63cd14d93ba..62126e686782 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -48,7 +48,7 @@ nvtx_range) from ...logger import logger from ...mapping import Mapping -from .config_utils import uses_vswa_kv_cache_layout +from .config_utils import has_fp4_kv_cache, uses_vswa_kv_cache_layout from .connectors.kv_cache_connector import KvCacheConnectorManager from .llm_request import LlmRequest, LlmRequestState, get_draft_token_length from .scheduler import ScheduledRequests @@ -1581,8 +1581,7 @@ def get_cache_size_per_token(model_config: ModelConfigPython, mla = hasattr(config, "kv_lora_rank") and config.kv_lora_rank is not None quant_config = model_config.quant_config - if (mla and quant_config is not None - and quant_config.quant_mode.has_fp4_kv_cache()): + if mla and has_fp4_kv_cache(quant_config): raise ValueError( "FP4 MLA cache sizing requires Fp4MlaKVCacheManagerV2; " "KVCacheManager V1 is not supported.") diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0ce6cde395b2..3d2f1704e1be 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -308,6 +308,7 @@ def test_nvfp4_mla_gsm8k(self, monkeypatch): from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttention from tensorrt_llm._torch.attention.mla import MLA + from tensorrt_llm._torch.pyexecutor.config_utils import has_fp4_kv_cache # Inspect the actual worker's quantization, cache pools and FMHA choice. monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") @@ -326,7 +327,7 @@ def test_nvfp4_mla_gsm8k(self, monkeypatch): executor = llm._executor.engine quant_config = executor.model_engine.model.model_config.quant_config assert quant_config.quant_algo == QuantAlgo.NVFP4 - assert quant_config.quant_mode.has_fp4_kv_cache() + assert has_fp4_kv_cache(quant_config) manager = executor.kv_cache_manager assert isinstance(manager, Fp4MlaKVCacheManagerV2) assert manager.get_mla_v_scale_pool().numel() > 0 diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index f4c9b06c5d84..521ee799c8bf 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -15,9 +15,11 @@ validate_encoder_decoder_kv_cache_config, validate_fp4_mla_config, ) +from tensorrt_llm._utils import QuantModeWrapper from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig +from tensorrt_llm.quantization.mode import QuantMode pytestmark = pytest.mark.cpu_only @@ -152,6 +154,7 @@ def _mock_sm107(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.parametrize("kv_cache_dtype", ["auto", "nvfp4"]) @pytest.mark.parametrize("chunked_prefill", [False, True]) +@pytest.mark.parametrize("quant_mode_representation", ["wrapper", "enum", "int"]) @pytest.mark.parametrize( "mla,backend,sparse,hybrid,expected", [ @@ -166,6 +169,7 @@ def test_validate_and_set_kv_cache_quant_resolves_nvfp4_on_sm107( monkeypatch: pytest.MonkeyPatch, kv_cache_dtype: str, chunked_prefill: bool, + quant_mode_representation: str, mla: bool, backend: str, sparse: bool, @@ -187,6 +191,12 @@ def test_validate_and_set_kv_cache_quant_resolves_nvfp4_on_sm107( validate_and_set_kv_cache_quant(model_config, kv_cache_dtype) assert model_config.quant_config.kv_cache_quant_algo == expected assert model_config.quant_config_dict["layer"].kv_cache_quant_algo == expected + quant_mode = QuantMode.from_quant_algo(kv_cache_quant_algo=expected) + if quant_mode_representation == "wrapper": + quant_mode = QuantModeWrapper([quant_mode]) + elif quant_mode_representation == "int": + quant_mode = int(quant_mode) + monkeypatch.setattr(model_config.quant_config, "quant_mode", quant_mode) assert uses_fp4_mla_attention(model_config) == (expected == QuantAlgo.NVFP4) llm_args = types.SimpleNamespace( enable_chunked_prefill=chunked_prefill, From cc982432caa1d853071c37c78848e2c433988c7b Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Thu, 10 Sep 2026 20:25:46 -0700 Subject: [PATCH 13/21] fix: preserve quant mode wrappers during config copying Do not forward missing Python protocol hooks through QuantModeWrapper query aggregation. This prevents an eagerly cached quant_mode from becoming an integer during deepcopy. Remove the integer-mode compatibility workaround and restore test_model_config.py to the PR base. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 15 +++-- .../_torch/pyexecutor/config_utils.py | 15 +---- .../_torch/pyexecutor/model_loader.py | 5 +- .../_torch/pyexecutor/resource_manager.py | 5 +- tensorrt_llm/_utils.py | 4 ++ .../defs/accuracy/test_llm_api_pytorch.py | 3 +- tests/unittest/_torch/test_model_config.py | 64 ++----------------- 7 files changed, 24 insertions(+), 87 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index eae939502fda..2cbe479c6254 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -55,11 +55,10 @@ from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, extract_mamba_kv_cache_params, extract_qwen4_exp_ple_cache_params, - get_layer_attention_window, has_fp4_kv_cache, - is_gemma4_hybrid, is_hybrid_linear, is_kimi_linear, - is_mla, is_nemotron_hybrid, is_qwen3_hybrid, - is_qwen4_exp, resolve_vocab_size, - uses_vswa_kv_cache_layout) + get_layer_attention_window, is_gemma4_hybrid, + is_hybrid_linear, is_kimi_linear, is_mla, + is_nemotron_hybrid, is_qwen3_hybrid, is_qwen4_exp, + resolve_vocab_size, uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder @@ -190,7 +189,8 @@ def get_kv_cache_manager_cls( sparse_attn_config = model_config.sparse_attention_config sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) quant_config = getattr(model_config, "quant_config", None) - if is_mla(config) and has_fp4_kv_cache(quant_config): + if (is_mla(config) and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache()): if kv_cache_config.use_kv_cache_manager_v2 is False: raise ValueError("FP4 MLA requires use_kv_cache_manager_v2=True.") if model_config.attn_backend != "TRTLLM": @@ -913,7 +913,8 @@ def _validate_or_fallback_kv_cache_manager_v2( f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") quant_config = getattr(model_config, "quant_config", None) - if is_mla(config) and has_fp4_kv_cache(quant_config): + if (is_mla(config) and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache()): raise NotImplementedError( "FP4 MLA requires Fp4MlaKVCacheManagerV2, which is " f"not yet supported with {incompat_str}. Disable these " diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 9ae4beed1cc5..268cf0f08c07 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -11,11 +11,9 @@ from tensorrt_llm._utils import str_dtype_to_torch from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig from tensorrt_llm.logger import logger -from tensorrt_llm.quantization.mode import QuantMode if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm.models.modeling_utils import QuantConfig def resolve_cache_transceiver_config( @@ -380,20 +378,11 @@ def supports_fp4_mla_attention(model_config: "ModelConfig") -> bool: and not is_hybrid_linear(model_config.pretrained_config)) -def has_fp4_kv_cache(quant_config: Optional["QuantConfig"]) -> bool: - """Check resolved KV quantization for integer modes and Python wrappers.""" - if quant_config is None: - return False - quant_mode = quant_config.quant_mode - if isinstance(quant_mode, int): - quant_mode = QuantMode(quant_mode) - return quant_mode.has_fp4_kv_cache() - - def uses_fp4_mla_attention(model_config: "ModelConfig") -> bool: """Use the resolved quantization, never the requested cache dtype.""" quant_config = getattr(model_config, "quant_config", None) - return (has_fp4_kv_cache(quant_config) + return (quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache() and supports_fp4_mla_attention(model_config)) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index f9da30f170dc..295e8f30d89b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -54,8 +54,7 @@ maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope -from .config_utils import (has_fp4_kv_cache, is_hybrid_linear, is_mla, - resolve_auto_ssm_cache_dtype, +from .config_utils import (is_hybrid_linear, is_mla, resolve_auto_ssm_cache_dtype, supports_fp4_mla_attention, uses_fp4_mla_attention, validate_kimi_kda_state_dtype) @@ -213,7 +212,7 @@ def validate_fp4_mla_config(model_config: ModelConfig, llm_args: TorchLlmArgs) -> None: """Validate FP4 MLA before model construction and KV-cache allocation.""" if not (is_mla(model_config.pretrained_config) - and has_fp4_kv_cache(model_config.quant_config)): + and model_config.quant_config.quant_mode.has_fp4_kv_cache()): return if not supports_fp4_mla_attention(model_config): raise ValueError( diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 62126e686782..c63cd14d93ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -48,7 +48,7 @@ nvtx_range) from ...logger import logger from ...mapping import Mapping -from .config_utils import has_fp4_kv_cache, uses_vswa_kv_cache_layout +from .config_utils import uses_vswa_kv_cache_layout from .connectors.kv_cache_connector import KvCacheConnectorManager from .llm_request import LlmRequest, LlmRequestState, get_draft_token_length from .scheduler import ScheduledRequests @@ -1581,7 +1581,8 @@ def get_cache_size_per_token(model_config: ModelConfigPython, mla = hasattr(config, "kv_lora_rank") and config.kv_lora_rank is not None quant_config = model_config.quant_config - if mla and has_fp4_kv_cache(quant_config): + if (mla and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache()): raise ValueError( "FP4 MLA cache sizing requires Fp4MlaKVCacheManagerV2; " "KVCacheManager V1 is not supported.") diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 372f4e1ee70e..bc835129c6f3 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -839,6 +839,10 @@ def __init__(self, objs): self.objs = objs def __getattr__(self, name): + # Missing Python protocol hooks must not be forwarded as quantization + # queries: reducing __deepcopy__ results turns the wrapper into an int. + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) def method_wrapper(*args, **kwargs): result = False diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 3d2f1704e1be..0ce6cde395b2 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -308,7 +308,6 @@ def test_nvfp4_mla_gsm8k(self, monkeypatch): from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttention from tensorrt_llm._torch.attention.mla import MLA - from tensorrt_llm._torch.pyexecutor.config_utils import has_fp4_kv_cache # Inspect the actual worker's quantization, cache pools and FMHA choice. monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") @@ -327,7 +326,7 @@ def test_nvfp4_mla_gsm8k(self, monkeypatch): executor = llm._executor.engine quant_config = executor.model_engine.model.model_config.quant_config assert quant_config.quant_algo == QuantAlgo.NVFP4 - assert has_fp4_kv_cache(quant_config) + assert quant_config.quant_mode.has_fp4_kv_cache() manager = executor.kv_cache_manager assert isinstance(manager, Fp4MlaKVCacheManagerV2) assert manager.get_mla_v_scale_pool().numel() > 0 diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py index 521ee799c8bf..19bfe98eebb8 100644 --- a/tests/unittest/_torch/test_model_config.py +++ b/tests/unittest/_torch/test_model_config.py @@ -1,6 +1,3 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - import json import struct import types @@ -9,17 +6,12 @@ import torch from tensorrt_llm._torch.model_config import _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT, ModelConfig -from tensorrt_llm._torch.pyexecutor.config_utils import uses_fp4_mla_attention from tensorrt_llm._torch.pyexecutor.model_loader import ( validate_and_set_kv_cache_quant, validate_encoder_decoder_kv_cache_config, - validate_fp4_mla_config, ) -from tensorrt_llm._utils import QuantModeWrapper -from tensorrt_llm.llmapi.llm_args import DeepSeekSparseAttentionConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig -from tensorrt_llm.quantization.mode import QuantMode pytestmark = pytest.mark.cpu_only @@ -152,61 +144,13 @@ def _mock_sm107(monkeypatch: pytest.MonkeyPatch) -> None: ) -@pytest.mark.parametrize("kv_cache_dtype", ["auto", "nvfp4"]) -@pytest.mark.parametrize("chunked_prefill", [False, True]) -@pytest.mark.parametrize("quant_mode_representation", ["wrapper", "enum", "int"]) -@pytest.mark.parametrize( - "mla,backend,sparse,hybrid,expected", - [ - (False, "TRTLLM", False, False, QuantAlgo.FP8), - (True, "TRTLLM", False, False, QuantAlgo.NVFP4), - (True, "FLASHINFER", False, False, QuantAlgo.FP8), - (True, "TRTLLM", True, False, QuantAlgo.FP8), - (True, "TRTLLM", False, True, QuantAlgo.FP8), - ], -) -def test_validate_and_set_kv_cache_quant_resolves_nvfp4_on_sm107( +def test_validate_and_set_kv_cache_quant_downgrades_nvfp4_on_sm107( monkeypatch: pytest.MonkeyPatch, - kv_cache_dtype: str, - chunked_prefill: bool, - quant_mode_representation: str, - mla: bool, - backend: str, - sparse: bool, - hybrid: bool, - expected: QuantAlgo, ) -> None: _mock_sm107(monkeypatch) - model_config = ModelConfig( - pretrained_config=types.SimpleNamespace( - kv_lora_rank=512 if mla else None, - qk_rope_head_dim=64 if mla else None, - hybrid_override_pattern="M*" if hybrid else None, - ), - attn_backend=backend, - sparse_attention_config=DeepSeekSparseAttentionConfig() if sparse else None, - quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4), - quant_config_dict={"layer": QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4)}, - ) - validate_and_set_kv_cache_quant(model_config, kv_cache_dtype) - assert model_config.quant_config.kv_cache_quant_algo == expected - assert model_config.quant_config_dict["layer"].kv_cache_quant_algo == expected - quant_mode = QuantMode.from_quant_algo(kv_cache_quant_algo=expected) - if quant_mode_representation == "wrapper": - quant_mode = QuantModeWrapper([quant_mode]) - elif quant_mode_representation == "int": - quant_mode = int(quant_mode) - monkeypatch.setattr(model_config.quant_config, "quant_mode", quant_mode) - assert uses_fp4_mla_attention(model_config) == (expected == QuantAlgo.NVFP4) - llm_args = types.SimpleNamespace( - enable_chunked_prefill=chunked_prefill, - kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2=True), - ) - if expected == QuantAlgo.NVFP4 and chunked_prefill: - with pytest.raises(ValueError, match="does not support chunked prefill"): - validate_fp4_mla_config(model_config, llm_args) - else: - validate_fp4_mla_config(model_config, llm_args) + model_config = _make_model_config_with_kv_quant(QuantAlgo.FP8) + validate_and_set_kv_cache_quant(model_config, "nvfp4") + assert model_config.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 def _make_mixed_precision_model_config(): From 941216ba9b8c38d66785fb04c3ff3128f1a7d1ed Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 11 Sep 2026 01:18:49 -0700 Subject: [PATCH 14/21] test: align MLA cache reuse mocks with model config Use a real QuantConfig and expose the backend, MLA dimensions, and sparse configuration required by FP4 MLA routing. Preserve the existing test cases and map the unquantized sentinel to the native None value. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- ...est_py_executor_creator_mla_cache_reuse_sync.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index 893a4673d7cd..778562a57da8 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -24,6 +24,7 @@ ) from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, ContextChunkingPolicy +from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantAlgo pytestmark = pytest.mark.cpu_only @@ -126,6 +127,7 @@ def __init__( *, attn_runtime_features, kv_cache_quant_algo, + attn_backend="TRTLLM", sparse_algorithm=None, enable_flash_mla=False, max_seq_len=128, @@ -135,6 +137,7 @@ def __init__( Args: attn_runtime_features: AttentionRuntimeFeatures instance. kv_cache_quant_algo: Quantization algorithm for KV cache. + attn_backend: Attention backend selected for the model. sparse_algorithm: Optional sparse-attention algorithm name. enable_flash_mla: Whether to emulate the FlashMLA block-size override. max_seq_len: Effective sequence length reported by the model engine. @@ -148,10 +151,16 @@ def __init__( self.attn_metadata = None self.model = SimpleNamespace( model_config=SimpleNamespace( + attn_backend=attn_backend, + sparse_attention_config=self.sparse_attention_config, enable_flash_mla=enable_flash_mla, is_generation=True, - pretrained_config=SimpleNamespace(), - quant_config=SimpleNamespace(kv_cache_quant_algo=kv_cache_quant_algo), + pretrained_config=SimpleNamespace(kv_lora_rank=512, qk_rope_head_dim=64), + quant_config=QuantConfig( + kv_cache_quant_algo=( + None if kv_cache_quant_algo == QuantAlgo.NO_QUANT else kv_cache_quant_algo + ) + ), ), vocab_size_padded=32000, ) @@ -305,6 +314,7 @@ def _create_model_engine(**kwargs): return _DummyModelEngine( attn_runtime_features=kwargs["attn_runtime_features"], kv_cache_quant_algo=kv_cache_quant_algo, + attn_backend=attn_backend, sparse_algorithm=sparse_algorithm, enable_flash_mla=enable_flash_mla, max_seq_len=model_max_seq_len, From f756d4c5afed60bf07bcb0ed380f91c905dc66d5 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Sun, 13 Sep 2026 22:15:43 -0700 Subject: [PATCH 15/21] refactor: address FP4 MLA capability and metadata review Centralize FP4 MLA FMHA capability checks and use the attention backend capability instead of concrete type checks. Cache invariant validation while preserving live-input checks on cache hits. Move batch-shared tensors and lifecycle bookkeeping into Fp4MlaState without changing metadata allocation geometry, and keep the layer-local FP8 attention view on the FMHA instance. Rename the accuracy test and update existing test helpers and documentation without adding test cases. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 14 +- .../attention/backends/fmha/fallback.py | 2 - .../_torch/attention/backends/fmha/fp4_mla.py | 112 ++++--- .../attention/backends/fmha/interface.py | 7 + .../attention/backends/fp4_mla/__init__.py | 236 +++++++------ .../backends/fp4_mla/fp4_mla_context.py | 9 +- .../attention/backends/fp4_mla/state.py | 206 ++++++++++++ .../_torch/attention/backends/interface.py | 5 + .../_torch/attention/backends/trtllm.py | 317 ++---------------- tensorrt_llm/_torch/attention/mla.py | 12 +- tensorrt_llm/_torch/speculative/mtp.py | 3 +- .../defs/accuracy/test_llm_api_pytorch.py | 2 +- .../test_lists/qa/llm_function_core.txt | 2 +- .../unittest/_torch/attention/test_fp4_mla.py | 138 ++++---- 14 files changed, 530 insertions(+), 535 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/state.py diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 461567fd7b85..cabf6be2101f 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -258,6 +258,7 @@ The core contract is: - `support_fused_rope()` - `support_fused_qkv()` - `support_mla()` + - `support_fp4_kv_cache()` - `runtime_workspace_bytes_per_token(model_config, mapping)` — the memory-accounting contract (default `0`); see below - `runtime_workspace_is_chunked_prefill_bounded(model_config)` — whether @@ -399,8 +400,8 @@ independently with `is_supported(..., phase=...)`; a phased library accepts only phases backed by its corresponding `run_*()` entry point. `Fmha` owns both entry points. Libraries declare shared capabilities through -class attributes, such as `supports_skip_correction` and -`supports_block_sparse_inputs`, and override only +class attributes, such as `supports_skip_correction`, `supports_block_sparse_inputs`, +and `supports_fp4_mla`, and override only `_is_available()` and `_is_supported()` for implementation-specific checks. `is_available()` rejects unsupported static capabilities before calling `_is_available()`. `is_supported()` provides the same boundary for shared @@ -428,6 +429,10 @@ The FMHA package is split by role: - `fmha/combined.py` composes different context and generation implementations for non-MLA mixed batches. - `fmha/fp4_mla.py` implements FP4 MLA context and no-dequant decode. + The shared FMHA availability check admits FP4 MLA only to libraries that + declare `supports_fp4_mla`; it does not restrict the existing FP4 GQA path. + Selection caches model-invariant and cache-key-covered validation. Dynamic + sparse/sinks inputs and prepared-state checks still run on cache hits. The core implementation requires dense TRTLLM MLA, BF16 absorption weights, fused RoPE with duplicated rotary tables, and KV Cache Manager V2. It uses FP8 context attention with an FP4 cache update and FP4 generation attention. @@ -437,6 +442,11 @@ The FMHA package is split by role: Disaggregated serving is reserved for the follow-up integration. On SM107, this dense TRTLLM path keeps NVFP4 KV quantization; unsupported profiles retain the existing FP8 fallback. + `fp4_mla/state.py` owns batch-shared page tables, HP/V-scale views, append + metadata, and scratch caches. `TrtllmAttentionMetadata` keeps one optional + state reference and forwards prepare/MTP lifecycle updates to it. These + buffers continue to use the metadata allocator for CUDA-graph address + stability; the layer-local FP8 attention view is allocated lazily by FMHA. - `fmha/triton_custom_mask.py` implements the Triton custom-mask context phase. Custom-mask data applies to context requests; for mixed batches, `TrtllmAttention` can pair it with a later causal-generation provider through diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 2f91b685fe6d..5631950e9e06 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -62,8 +62,6 @@ class FallbackFmha(Fmha): @classmethod def _is_available(cls, attn: "TrtllmAttention") -> bool: - if attn.is_mla_enable and attn.has_fp4_kv_cache: - return False sparse_algorithm = getattr(attn.sparse_params, "algorithm", None) if sparse_algorithm in ("deepseek_v4", "dsa"): if getattr(attn, "kv_cache_dtype", None) == "fp8_ds_mla": diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index 7c436abcd92d..1774b591d35a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -23,7 +23,6 @@ scatter_fp4_mla_kv_cache, ) from tensorrt_llm._torch.attention.backends.fp4_mla.fp4_mla_context import ( - _FP8_CONTEXT_ATTN_ATTR, _FP8_CONTEXT_SCRATCH_ATTR, _build_fp8_mla_context_attn, _execute_fp8_context_with_cache_update, @@ -38,6 +37,7 @@ ) from tensorrt_llm.bindings import DataType +from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha if TYPE_CHECKING: @@ -50,52 +50,37 @@ class Fp4MlaFmha(PhasedFmha): """TRTLLM FMHA library for FP4 MLA context and no-dequant decode.""" + supports_fp4_mla = True + + def __init__(self, attn: "TrtllmAttention") -> None: + super().__init__(attn) + self._fp8_attention: Optional["TrtllmAttention"] = None + @classmethod def _is_available(cls, attn: "TrtllmAttention") -> bool: return attn.is_mla_enable and attn.has_fp4_kv_cache - def forward( + def _is_supported( self, q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, - ) -> None: - self._validate_request(k, v, metadata, forward_args) - super().forward(q, k, v, metadata, forward_args) - - def _validate_request( - self, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata: "TrtllmAttentionMetadata", - forward_args: AttentionForwardArgs, - ) -> None: + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + # Masks/output formats are in the selection cache key; cache geometry + # and sparse compression are model invariants. Live inputs stay in forward. + del q, k, v, phase if forward_args.output_sf is not None: raise NotImplementedError("FP4 MLA does not support quantized attention output.") if forward_args.attention_mask != PredefinedAttentionMask.CAUSAL: raise NotImplementedError("FP4 MLA requires a causal attention mask.") if forward_args.attention_mask_data is not None: raise NotImplementedError("FP4 MLA does not support custom attention masks.") - if forward_args.attention_sinks is not None: - raise NotImplementedError("FP4 MLA does not support attention sinks.") - - sparse_prediction = forward_args.sparse_prediction sparse_params = self.attn.sparse_params - uses_spcompress = getattr(sparse_params, "uses_spcompress", None) - if ( - ( - sparse_prediction.sparse_kv_indices is not None - and sparse_prediction.sparse_kv_indices.numel() > 0 - ) - or ( - sparse_prediction.sparse_attn_indices is not None - and sparse_prediction.sparse_attn_indices.numel() > 0 - ) - or metadata.num_sparse_topk > 0 - or uses_spcompress - ): + if getattr(sparse_params, "uses_spcompress", False): raise NotImplementedError("FP4 MLA does not support sparse attention.") kv_cache_manager = metadata.kv_cache_manager @@ -105,23 +90,52 @@ def _validate_request( raise RuntimeError("FP4 MLA requires NVFP4 KV cache storage.") if kv_cache_manager.kv_factor != 1: raise RuntimeError("FP4 MLA requires a SELF-K-only KV cache.") - if metadata.high_precision_kv_pool is None: - raise RuntimeError("FP4 MLA requires the high-precision KV pool.") - if metadata.fp4_mla_v_scale_pool is None: - raise RuntimeError("FP4 MLA requires the V-scale pool.") if metadata.beam_width != 1: raise NotImplementedError("FP4 MLA does not support beam search.") + return True + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + # These inputs/readiness conditions are not part of the FMHA cache + # key. Keep checking them even when selection reuses a cached library. + if forward_args.attention_sinks is not None: + raise NotImplementedError("FP4 MLA does not support attention sinks.") + sparse_inputs = forward_args.sparse_runtime_params + if ( + metadata.num_sparse_topk > 0 + or sparse_inputs.block_sparse_inputs is not None + or ( + sparse_inputs.sparse_kv_indices is not None + and sparse_inputs.sparse_kv_indices.numel() > 0 + ) + or ( + sparse_inputs.sparse_attn_indices is not None + and sparse_inputs.sparse_attn_indices.numel() > 0 + ) + ): + raise NotImplementedError("FP4 MLA does not support sparse attention.") + state = metadata.fp4_mla_state + if state is None or state.hp_pool is None: + raise RuntimeError("FP4 MLA requires prepared high-precision KV state.") + if state.v_scale_pool is None: + raise RuntimeError("FP4 MLA requires the V-scale pool.") attention_input_type = forward_args.attention_input_type if attention_input_type == AttentionInputType.context_only: if k is None or v is None: raise RuntimeError("FP4 MLA context requires expanded K and V tensors.") - return - if attention_input_type == AttentionInputType.generation_only: + elif attention_input_type == AttentionInputType.generation_only: if k is not None or v is not None: raise RuntimeError("FP4 MLA generation expects a fused query input.") - return - raise NotImplementedError("FP4 MLA requires a context-only or generation-only call.") + else: + raise NotImplementedError("FP4 MLA requires a context-only or generation-only call.") + super().forward(q, k, v, metadata, forward_args) def run_mla_context(self, params: FmhaParams) -> None: attn = params.attn @@ -137,7 +151,7 @@ def run_mla_context(self, params: FmhaParams) -> None: raise RuntimeError("FP4 MLA context requires an output buffer.") if forward_args.latent_cache is None: raise RuntimeError("FP4 MLA context requires latent_cache.") - if metadata.positions is None: + if metadata.fp4_mla_state.positions is None: raise RuntimeError("FP4 MLA context requires token positions.") if metadata.num_contexts <= 0: raise RuntimeError("FP4 MLA context requires context requests.") @@ -194,10 +208,10 @@ def update_fp4_cache() -> None: ) setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch) - fp8_attention = getattr(attn, _FP8_CONTEXT_ATTN_ATTR, None) + fp8_attention = self._fp8_attention if fp8_attention is None: fp8_attention = _build_fp8_mla_context_attn(attn) - setattr(attn, _FP8_CONTEXT_ATTN_ATTR, fp8_attention) + self._fp8_attention = fp8_attention fp8_attention.rotary_inv_freq = attn.rotary_inv_freq fp8_attention.rotary_cos_sin = attn.rotary_cos_sin fp8_metadata = _get_fp8_mla_context_metadata(metadata, scratch) @@ -240,18 +254,18 @@ def run_mla_generation(self, params: FmhaParams) -> None: qk_rope_head_dim = attn.qk_rope_head_dim or 0 fused_head_dim = kv_lora_rank + qk_rope_head_dim - if not bool(getattr(metadata, "_fp4_mla_generation_cache_scattered", False)): + if not bool(getattr(metadata.fp4_mla_state, "generation_cache_scattered", False)): raise RuntimeError( "FP4 MLA generation requires fused RoPE/Q quantization and " "cache/HP-pool update before attention." ) - metadata._fp4_mla_generation_cache_scattered = False + metadata.fp4_mla_state.generation_cache_scattered = False query = q.view(q.shape[0], attn.num_heads, fused_head_dim) output_view = output.view(q.shape[0], attn.num_heads, kv_lora_rank) sm_scale = 1.0 / (attn.q_scaling * ((attn.qk_nope_head_dim or 0) + qk_rope_head_dim) ** 0.5) - prequantized_q = getattr(metadata, "_fp4_mla_prequantized_q", None) - prequantized_q_sf = getattr(metadata, "_fp4_mla_prequantized_q_sf", None) - q_batch_capacity = getattr(metadata, "_fp4_mla_q_batch_capacity", None) + prequantized_q = getattr(metadata.fp4_mla_state, "prequantized_q", None) + prequantized_q_sf = getattr(metadata.fp4_mla_state, "prequantized_q_sf", None) + q_batch_capacity = getattr(metadata.fp4_mla_state, "q_batch_capacity", None) try: run_fp4_mla_attention_decode( metadata, @@ -268,9 +282,9 @@ def run_mla_generation(self, params: FmhaParams) -> None: softmax_stats_tensor=params.fwd.softmax_stats_tensor, ) finally: - metadata._fp4_mla_prequantized_q = None - metadata._fp4_mla_prequantized_q_sf = None - metadata._fp4_mla_q_batch_capacity = None + metadata.fp4_mla_state.prequantized_q = None + metadata.fp4_mla_state.prequantized_q_sf = None + metadata.fp4_mla_state.q_batch_capacity = None __all__ = ["Fp4MlaFmha"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index b8b1ce4cad07..efac2712c932 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/interface.py @@ -43,6 +43,7 @@ class Fmha(ABC): supports_skip_correction: ClassVar[bool] = False supports_block_sparse_inputs: ClassVar[bool] = False supports_workspace_reclamation: bool = False + supports_fp4_mla: ClassVar[bool] = False def __init__(self, attn: "TrtllmAttention"): self._attn_ref: weakref.ReferenceType["TrtllmAttention"] = weakref.ref(attn) @@ -76,6 +77,12 @@ def is_available(cls, attn: "TrtllmAttention") -> bool: f"{cls.__name__} is unavailable: skip-correction is enabled and unsupported." ) return False + if ( + getattr(attn, "is_mla_enable", False) + and getattr(attn, "has_fp4_kv_cache", False) + and not cls.supports_fp4_mla + ): + return False return cls._is_available(attn) @classmethod diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py index 9924a4a7e832..7c81954f85c1 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -572,10 +572,10 @@ def configure_fp4_mla_device_page_table( table on the GPU. The materialization kernel decodes V2 page indices and refreshes rows from the final device KV lengths before cache update. """ - metadata._fp4_mla_device_page_table = False - metadata._fp4_mla_device_page_table_valid = False - metadata.fp4_mla_page_table_stride = 0 - metadata.fp4_mla_context_repack_max_touched_pages = 1 + metadata.fp4_mla_state.device_page_table = False + metadata.fp4_mla_state.device_page_table_valid = False + metadata.fp4_mla_state.page_table_stride = 0 + metadata.fp4_mla_state.context_repack_max_touched_pages = 1 kv_cache_manager = getattr(metadata, "kv_cache_manager", None) num_contexts = int(getattr(metadata, "num_contexts", 0)) @@ -585,9 +585,9 @@ def configure_fp4_mla_device_page_table( num_context_tokens = int(getattr(metadata, "num_ctx_tokens", 0)) num_generation_tokens = num_tokens - num_context_tokens block_offsets = getattr(metadata, "kv_cache_block_offsets", None) - page_ids = getattr(metadata, "_paged_kv_indices", None) - paged_kv_indptr = getattr(metadata, "_paged_kv_indptr", None) - paged_kv_indptr_decode = getattr(metadata, "paged_kv_indptr_decode", None) + page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) + paged_kv_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) + paged_kv_indptr_decode = getattr(metadata.fp4_mla_state, "paged_kv_indptr_decode", None) max_page_capacity = int(getattr(kv_cache_manager, "max_blocks_per_seq", 0) or 0) page_spec = _fp4_mla_page_table_spec(kv_cache_manager) page_index_scale = int(page_spec.cache_page_index_scale) @@ -631,7 +631,7 @@ def configure_fp4_mla_device_page_table( if not supported: return False - hp_page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) + hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) max_pool_id = max(page_spec.cache_pool_id, page_spec.hp_pool_id) if ( not isinstance(hp_page_ids, torch.Tensor) @@ -640,10 +640,10 @@ def configure_fp4_mla_device_page_table( or block_offsets.shape[0] <= max_pool_id ): return False - metadata._fp4_mla_cache_pool_id = int(page_spec.cache_pool_id) - metadata._fp4_mla_cache_page_index_scale = int(page_spec.cache_page_index_scale) - metadata._fp4_mla_hp_pool_id = int(page_spec.hp_pool_id) - metadata._fp4_mla_hp_page_index_scale = int(page_spec.hp_page_index_scale) + metadata.fp4_mla_state.cache_pool_id = int(page_spec.cache_pool_id) + metadata.fp4_mla_state.cache_page_index_scale = int(page_spec.cache_page_index_scale) + metadata.fp4_mla_state.hp_pool_id = int(page_spec.hp_pool_id) + metadata.fp4_mla_state.hp_page_index_scale = int(page_spec.hp_page_index_scale) max_pages = max_page_capacity host_kv_lens_available = ( @@ -675,7 +675,7 @@ def configure_fp4_mla_device_page_table( max_context_len, FP4_MLA_TOKENS_PER_BLOCK, ) - metadata.fp4_mla_context_repack_max_touched_pages = min( + metadata.fp4_mla_state.context_repack_max_touched_pages = min( max_pages, triton.next_power_of_2(max(1, max_context_pages)), ) @@ -700,14 +700,15 @@ def configure_fp4_mla_device_page_table( ) if not buffers_cover_table: return False - if metadata._fp4_mla_hp_page_indices.numel() < required_page_ids: + if metadata.fp4_mla_state.hp_page_indices.numel() < required_page_ids: return False - metadata._fp4_mla_device_page_table = True - metadata.fp4_mla_page_table_stride = max_pages - metadata.num_blocks = None - metadata.num_context_blocks = num_contexts * max_pages - metadata.num_generation_blocks = num_generation_sequences * max_pages + metadata.fp4_mla_state.device_page_table = True + metadata.fp4_mla_state.num_sequences = num_sequences + metadata.fp4_mla_state.page_table_stride = max_pages + metadata.fp4_mla_state.num_blocks = None + metadata.fp4_mla_state.num_context_blocks = num_contexts * max_pages + metadata.fp4_mla_state.num_generation_blocks = num_generation_sequences * max_pages return True @@ -717,15 +718,15 @@ def materialize_fp4_mla_device_page_table( generation_kv_lens: Optional[torch.Tensor] = None, ) -> None: """Refresh the fixed-stride context and generation page table once per forward.""" - if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): + if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): raise RuntimeError("FP4 MLA requires fixed-stride device page metadata.") - if bool(getattr(metadata, "_fp4_mla_device_page_table_valid", False)): + if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): return num_contexts = int(metadata.num_contexts) num_sequences = int(metadata.num_seqs) num_generation_sequences = num_sequences - num_contexts - max_pages = int(metadata.fp4_mla_page_table_stride) + max_pages = int(metadata.fp4_mla_state.page_table_stride) if num_sequences <= 0 or max_pages <= 0: raise RuntimeError( "FP4 MLA device page metadata requires positive sequence and page capacities." @@ -757,15 +758,15 @@ def materialize_fp4_mla_device_page_table( f"{num_generation_sequences} entries." ) - cache_pool_id = int(getattr(metadata, "_fp4_mla_cache_pool_id", 0)) + cache_pool_id = int(getattr(metadata.fp4_mla_state, "cache_pool_id", 0)) block_offsets = metadata.kv_cache_block_offsets[ cache_pool_id, :num_sequences, 0, :max_pages, ] - page_ids = metadata._paged_kv_indices[: num_sequences * max_pages] - page_index_scale = int(metadata._fp4_mla_cache_page_index_scale) + page_ids = metadata.fp4_mla_state._paged_kv_indices[: num_sequences * max_pages] + page_index_scale = int(metadata.fp4_mla_state.cache_page_index_scale) if page_index_scale <= 0: raise RuntimeError("FP4 MLA device page metadata requires a positive page-index scale.") grid = ( @@ -774,8 +775,8 @@ def materialize_fp4_mla_device_page_table( ) _fp4_mla_materialize_page_table_kernel[grid]( page_ids, - metadata._paged_kv_indptr, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state._paged_kv_indptr, + metadata.fp4_mla_state.paged_kv_indptr_decode, block_offsets, kv_lens, generation_kv_lens, @@ -788,11 +789,11 @@ def materialize_fp4_mla_device_page_table( PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, num_warps=4, ) - hp_page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) + hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) if not isinstance(hp_page_ids, torch.Tensor): raise RuntimeError("FP4 MLA requires an HP page-table output tensor.") - hp_pool_id = int(metadata._fp4_mla_hp_pool_id) - hp_page_index_scale = int(metadata._fp4_mla_hp_page_index_scale) + hp_pool_id = int(metadata.fp4_mla_state.hp_pool_id) + hp_page_index_scale = int(metadata.fp4_mla_state.hp_page_index_scale) hp_block_offsets = metadata.kv_cache_block_offsets[ hp_pool_id, :num_sequences, @@ -801,8 +802,8 @@ def materialize_fp4_mla_device_page_table( ] _fp4_mla_materialize_page_table_kernel[grid]( hp_page_ids, - metadata._paged_kv_indptr, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state._paged_kv_indptr, + metadata.fp4_mla_state.paged_kv_indptr_decode, hp_block_offsets, kv_lens, generation_kv_lens, @@ -815,7 +816,7 @@ def materialize_fp4_mla_device_page_table( PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, num_warps=4, ) - metadata._fp4_mla_device_page_table_valid = True + metadata.fp4_mla_state.device_page_table_valid = True @triton.jit @@ -1126,7 +1127,7 @@ def get_fp4_mla_v_scale_pool_view( v_head_dim: int, ) -> torch.Tensor: """View the auxiliary MLA V-scale pool in Triton's block-scaled layout.""" - pool = getattr(metadata, "fp4_mla_v_scale_pool", None) + pool = getattr(metadata.fp4_mla_state, "v_scale_pool", None) if pool is None: raise RuntimeError("FP4 MLA V scale pool is not allocated.") @@ -1157,7 +1158,7 @@ def get_fp4_mla_v_scale_pool_view( def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata, "_fp4_mla_kv_global_scale", None) + global_scale = getattr(metadata.fp4_mla_state, "kv_global_scale", None) if ( not isinstance(global_scale, torch.Tensor) or global_scale.device != device @@ -1169,7 +1170,7 @@ def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tens def _get_fp4_mla_q_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata, "_fp4_mla_q_global_scale", None) + global_scale = getattr(metadata.fp4_mla_state, "q_global_scale", None) if ( not isinstance(global_scale, torch.Tensor) or global_scale.device != device @@ -1287,7 +1288,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( ) rotary_cos_sin_ptr = rotary_cos_sin if rotary_cos_sin is not None else latent_cache - hp_pool = getattr(metadata, "high_precision_kv_pool", None) + hp_pool = getattr(metadata.fp4_mla_state, "hp_pool", None) if not isinstance(hp_pool, torch.Tensor): raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") if hp_pool.device != latent_cache.device: @@ -1306,7 +1307,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( ) if hp_pool.stride(-1) != 1: raise ValueError("FP4 MLA high-precision KV pool must be contiguous in head_dim.") - hp_page_ids = metadata._fp4_mla_hp_page_indices + hp_page_ids = metadata.fp4_mla_state.hp_page_indices if not isinstance(hp_page_ids, torch.Tensor): raise RuntimeError("FP4 MLA context cache update requires HP page metadata.") store_hp_tail = num_contexts > 0 @@ -1334,13 +1335,13 @@ def _scatter_fp4_mla_kv_cache_2d_context( rotary_cos_sin_ptr, hp_pool, hp_page_ids, - metadata.batch_indices, - metadata.positions, - metadata.paged_kv_indices, - metadata.paged_kv_indptr, - metadata.paged_kv_indices.shape[0], - metadata.paged_kv_indptr.shape[0], - metadata.batch_indices.shape[0], + metadata.fp4_mla_state.batch_indices, + metadata.fp4_mla_state.positions, + metadata.fp4_mla_state.paged_kv_indices, + metadata.fp4_mla_state.paged_kv_indptr, + metadata.fp4_mla_state.paged_kv_indices.shape[0], + metadata.fp4_mla_state.paged_kv_indptr.shape[0], + metadata.fp4_mla_state.batch_indices.shape[0], v_sf.shape[1], v_sf.shape[0], num_contexts, @@ -1388,7 +1389,7 @@ def _fp4_mla_generation_num_blocks_device(metadata: Any) -> torch.Tensor: slots are never consumed. """ num_gen = metadata.num_seqs - metadata.num_contexts - return metadata.paged_kv_indptr_decode[num_gen : num_gen + 1] + return metadata.fp4_mla_state.paged_kv_indptr_decode[num_gen : num_gen + 1] def _fp4_mla_uniform_generation_lengths( @@ -1402,8 +1403,8 @@ def _fp4_mla_uniform_generation_lengths( num_seqs = metadata.num_seqs kv_lens_gen = metadata.kv_lens_cuda_runtime[num_contexts:num_seqs] prompt_lens_gen = metadata.prompt_lens_cuda_runtime[num_contexts:num_seqs] - corrected_kv_lens = getattr(metadata, "fp4_mla_generation_kv_lens", None) - generation_lens = getattr(metadata, "fp4_mla_generation_append_lens", None) + corrected_kv_lens = getattr(metadata.fp4_mla_state, "generation_kv_lens", None) + generation_lens = getattr(metadata.fp4_mla_state, "generation_append_lens", None) tensors = ( kv_lens_gen, prompt_lens_gen, @@ -1421,17 +1422,13 @@ def _fp4_mla_uniform_generation_lengths( record_for_capture = bool( getattr(metadata, "is_cuda_graph", False) and torch.cuda.is_current_stream_capturing() - and not getattr( - metadata, - "_fp4_mla_generation_lengths_capture_recorded", - False, - ) + and not getattr(metadata.fp4_mla_state, "generation_lengths_capture_recorded", False) ) precomputed = ( not record_for_capture - and metadata.fp4_mla_generation_lengths_num_tokens == num_gen_tokens - and metadata.fp4_mla_generation_lengths_num_seqs == num_gen - and metadata.fp4_mla_generation_lengths_num_contexts == num_contexts + and metadata.fp4_mla_state.generation_lengths_num_tokens == num_gen_tokens + and metadata.fp4_mla_state.generation_lengths_num_seqs == num_gen + and metadata.fp4_mla_state.generation_lengths_num_contexts == num_contexts ) if not precomputed: populate_fp4_mla_generation_lengths( @@ -1442,11 +1439,11 @@ def _fp4_mla_uniform_generation_lengths( num_gen_tokens=num_gen_tokens, num_gen=num_gen, ) - metadata.fp4_mla_generation_lengths_num_tokens = num_gen_tokens - metadata.fp4_mla_generation_lengths_num_seqs = num_gen - metadata.fp4_mla_generation_lengths_num_contexts = num_contexts + metadata.fp4_mla_state.generation_lengths_num_tokens = num_gen_tokens + metadata.fp4_mla_state.generation_lengths_num_seqs = num_gen + metadata.fp4_mla_state.generation_lengths_num_contexts = num_contexts if record_for_capture: - metadata._fp4_mla_generation_lengths_capture_recorded = True + metadata.fp4_mla_state.generation_lengths_capture_recorded = True return corrected_kv_lens[:num_gen], generation_lens[:num_gen] @@ -1455,9 +1452,9 @@ def _materialize_fp4_mla_device_page_table_for_forward( generation_kv_lens: Optional[torch.Tensor] = None, ) -> None: """Materialize all fixed-stride rows from final per-forward device lengths.""" - if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): + if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): raise RuntimeError("FP4 MLA cache update requires fixed-stride device page metadata.") - if bool(getattr(metadata, "_fp4_mla_device_page_table_valid", False)): + if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): return num_contexts = int(metadata.num_contexts) @@ -1528,7 +1525,7 @@ def _scatter_fp4_mla_kv_cache_2d_generation( kv_lens_gen, gen_lens_gen = _fp4_mla_uniform_generation_lengths(metadata, num_tokens, num_gen) _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens_gen) - pool = getattr(metadata, "high_precision_kv_pool", None) + pool = getattr(metadata.fp4_mla_state, "hp_pool", None) if pool is None: raise RuntimeError("FP4 MLA 2D generation scatter requires the HP KV pool.") try: @@ -1661,7 +1658,7 @@ def launch_generation_update( gen_lens_gen, page_ids, hp_page_ids, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state.paged_kv_indptr_decode, page_ids_len, hp_page_ids.numel(), indptr_len, @@ -1768,16 +1765,16 @@ def launch_generation_update( ) preload_keys = _fp4_mla_triton_preload_key_set(metadata) if preload_key not in preload_keys: - context_page_ids = getattr(metadata, "_paged_kv_indices", None) + context_page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) if not isinstance(context_page_ids, torch.Tensor): context_page_ids = page_ids - context_indptr = getattr(metadata, "_paged_kv_indptr", None) + context_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) if not isinstance(context_indptr, torch.Tensor): - context_indptr = metadata.paged_kv_indptr_decode - context_batch_indices = getattr(metadata, "batch_indices", None) + context_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode + context_batch_indices = getattr(metadata.fp4_mla_state, "batch_indices", None) if not isinstance(context_batch_indices, torch.Tensor): context_batch_indices = hp_page_ids - context_positions = getattr(metadata, "positions", None) + context_positions = getattr(metadata.fp4_mla_state, "positions", None) if not isinstance(context_positions, torch.Tensor): context_positions = hp_page_ids _fp4_mla_context_cache_update_kernel[(1, num_dim_blocks)]( @@ -1863,7 +1860,7 @@ def launch_generation_update( launch_generation_update( launch_grid, page_ids_len=page_ids.shape[0], - indptr_len=metadata.paged_kv_indptr_decode.shape[0], + indptr_len=metadata.fp4_mla_state.paged_kv_indptr_decode.shape[0], max_gen_tiles_variant=max(max_gen_tiles, 1), q_prefix_block_dim_variant=q_prefix_block_dim, q_prefix_blocks_variant=q_prefix_blocks, @@ -1919,9 +1916,9 @@ def scatter_fp4_mla_kv_cache( phase updated the HP pool. """ if phase == "generation": - metadata._fp4_mla_prequantized_q = None - metadata._fp4_mla_prequantized_q_sf = None - metadata._fp4_mla_q_batch_capacity = None + metadata.fp4_mla_state.prequantized_q = None + metadata.fp4_mla_state.prequantized_q_sf = None + metadata.fp4_mla_state.q_batch_capacity = None if latent_cache.numel() == 0: raise ValueError("FP4 MLA cache scatter requires at least one latent token.") @@ -1932,8 +1929,8 @@ def scatter_fp4_mla_kv_cache( raise ValueError( f"FP4 MLA KV head_dim must be divisible by {FP4_BLOCK_SIZE}, got {head_dim}." ) - indices_len = metadata.batch_indices.shape[0] - positions_len = metadata.positions.shape[0] + indices_len = metadata.fp4_mla_state.batch_indices.shape[0] + positions_len = metadata.fp4_mla_state.positions.shape[0] if token_offset + num_tokens > indices_len or token_offset + num_tokens > positions_len: raise RuntimeError( f"FP4 MLA scatter would read batch_indices[{token_offset}:" @@ -1959,7 +1956,7 @@ def scatter_fp4_mla_kv_cache( if phase not in ("context", "generation"): raise ValueError("FP4 MLA scatter requires phase='context' or 'generation'.") - if getattr(metadata, "fp4_mla_v_scale_pool", None) is None: + if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: raise RuntimeError("FP4 MLA scatter requires the auxiliary V scale pool.") if metadata.page_size % FP4_BLOCK_SIZE != 0: raise ValueError( @@ -2009,9 +2006,9 @@ def scatter_fp4_mla_kv_cache( q_quant_input.shape[1], q_quant_input.device, ) - metadata._fp4_mla_prequantized_q = q_fp4_out - metadata._fp4_mla_prequantized_q_sf = q_sf_out - metadata._fp4_mla_q_batch_capacity = q_batch_capacity + metadata.fp4_mla_state.prequantized_q = q_fp4_out + metadata.fp4_mla_state.prequantized_q_sf = q_sf_out + metadata.fp4_mla_state.q_batch_capacity = q_batch_capacity hp_pool_updated = True cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND fused_v_transpose = cutedsl_backend and _fp4_mla_cutedsl_fused_v_transpose_enabled() @@ -2104,7 +2101,7 @@ def scatter_fp4_mla_kv_cache( v_packed_base=v_packed_base, v_page_offset=v_page_offset, ) - v_pack_page_ids = metadata.paged_kv_indices + v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices else: generation_state = _scatter_fp4_mla_kv_cache_2d_generation( metadata, @@ -2144,18 +2141,20 @@ def scatter_fp4_mla_kv_cache( if cutedsl_backend and not fused_v_transpose and not direct_v_packed_write: if phase == "context": num_contexts = metadata.num_contexts - cutedsl_v_pack_page_ids = metadata.paged_kv_indices[: metadata.num_context_blocks] - cutedsl_repack_page_indptr = metadata.paged_kv_indptr[: num_contexts + 1] + cutedsl_v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices[ + : metadata.fp4_mla_state.num_context_blocks + ] + cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr[: num_contexts + 1] cutedsl_repack_kv_lens = metadata.kv_lens_cuda_runtime[:num_contexts] cutedsl_repack_generation_lens = metadata.prompt_lens_cuda_runtime[:num_contexts] cutedsl_repack_max_touched_pages = int( - metadata.fp4_mla_context_repack_max_touched_pages + metadata.fp4_mla_state.context_repack_max_touched_pages ) elif generation_state is not None: kv_lens_gen, gen_lens_gen, generation_page_ids = generation_state num_gen = kv_lens_gen.numel() cutedsl_v_pack_page_ids = generation_page_ids - cutedsl_repack_page_indptr = metadata.paged_kv_indptr_decode + cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode cutedsl_repack_kv_lens = kv_lens_gen cutedsl_repack_generation_lens = gen_lens_gen cutedsl_repack_max_touched_pages = _ceil_div( @@ -2235,7 +2234,8 @@ def _ensure_workspace_tensor( dtype: torch.dtype, device: torch.device, ) -> torch.Tensor: - tensor = getattr(metadata, attr_name, None) + workspaces = metadata.fp4_mla_state.workspaces + tensor = workspaces.get(attr_name) needs_alloc = ( tensor is None or tensor.dtype != dtype @@ -2250,7 +2250,7 @@ def _ensure_workspace_tensor( "Run a warmup prepare/forward first." ) tensor = torch.empty(shape, dtype=dtype, device=device) - setattr(metadata, attr_name, tensor) + workspaces[attr_name] = tensor slices = tuple(slice(0, dim) for dim in shape) return tensor[slices] @@ -2483,19 +2483,15 @@ def _set_triton_v_packed_cache_valid( if _shared_v_pack_storage_enabled() else _triton_v_packed_valid_attr(layer_idx) ) - setattr( - metadata, - valid_attr, - _triton_v_packed_cache_tag( - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ), + metadata.fp4_mla_state.v_packed_cache_tags[valid_attr] = _triton_v_packed_cache_tag( + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, ) @@ -2516,7 +2512,7 @@ def _is_triton_v_packed_cache_valid( if _shared_v_pack_storage_enabled() else _triton_v_packed_valid_attr(layer_idx) ) - return getattr(metadata, valid_attr, None) == _triton_v_packed_cache_tag( + return metadata.fp4_mla_state.v_packed_cache_tags.get(valid_attr) == _triton_v_packed_cache_tag( layer_idx, kv_cache, v_head_dim=v_head_dim, @@ -2554,7 +2550,7 @@ def _get_triton_v_packed_cache( page_ids=page_ids, ): return None - v_packed = getattr(metadata, _triton_v_packed_attr(layer_idx), None) + v_packed = metadata.fp4_mla_state.workspaces.get(_triton_v_packed_attr(layer_idx)) expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) if ( v_packed is None @@ -2654,9 +2650,9 @@ def _max_generation_pages(metadata: Any) -> int: num_gen = metadata.num_seqs - metadata.num_contexts if num_gen <= 0: return 0 - if not getattr(metadata, "_fp4_mla_device_page_table", False): + if not getattr(metadata.fp4_mla_state, "device_page_table", False): raise RuntimeError("FP4 MLA generation requires fixed-stride device page metadata.") - max_pages = int(metadata.fp4_mla_page_table_stride) + max_pages = int(metadata.fp4_mla_state.page_table_stride) if max_pages <= 0: raise RuntimeError("FP4 MLA device page-table stride must be positive.") return max_pages @@ -2671,7 +2667,7 @@ def _fp4_mla_generation_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tens f"{num_gen_seqs} != {expected_num_gen}." ) max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata, "_paged_kv_indices", None) + page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) start = metadata.num_contexts * max_pages end = start + num_gen_seqs * max_pages if ( @@ -2694,7 +2690,7 @@ def _fp4_mla_generation_hp_page_ids(metadata: Any, num_gen_seqs: int) -> torch.T f"{num_gen_seqs} != {expected_num_gen}." ) max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) + page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) start = metadata.num_contexts * max_pages end = start + num_gen_seqs * max_pages if ( @@ -2727,7 +2723,7 @@ def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> b start = metadata.num_contexts end = metadata.num_seqs - block_counts = _host_int_list(getattr(metadata, "num_blocks", None), start, end) + block_counts = _host_int_list(getattr(metadata.fp4_mla_state, "num_blocks", None), start, end) if block_counts is not None and ( not block_counts or min(block_counts) != max_pages or max(block_counts) != max_pages ): @@ -2743,12 +2739,12 @@ def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> b tuple(block_counts) if block_counts is not None else None, kv_lens_cuda.data_ptr(), ) - cache = getattr(metadata, "_fp4_mla_full_pages_cache", None) + cache = getattr(metadata.fp4_mla_state, "full_pages_cache", None) if cache is not None and cache[0] == cache_key: return bool(cache[1]) kv_lens = [int(item) for item in kv_lens_cuda[start:end].detach().cpu().tolist()] result = bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size - setattr(metadata, "_fp4_mla_full_pages_cache", (cache_key, result)) + setattr(metadata.fp4_mla_state, "full_pages_cache", (cache_key, result)) return result kv_cache_params = getattr(metadata, "kv_cache_params", None) @@ -2990,7 +2986,7 @@ def _tma_alloc(size: int, alignment: int, stream): global_scale, q_global_scale, src_page_ids, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -3055,8 +3051,8 @@ def _tma_alloc(size: int, alignment: int, stream): # reduce during capture so we never allocate mid-capture. The single-level # reduce is numerically identical (it just launches fewer CTAs). if num_reduce_groups > 1 and torch.cuda.is_current_stream_capturing(): - gmax = getattr(metadata, "_fp4_mla_attention_group_max_buf", None) - gsum = getattr(metadata, "_fp4_mla_attention_group_sum_buf", None) + gmax = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_max_buf") + gsum = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_sum_buf") groups_ready = ( gmax is not None and gsum is not None @@ -3136,7 +3132,7 @@ def _tma_alloc(size: int, alignment: int, stream): max_scores, denom, page_max, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], max_scores.stride(0), @@ -3206,7 +3202,7 @@ def _tma_alloc(size: int, alignment: int, stream): # Fall back to the unsplit PV (numerically identical) during capture unless a # warmup forward already sized that workspace, so capture never allocates. if page_split > 1 and torch.cuda.is_current_stream_capturing(): - pbuf = getattr(metadata, "_fp4_mla_attention_pv_partial_buf", None) + pbuf = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_pv_partial_buf") partial_ready = ( pbuf is not None and pbuf.shape[0] >= num_queries @@ -3236,7 +3232,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -3285,7 +3281,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -3359,7 +3355,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -3400,7 +3396,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.paged_kv_indptr_decode, + metadata.fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -3497,7 +3493,7 @@ def run_fp4_mla_attention_decode( raise ValueError("FP4 MLA attention output batch dimensions do not match.") backend = _fp4_mla_attention_backend() - if getattr(metadata, "fp4_mla_v_scale_pool", None) is None: + if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: raise RuntimeError( "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." ) @@ -3549,7 +3545,7 @@ def run_fp4_mla_attention_decode( ) sf_cache = sf_cache.view(torch.float8_e4m3fn) - v_sf_pool = metadata.fp4_mla_v_scale_pool + v_sf_pool = metadata.fp4_mla_state.v_scale_pool v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=kv_lora_rank)[local_layer].view( torch.float8_e4m3fn ) @@ -3717,7 +3713,7 @@ def run_fp4_mla_attention_decode( core_v_sf, global_scale, src_page_ids, - metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], + metadata.fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], kv_lens, kernel_output, max_kv_len=max_pages * metadata.page_size, diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py index 0cfd34cb55ec..4b50fa3de793 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py @@ -34,7 +34,6 @@ _FP8_CONTEXT_SUPPORTED_SMS = {90, 100, 103, 107, 120} _FP8_CONTEXT_SCRATCH_ATTR = "_fp4_mla_fp8_context_scratch" -_FP8_CONTEXT_ATTN_ATTR = "_fp4_mla_fp8_context_attn" def require_fp4_mla_fp8_context_support() -> None: @@ -303,7 +302,7 @@ def _build_fp8_mla_context_metadata( ) -> "TrtllmAttentionMetadata": """Route the mandatory FP8 cache write through a direct metadata view.""" fp8_meta = copy.copy(meta) - fp8_meta._fp4_mla_fp8_context_state = None + fp8_meta.fp4_mla_state = None fp8_meta.kv_cache_manager = scratch.cache_manager_view fp8_meta.kv_cache_block_offsets = scratch.block_offsets fp8_meta.block_ids_per_seq = scratch.block_ids_per_seq @@ -319,7 +318,7 @@ def _build_fp8_mla_context_metadata( # Scratch lengths intentionally start from zero. Preserve the actual # absolute positions for Q/K RoPE through the native kernel's explicit # per-token position-offset input. - fp8_meta.helix_position_offsets = meta.positions[: meta.num_ctx_tokens] + fp8_meta.helix_position_offsets = meta.fp4_mla_state.positions[: meta.num_ctx_tokens] return fp8_meta @@ -328,11 +327,11 @@ def _get_fp8_mla_context_metadata( scratch: _Fp8MlaContextScratch, ) -> "TrtllmAttentionMetadata": """Prepare and reuse one direct metadata view for all layers in a step.""" - state = meta._fp4_mla_fp8_context_state + state = meta.fp4_mla_state.fp8_context_state if state is not None and state[0] is scratch: return state[1] scratch.prepare(meta) fp8_meta = _build_fp8_mla_context_metadata(meta, scratch) - meta._fp4_mla_fp8_context_state = (scratch, fp8_meta) + meta.fp4_mla_state.fp8_context_state = (scratch, fp8_meta) return fp8_meta diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py new file mode 100644 index 000000000000..ea5465a1b2a5 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Batch-shared FP4 MLA state with metadata/CUDA-graph buffer lifetimes.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import torch + +from . import ( + FP4_MLA_KV_GLOBAL_SCALE, + FP4_MLA_Q_GLOBAL_SCALE, + HP_BLOCK_SIZE, + configure_fp4_mla_device_page_table, + populate_fp4_mla_append_metadata, +) + +if TYPE_CHECKING: + from ....memory_buffer_utils import Buffers + from ..trtllm import TrtllmAttentionMetadata + from .fp4_mla_context import _Fp8MlaContextScratch + + +@dataclass +class Fp4MlaState: + """State shared by RoPE/cache append and all FP4 FMHAs in one batch. + + Allocate through the metadata buffer allocator so eager forwards and each + CUDA graph retain their existing storage/aliasing rules. This state is not + layer-local: MTP and metadata preparation update it before any FMHA runs. + """ + + hp_pool: torch.Tensor | None = None + hp_page_indices: torch.Tensor | None = None + v_scale_pool: torch.Tensor | None = None + q_global_scale: torch.Tensor | None = None + kv_global_scale: torch.Tensor | None = None + batch_indices: torch.Tensor | None = None + positions: torch.Tensor | None = None + generation_kv_lens: torch.Tensor | None = None + generation_append_lens: torch.Tensor | None = None + generation_lengths_num_tokens: int = -1 + generation_lengths_num_seqs: int = -1 + generation_lengths_num_contexts: int = -1 + generation_lengths_capture_recorded: bool = False + generation_cache_scattered: bool = False + device_page_table: bool = False + device_page_table_valid: bool = False + page_table_stride: int = 0 + context_repack_max_touched_pages: int = 1 + cache_pool_id: int = 0 + cache_page_index_scale: int = 0 + hp_pool_id: int = 0 + hp_page_index_scale: int = 0 + prequantized_q: torch.Tensor | None = None + prequantized_q_sf: torch.Tensor | None = None + q_batch_capacity: int | None = None + fp8_context_state: tuple[_Fp8MlaContextScratch, TrtllmAttentionMetadata] | None = None + _paged_kv_indptr: torch.Tensor | None = None + paged_kv_indptr_decode: torch.Tensor | None = None + _paged_kv_indices: torch.Tensor | None = None + num_sequences: int = 0 + num_blocks: list[int] | None = None + num_context_blocks: int = 0 + num_generation_blocks: int = 0 + full_pages_cache: tuple[tuple, bool] | None = None + workspaces: dict[str, torch.Tensor] = field(default_factory=dict) + v_packed_cache_tags: dict[str, tuple] = field(default_factory=dict) + + @property + def paged_kv_indices(self) -> torch.Tensor: + if self._paged_kv_indices is None: + raise RuntimeError("FP4 MLA paged_kv_indices is not allocated.") + return self._paged_kv_indices[: self.num_context_blocks + self.num_generation_blocks] + + @property + def paged_kv_indptr(self) -> torch.Tensor: + if self._paged_kv_indptr is None: + raise RuntimeError("FP4 MLA paged_kv_indptr is not allocated.") + return self._paged_kv_indptr[: self.num_sequences + 1] + + @classmethod + def create(cls, metadata: TrtllmAttentionMetadata, buffers: Buffers | None) -> Fp4MlaState: + manager = metadata.kv_cache_manager + hp_ring_size = manager.fp4_mla_hp_pool_size + if hp_ring_size < HP_BLOCK_SIZE: + raise RuntimeError( + "FP4 MLA high-precision ring must retain at least one " + f"{HP_BLOCK_SIZE}-token quantization tile, got {hp_ring_size} slots." + ) + hp_pool = manager.get_fp4_mla_hp_pool() + if ( + not isinstance(hp_pool, torch.Tensor) + or hp_pool.dtype != torch.bfloat16 + or hp_pool.device.type != "cuda" + or hp_pool.ndim != 4 + or hp_pool.shape[1] != manager.num_local_layers + or hp_pool.shape[2] != manager.kv_factor + or hp_pool.shape[3] != hp_ring_size * manager.head_dim + ): + raise RuntimeError( + "FP4 MLA V2 HP pool must be a CUDA BF16 tensor shaped " + "[pages, local_layers, kv_factor, ring * head_dim], got " + f"{getattr(hp_pool, 'shape', None)}." + ) + + def allocate( + name: str, shape: tuple[int, ...], dtype: torch.dtype = torch.int32 + ) -> torch.Tensor: + return metadata.get_empty( + buffers, + shape, + cache_name=f"fp4_mla_{name}", + dtype=dtype, + capture_graph=metadata.is_cuda_graph, + ) + + state = cls(hp_pool=hp_pool, v_scale_pool=manager.get_mla_v_scale_pool()) + tokens = (metadata.max_num_tokens,) + sequences = (metadata.max_num_sequences,) + indptr = (metadata.max_num_sequences + 1,) + pages = (metadata.max_num_sequences * int(manager.max_blocks_per_seq),) + state.batch_indices = allocate("batch_indices", tokens) + state.positions = allocate("positions", tokens) + state.generation_kv_lens = allocate("generation_kv_lens", sequences) + state.generation_append_lens = allocate("generation_append_lens", sequences) + state._paged_kv_indices = allocate("paged_kv_indices", pages) + state.hp_page_indices = allocate("hp_page_indices", pages) + state._paged_kv_indptr = allocate("paged_kv_indptr", indptr) + state.paged_kv_indptr_decode = allocate("paged_kv_indptr_decode", indptr) + state.q_global_scale = allocate("q_global_scale", (1,), torch.float32) + state.q_global_scale.fill_(FP4_MLA_Q_GLOBAL_SCALE) + state.kv_global_scale = allocate("kv_global_scale", (1,), torch.float32) + state.kv_global_scale.fill_(FP4_MLA_KV_GLOBAL_SCALE) + return state + + def invalidate_generation_lengths(self) -> None: + self.generation_lengths_num_tokens = -1 + self.generation_lengths_num_seqs = -1 + self.generation_lengths_num_contexts = -1 + self.generation_lengths_capture_recorded = False + + def prepare(self, metadata: TrtllmAttentionMetadata, kv_lens: torch.Tensor) -> None: + self.invalidate_generation_lengths() + if metadata.kv_cache_manager is None or metadata.request_ids is None: + raise RuntimeError( + "FP4 MLA device page metadata requires a KV cache manager and request IDs." + ) + if not configure_fp4_mla_device_page_table(metadata, kv_lens): + raise RuntimeError( + "FP4 MLA requires fixed-stride device page metadata; the " + "current KV-cache manager or batch layout is unsupported." + ) + # Graph capture's on_update_kv_lens owns the append kernel. Uniform + # generation derives token positions in the fused update instead. + if not metadata.is_cuda_graph and metadata.num_tokens > 0 and metadata.num_contexts > 0: + self.populate_append_metadata(metadata) + + def on_update_kv_lens(self, metadata: TrtllmAttentionMetadata) -> None: + self.device_page_table_valid = False + if metadata.num_tokens > 0: + self.invalidate_generation_lengths() + if metadata.num_contexts > 0: + self.populate_append_metadata(metadata) + + def update_for_spec_dec(self, metadata: TrtllmAttentionMetadata) -> None: + num_seqs = metadata.num_seqs + metadata.prompt_lens_cuda_runtime = metadata.seq_lens_kv_cuda[:num_seqs] + if not torch.cuda.is_current_stream_capturing(): + metadata.prompt_lens_cpu_runtime = metadata.seq_lens_kv[:num_seqs] + self.on_update_kv_lens(metadata) + + def restore_from_spec_dec(self, metadata: TrtllmAttentionMetadata) -> None: + # Rebind aliases to the restored stable tensors, never to the temporary + # spec-dec clone whose storage may be released after graph capture. + num_seqs = metadata.num_seqs + metadata.kv_lens_cuda_runtime = metadata.kv_lens_cuda[:num_seqs] + metadata.prompt_lens_cuda_runtime = metadata.seq_lens_kv_cuda[:num_seqs] + if not torch.cuda.is_current_stream_capturing(): + metadata.prompt_lens_cpu_runtime = metadata.seq_lens_kv[:num_seqs] + + def populate_append_metadata(self, metadata: TrtllmAttentionMetadata) -> None: + num_seqs = metadata.num_contexts + metadata.num_generations + if num_seqs == 0 or metadata.num_tokens == 0: + return + assert self.batch_indices is not None + assert self.positions is not None + assert metadata.kv_lens_cuda_runtime is not None + assert metadata.prompt_lens_cuda_runtime is not None + # Canonical tensors can change between MTP sub-steps. Do not capture + # stale *_runtime views when writing token positions and batch indices. + append_lens = metadata.seq_lens_kv_cuda[:num_seqs] + if not append_lens.is_cuda: + raise RuntimeError("FP4 MLA append metadata requires CUDA sequence lengths.") + populate_fp4_mla_append_metadata( + append_lens, + metadata.kv_lens_cuda[:num_seqs], + self.batch_indices, + self.positions, + num_tokens=metadata.num_tokens, + num_sequences=num_seqs, + num_contexts=metadata.num_contexts, + num_context_tokens=metadata.num_ctx_tokens, + ) diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index bf2a98487b7e..55b88face9d4 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -1080,6 +1080,11 @@ def support_fused_qkv(cls) -> bool: def support_mla(cls) -> bool: return False + @classmethod + def support_fp4_kv_cache(cls) -> bool: + """Whether the backend can execute attention with an FP4 KV cache.""" + return False + @classmethod def support_multi_item_scoring(cls) -> bool: return False diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index ba51b5ec1e1a..7a6a83dc0a7c 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -18,7 +18,7 @@ import os import weakref from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -41,11 +41,8 @@ from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from .fmha.manager import FmhaManager -from .fp4_mla import (FP4_MLA_KV_GLOBAL_SCALE, FP4_MLA_Q_GLOBAL_SCALE, - HP_BLOCK_SIZE, can_fuse_fp4_mla_q_quant, - configure_fp4_mla_device_page_table, - populate_fp4_mla_append_metadata, - scatter_fp4_mla_kv_cache) +from .fp4_mla import can_fuse_fp4_mla_q_quant, scatter_fp4_mla_kv_cache +from .fp4_mla.state import Fp4MlaState from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, KVCacheParams, MLAParams, PositionalEmbeddingParams, @@ -230,55 +227,11 @@ def effective_beam_width(self) -> int: # True during warmup forward passes (dummy requests, no real data). is_warmup: bool = False - # High-precision BF16 KV pool for MLA FP4 models. The FP4 MLA V2 manager - # exposes compact paged rings. - high_precision_kv_pool: Optional[torch.Tensor] = None - _fp4_mla_hp_page_indices: Optional[torch.Tensor] = None - fp4_mla_v_scale_pool: Optional[torch.Tensor] = None - _fp4_mla_q_global_scale: Optional[torch.Tensor] = None - _fp4_mla_kv_global_scale: Optional[torch.Tensor] = None - batch_indices: Optional[torch.Tensor] = None - positions: Optional[torch.Tensor] = None - fp4_mla_generation_kv_lens: Optional[torch.Tensor] = None - fp4_mla_generation_append_lens: Optional[torch.Tensor] = None - fp4_mla_generation_lengths_num_tokens: int = field(init=False, default=-1) - fp4_mla_generation_lengths_num_seqs: int = field(init=False, default=-1) - fp4_mla_generation_lengths_num_contexts: int = field(init=False, default=-1) - _fp4_mla_generation_lengths_capture_recorded: bool = field(init=False, - default=False, - repr=False) - _fp4_mla_generation_cache_scattered: bool = field(init=False, - default=False, - repr=False) - _fp4_mla_device_page_table: bool = field(init=False, - default=False, - repr=False) - _fp4_mla_device_page_table_valid: bool = field(init=False, - default=False, - repr=False) - fp4_mla_page_table_stride: int = field(init=False, default=0, repr=False) - fp4_mla_context_repack_max_touched_pages: int = field(init=False, - default=1, - repr=False) - _fp4_mla_prequantized_q: Optional[torch.Tensor] = field(init=False, - default=None, - repr=False) - _fp4_mla_prequantized_q_sf: Optional[torch.Tensor] = field(init=False, - default=None, - repr=False) - _fp4_mla_q_batch_capacity: Optional[int] = field(init=False, - default=None, - repr=False) - _fp4_mla_fp8_context_state: Optional[Tuple[Any, Any]] = field(init=False, - default=None, - repr=False, - compare=False) - _paged_kv_indptr: Optional[torch.Tensor] = None - paged_kv_indptr_decode: Optional[torch.Tensor] = None - _paged_kv_indices: Optional[torch.Tensor] = None - num_blocks: Optional[List[int]] = None - num_context_blocks: int = 0 - num_generation_blocks: int = 0 + # Batch-shared FP4 state; other attention paths allocate none of it. + fp4_mla_state: Optional[Fp4MlaState] = field(init=False, + default=None, + repr=False, + compare=False) # Pre-computed FlashMLA tile-scheduler metadata and num_splits. # Computed once per forward pass in TrtllmAttention.forward() and reused across layers. @@ -378,25 +331,6 @@ def page_size(self) -> int: def page_size(self, value: int) -> None: self._page_size_override = value - @property - def paged_kv_indices(self) -> torch.Tensor: - """ - Flattened page table used by FP4 MLA helper kernels. - """ - if self._paged_kv_indices is None: - raise RuntimeError("paged_kv_indices is not allocated.") - total_blocks = self.num_context_blocks + self.num_generation_blocks - return self._paged_kv_indices[:total_blocks] - - @property - def paged_kv_indptr(self) -> torch.Tensor: - """ - Page-table indptr used by FP4 MLA helper kernels. - """ - if self._paged_kv_indptr is None: - raise RuntimeError("paged_kv_indptr is not allocated.") - return self._paged_kv_indptr[:self.num_seqs + 1] - @property def host_kv_cache_pool_pointers(self) -> Optional[torch.Tensor]: """ @@ -653,109 +587,10 @@ def _post_init_with_buffers(self, buffers) -> None: pin_memory=prefer_pinned(), ) - # Bind the native paged high-precision BF16 KV pool for MLA FP4 models. if (self.kv_cache_manager is not None and self.kv_cache_manager.kv_factor == 1 and self.kv_cache_manager.dtype == DataType.NVFP4): - num_local_layers = self.kv_cache_manager.num_local_layers - head_dim = self.kv_cache_manager.head_dim - kv_factor = self.kv_cache_manager.kv_factor - hp_ring_size = self.kv_cache_manager.fp4_mla_hp_pool_size - if hp_ring_size < HP_BLOCK_SIZE: - raise RuntimeError( - "FP4 MLA high-precision ring must retain at least one " - f"{HP_BLOCK_SIZE}-token quantization tile, got " - f"{hp_ring_size} slots.") - hp_pool = self.kv_cache_manager.get_fp4_mla_hp_pool() - expected_tail = hp_ring_size * head_dim - if (not isinstance(hp_pool, torch.Tensor) - or hp_pool.dtype != torch.bfloat16 - or hp_pool.device.type != 'cuda' or hp_pool.ndim != 4 - or hp_pool.shape[1] != num_local_layers - or hp_pool.shape[2] != kv_factor - or hp_pool.shape[3] != expected_tail): - raise RuntimeError( - "FP4 MLA V2 HP pool must be a CUDA BF16 tensor shaped " - "[pages, local_layers, kv_factor, ring * head_dim], got " - f"{getattr(hp_pool, 'shape', None)}.") - self.high_precision_kv_pool = hp_pool - self.batch_indices = self.get_empty( - buffers, - (self.max_num_tokens, ), - cache_name="fp4_mla_batch_indices", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.positions = self.get_empty( - buffers, - (self.max_num_tokens, ), - cache_name="fp4_mla_positions", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.fp4_mla_generation_kv_lens = self.get_empty( - buffers, - (self.max_num_sequences, ), - cache_name="fp4_mla_generation_kv_lens", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.fp4_mla_generation_append_lens = self.get_empty( - buffers, - (self.max_num_sequences, ), - cache_name="fp4_mla_generation_append_lens", - dtype=torch.int32, - capture_graph=capture_graph, - ) - page_table_capacity = ( - self.max_num_sequences * - int(self.kv_cache_manager.max_blocks_per_seq)) - self._paged_kv_indices = self.get_empty( - buffers, - (page_table_capacity, ), - cache_name="fp4_mla_paged_kv_indices", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self._fp4_mla_hp_page_indices = self.get_empty( - buffers, - (page_table_capacity, ), - cache_name="fp4_mla_hp_page_indices", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self._paged_kv_indptr = self.get_empty( - buffers, - (self.max_num_sequences + 1, ), - cache_name="fp4_mla_paged_kv_indptr", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.paged_kv_indptr_decode = self.get_empty( - buffers, - (self.max_num_sequences + 1, ), - cache_name="fp4_mla_paged_kv_indptr_decode", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self._fp4_mla_q_global_scale = self.get_empty( - buffers, - (1, ), - cache_name="fp4_mla_q_global_scale", - dtype=torch.float32, - capture_graph=capture_graph, - ) - self._fp4_mla_q_global_scale.fill_(FP4_MLA_Q_GLOBAL_SCALE) - self._fp4_mla_kv_global_scale = self.get_empty( - buffers, - (1, ), - cache_name="fp4_mla_kv_global_scale", - dtype=torch.float32, - capture_graph=capture_graph, - ) - self._fp4_mla_kv_global_scale.fill_(FP4_MLA_KV_GLOBAL_SCALE) - self.fp4_mla_v_scale_pool = self.kv_cache_manager.get_mla_v_scale_pool( - ) + self.fp4_mla_state = Fp4MlaState.create(self, buffers) # Allocate static buffers for helix parallelism support. if self.enable_helix: @@ -820,37 +655,21 @@ def _post_init_with_buffers(self, buffers) -> None: self._helix_spec_tokens_valid = False def on_update_kv_lens(self): - # After changing the kv_lens/kv_lens_cuda, we may need to update other metadata. - # Especially for the changes in the _preprocess_inputs() of model_engine.py. + # KV lengths can change between speculative decoding sub-steps. if self.enable_flash_mla: self._flash_mla_metadata_valid = False self._invalidate_mla_scheduler_buffers() - if getattr(self, '_fp4_mla_device_page_table', False): - self._fp4_mla_device_page_table_valid = False - if getattr(self, 'high_precision_kv_pool', None) is not None: - self._update_fp4_mla_append_metadata() - - def _update_fp4_mla_append_metadata(self) -> None: - if self.high_precision_kv_pool is not None and self.num_tokens > 0: - self._invalidate_fp4_mla_generation_lengths() - if self._needs_fp4_mla_append_metadata(): - self._populate_fp4_mla_batch_indices_positions() + state = getattr(self, "fp4_mla_state", None) + if state is not None: + state.on_update_kv_lens(self) def update_for_spec_dec(self) -> None: - # MTP updates kv_lens_cuda in-place between sub-steps, which changes - # cache_seq_lens seen by the C++ attention op. Invalidate the metadata - # so that forward() recomputes it for the next sub-step. if self.enable_flash_mla: self._flash_mla_metadata_valid = False self._invalidate_mla_scheduler_buffers() - if getattr(self, 'high_precision_kv_pool', None) is None: - return - - num_seqs = self.num_seqs - self.prompt_lens_cuda_runtime = self.seq_lens_kv_cuda[:num_seqs] - if not torch.cuda.is_current_stream_capturing(): - self.prompt_lens_cpu_runtime = self.seq_lens_kv[:num_seqs] - self._update_fp4_mla_append_metadata() + state = getattr(self, "fp4_mla_state", None) + if state is not None: + state.update_for_spec_dec(self) def _invalidate_mla_scheduler_buffers(self) -> None: # Spec-dec rewrites q_lens and kv_lens between sub-steps, so the cumulative @@ -859,21 +678,9 @@ def _invalidate_mla_scheduler_buffers(self) -> None: self._mla_ctx_cu_seqlens_valid = False def restore_from_spec_dec(self) -> None: - # The spec-dec draft loop pointed the FP4 MLA length aliases at the - # temporary _seq_lens_cuda clone made by prepare_for_spec_dec. Rebind - # them to the restored stable buffers; otherwise the next forward's - # captured ops (CUDA graph) bake pointers to the clone, which is freed - # after capture, and every replay reads freed memory (garbage - # positions/kv lens -> corrupted KV writes and illegal memory access - # in the RoPE table lookup). super().restore_from_spec_dec() - if getattr(self, 'high_precision_kv_pool', None) is None: - return - num_seqs = self.num_seqs - self.kv_lens_cuda_runtime = self.kv_lens_cuda[:num_seqs] - self.prompt_lens_cuda_runtime = self.seq_lens_kv_cuda[:num_seqs] - if not torch.cuda.is_current_stream_capturing(): - self.prompt_lens_cpu_runtime = self.seq_lens_kv[:num_seqs] + if self.fp4_mla_state is not None: + self.fp4_mla_state.restore_from_spec_dec(self) def update_helix_param( self, @@ -1067,7 +874,8 @@ def restore_after_draft_forward(self, saved_state: dict | None) -> None: def prepare(self) -> None: # The FP8 scratch metadata view is shared by every local FP4 MLA layer # in one eager context forward and must be rebuilt for the next batch. - self._fp4_mla_fp8_context_state = None + if self.fp4_mla_state is not None: + self.fp4_mla_state.fp8_context_state = None super().prepare() # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers. self._invalidate_mla_scheduler_buffers() @@ -1212,7 +1020,7 @@ def prepare(self) -> None: # tokens. Use the actual KV length (without extra tokens) for # kv_lens_runtime, which becomes host_past_key_value_lengths and # eventually mMaxSeqLenKv. - if self.high_precision_kv_pool is not None: + if self.fp4_mla_state is not None: # FP4 MLA needs the per-forward append lengths rather than the # original prompt lengths. Context scratch-cache metadata and MTP # generation both consume these runtime views. @@ -1231,10 +1039,8 @@ def prepare(self) -> None: host_request_types=self.host_request_types[:self.num_seqs], ) - if self.high_precision_kv_pool is not None: - self._invalidate_fp4_mla_generation_lengths() - self._configure_fp4_mla_page_metadata(kv_lens) - self._prepare_fp4_mla_append_metadata() + if self.fp4_mla_state is not None: + self.fp4_mla_state.prepare(self, kv_lens) def prepare_encoder_decoder_from_precomputed_lengths( self, prompt_lens: torch.Tensor, kv_lens: torch.Tensor, @@ -1345,75 +1151,6 @@ def prepare_encoder_cuda_graph_replay(self, seq_lens: List[int], self._num_ctx_tokens = padded_num_tokens self.host_total_kv_lens[0] = padded_num_tokens - def _prepare_fp4_mla_append_metadata(self) -> None: - """Populate eager append metadata or defer it to CUDA graph replay. - - CUDA graph ``_forward_step`` captures ``on_update_kv_lens`` before the - model forward. Let that captured update own these buffers instead of - launching the same metadata kernel eagerly during input preparation. - - One-token generation is request-major in the fused FP4 update kernel, - so it derives the token position directly from the sequence length and - does not need these per-token buffers at capture or replay time. - """ - if (self.is_cuda_graph or self.num_tokens == 0 - or not self._needs_fp4_mla_append_metadata()): - return - self._populate_fp4_mla_batch_indices_positions() - - def _needs_fp4_mla_append_metadata(self) -> bool: - """Return whether a forward still consumes materialized token indices.""" - return self.num_contexts > 0 - - def _invalidate_fp4_mla_generation_lengths(self) -> None: - """Invalidate generation lengths before the next FP4 MLA forward step.""" - self.fp4_mla_generation_lengths_num_tokens = -1 - self.fp4_mla_generation_lengths_num_seqs = -1 - self.fp4_mla_generation_lengths_num_contexts = -1 - self._fp4_mla_generation_lengths_capture_recorded = False - - def _configure_fp4_mla_page_metadata(self, kv_lens: torch.Tensor) -> None: - """Configure fixed-stride device page metadata.""" - if self.kv_cache_manager is None or self.request_ids is None: - raise RuntimeError( - "FP4 MLA device page metadata requires a KV cache manager " - "and request IDs.") - assert self._paged_kv_indices is not None - assert self._paged_kv_indptr is not None - assert self.paged_kv_indptr_decode is not None - if not configure_fp4_mla_device_page_table(self, kv_lens): - raise RuntimeError( - "FP4 MLA requires fixed-stride device page metadata; the " - "current KV-cache manager or batch layout is unsupported.") - - def _populate_fp4_mla_batch_indices_positions(self) -> None: - """Populate FP4 MLA scatter/HP append metadata in one CUDA launch.""" - num_seqs = self.num_contexts + self.num_generations - if num_seqs == 0 or self.num_tokens == 0: - return - assert self.batch_indices is not None - assert self.positions is not None - assert self.kv_lens_cuda_runtime is not None - assert self.prompt_lens_cuda_runtime is not None - - append_lens = self.seq_lens_kv_cuda[:num_seqs] - if not append_lens.is_cuda: - raise RuntimeError( - "FP4 MLA append metadata requires CUDA sequence lengths.") - # Use the canonical tensors rather than the *_runtime aliases. A - # spec-dec sub-step can replace the canonical append lengths, and CUDA - # graph capture requires their stable, live storage. - populate_fp4_mla_append_metadata( - append_lens, - self.kv_lens_cuda[:num_seqs], - self.batch_indices, - self.positions, - num_tokens=self.num_tokens, - num_sequences=num_seqs, - num_contexts=self.num_contexts, - num_context_tokens=self.num_ctx_tokens, - ) - def prepare_flash_mla(self) -> None: self._flash_mla_metadata_valid = False # Request-specific fills and H2D copies must happen before replay, not @@ -2627,6 +2364,10 @@ def support_fused_qkv(cls) -> bool: def support_mla(cls) -> bool: return True + @classmethod + def support_fp4_kv_cache(cls) -> bool: + return True + def has_cached_kv_for_mla_context( self, metadata: TrtllmAttentionMetadata, @@ -2999,7 +2740,7 @@ def _fp4_mla_rope_generation( "FP4 MLA generation requires fused BF16 RoPE/cache update " "with an FP32 rotary table.") - metadata._fp4_mla_generation_cache_scattered = False + metadata.fp4_mla_state.generation_cache_scattered = False hp_pool_updated = scatter_fp4_mla_kv_cache( metadata, latent_cache, @@ -3016,7 +2757,7 @@ def _fp4_mla_rope_generation( if not hp_pool_updated: raise RuntimeError( "Fused FP4 MLA RoPE/cache scatter did not update the HP pool.") - metadata._fp4_mla_generation_cache_scattered = True + metadata.fp4_mla_state.generation_cache_scattered = True def can_fuse_fp4_mla_q_quant( self, diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index 64d9c8de7b1f..49903165372d 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -697,7 +697,11 @@ def create_weights(self): and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales() ) mla_weight_dtype = torch.float8_e4m3fn if has_fp8_block_scales else self.dtype - if isinstance(self.mqa, TrtllmAttention) and self.mqa.has_fp4_kv_cache: + if ( + self.mqa.support_fp4_kv_cache() + and self.quant_config is not None + and self.quant_config.layer_quant_mode.has_fp4_kv_cache() + ): if ( mla_weight_dtype != torch.bfloat16 or self.mapping.cp_size != 1 @@ -1593,7 +1597,11 @@ def forward_absorption_generation( device=q.device, ) - fp4_mla = isinstance(self.mqa, TrtllmAttention) and self.mqa.has_fp4_kv_cache + fp4_mla = ( + self.mqa.support_fp4_kv_cache() + and self.quant_config is not None + and self.quant_config.layer_quant_mode.has_fp4_kv_cache() + ) if fp4_mla and latent_cache is None: raise RuntimeError( "FP4 MLA generation requires a latent cache for fused " diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 636d639790f6..7363ace81da9 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -945,8 +945,7 @@ def change_attn_metadata(self, num_accepted_tokens: torch.Tensor, attn_metadata.kv_lens_cuda[num_contexts:batch_size].clamp_( min=runtime_draft_len) attn_metadata.on_update_kv_lens() - if getattr(attn_metadata, "high_precision_kv_pool", - None) is not None: + if getattr(attn_metadata, "fp4_mla_state", None) is not None: attn_metadata.update_for_spec_dec() if attn_metadata.kv_cache_params is not None and not attn_metadata.is_cuda_graph: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0ce6cde395b2..b3a8eb5f368c 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -300,7 +300,7 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, @skip_no_rubin @pytest.mark.skip_less_device_memory(60000) - def test_nvfp4_mla_gsm8k(self, monkeypatch): + def test_nvfp4_mla(self, monkeypatch): from tensorrt_llm._torch.attention.backends.fmha.fp4_mla import \ Fp4MlaFmha from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import \ diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index d360fb6bc84f..02ba64e2f551 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -779,7 +779,7 @@ test_e2e.py::test_trtllm_benchmark_serving[gpt_oss/gpt-oss-20b] test_e2e.py::test_trtllm_multimodal_benchmark_serving # Dense FP4 MLA accuracy (SM107 only) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_mla_gsm8k +accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_mla # fine-grained sync tests (SM107+ only) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_fine_grained_sync[enable_autotuner=False-moe_backend=TRTLLM-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index 3b3572e6aca6..c6df26fd383e 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -28,6 +28,7 @@ scatter_fp4_mla_kv_cache, ) from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 +from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, MTPDecodingConfig from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig @@ -71,9 +72,12 @@ def test_fp4_mla_generation_hp_page_ids_skip_mixed_batch_context_rows() -> None: metadata = SimpleNamespace( num_contexts=2, num_seqs=5, - _fp4_mla_device_page_table=True, - fp4_mla_page_table_stride=3, - _fp4_mla_hp_page_indices=torch.arange(15, dtype=torch.int32), + fp4_mla_state=Fp4MlaState( + num_sequences=5, + device_page_table=True, + page_table_stride=3, + hp_page_indices=torch.arange(15, dtype=torch.int32), + ), ) page_ids = fp4_mla_backend._fp4_mla_generation_hp_page_ids(metadata, 3) @@ -90,13 +94,16 @@ def test_fp4_mla_cuda_graph_generation_lengths_records_capture_once(monkeypatch) num_seqs=3, kv_lens_cuda_runtime=torch.tensor([9, 17, 25], dtype=torch.int32, device="cuda"), prompt_lens_cuda_runtime=torch.tensor([9, 1, 1], dtype=torch.int32, device="cuda"), - fp4_mla_generation_kv_lens=corrected_kv_lens, - fp4_mla_generation_append_lens=generation_lens, - fp4_mla_generation_lengths_num_tokens=4, - fp4_mla_generation_lengths_num_seqs=2, - fp4_mla_generation_lengths_num_contexts=1, - _fp4_mla_generation_lengths_capture_recorded=False, is_cuda_graph=True, + fp4_mla_state=Fp4MlaState( + num_sequences=3, + generation_kv_lens=corrected_kv_lens, + generation_append_lens=generation_lens, + generation_lengths_num_tokens=4, + generation_lengths_num_seqs=2, + generation_lengths_num_contexts=1, + generation_lengths_capture_recorded=False, + ), ) populate_calls = [] @@ -122,7 +129,7 @@ def populate_generation_lengths(*args, **kwargs) -> None: assert args[2].data_ptr() == corrected_kv_lens.data_ptr() assert args[3].data_ptr() == generation_lens.data_ptr() assert kwargs == {"num_gen_tokens": 4, "num_gen": 2} - assert metadata._fp4_mla_generation_lengths_capture_recorded + assert metadata.fp4_mla_state.generation_lengths_capture_recorded assert first[0].data_ptr() == second[0].data_ptr() == corrected_kv_lens.data_ptr() assert first[1].data_ptr() == second[1].data_ptr() == generation_lens.data_ptr() @@ -416,41 +423,42 @@ def _build_multi_seq_metadata( return SimpleNamespace( kv_cache_manager=kv_cache_manager, - batch_indices=batch_indices, - positions=positions, - paged_kv_indices=paged_kv_indices, - _paged_kv_indices=paged_kv_indices, - _fp4_mla_hp_page_indices=hp_page_indices, - paged_kv_indptr=paged_kv_indptr, - _paged_kv_indptr=paged_kv_indptr, - paged_kv_indptr_decode=paged_kv_indptr.clone(), - _fp4_mla_device_page_table=True, - _fp4_mla_device_page_table_valid=True, - fp4_mla_page_table_stride=max_blocks_per_seq, - fp4_mla_context_repack_max_touched_pages=max_blocks_per_seq, page_size=page_size, - num_context_blocks=num_seqs * max_blocks_per_seq, - num_generation_blocks=0, num_contexts=num_seqs, num_seqs=num_seqs, - num_blocks=None, - high_precision_kv_pool=hp_pool, - fp4_mla_v_scale_pool=kv_cache_manager.get_mla_v_scale_pool(), kv_lens_cuda_runtime=kv_lens, prompt_lens_cuda_runtime=prompt_lens_cuda, prompt_lens_cpu_runtime=prompt_lens_cpu, - fp4_mla_generation_kv_lens=torch.empty(num_seqs, dtype=torch.int32, device=device), - fp4_mla_generation_append_lens=torch.empty(num_seqs, dtype=torch.int32, device=device), - fp4_mla_generation_lengths_num_tokens=-1, - fp4_mla_generation_lengths_num_seqs=-1, - fp4_mla_generation_lengths_num_contexts=-1, - _fp4_mla_generation_lengths_capture_recorded=False, - _fp4_mla_q_global_scale=q_global_scale, - _fp4_mla_kv_global_scale=kv_global_scale, request_ids=request_ids, runtime_features=SimpleNamespace(has_speculative_draft_tokens=False), is_cuda_graph=False, is_warmup=False, + fp4_mla_state=Fp4MlaState( + batch_indices=batch_indices, + positions=positions, + _paged_kv_indices=paged_kv_indices, + hp_page_indices=hp_page_indices, + _paged_kv_indptr=paged_kv_indptr, + paged_kv_indptr_decode=paged_kv_indptr.clone(), + device_page_table=True, + device_page_table_valid=True, + page_table_stride=max_blocks_per_seq, + context_repack_max_touched_pages=max_blocks_per_seq, + num_context_blocks=num_seqs * max_blocks_per_seq, + num_generation_blocks=0, + num_sequences=num_seqs, + num_blocks=None, + hp_pool=hp_pool, + v_scale_pool=kv_cache_manager.get_mla_v_scale_pool(), + generation_kv_lens=torch.empty(num_seqs, dtype=torch.int32, device=device), + generation_append_lens=torch.empty(num_seqs, dtype=torch.int32, device=device), + generation_lengths_num_tokens=-1, + generation_lengths_num_seqs=-1, + generation_lengths_num_contexts=-1, + generation_lengths_capture_recorded=False, + q_global_scale=q_global_scale, + kv_global_scale=kv_global_scale, + ), ) @@ -460,13 +468,15 @@ def _materialize_reference_cache_storage(metadata, layer_idx: int, head_dim: int storage_head_dim = kv_cache.shape[-1] * 2 static_global_scale = float(_get_fp4_mla_global_scale(metadata, kv_cache.device).item()) num_generation_sequences = metadata.num_seqs - metadata.num_contexts - page_stride = metadata.fp4_mla_page_table_stride - page_rows = metadata.paged_kv_indices.view( + page_stride = metadata.fp4_mla_state.page_table_stride + page_rows = metadata.fp4_mla_state.paged_kv_indices.view( metadata.num_seqs, page_stride, ) generation_page_rows = page_rows[metadata.num_contexts : metadata.num_seqs] - indptr = metadata.paged_kv_indptr_decode[: num_generation_sequences + 1].cpu().tolist() + indptr = ( + metadata.fp4_mla_state.paged_kv_indptr_decode[: num_generation_sequences + 1].cpu().tolist() + ) expected_indptr = [seq_idx * page_stride for seq_idx in range(num_generation_sequences + 1)] assert indptr == expected_indptr assert generation_page_rows.shape == (num_generation_sequences, page_stride) @@ -523,9 +533,9 @@ def _materialize_reference_cache_tokens( sf_cache = sf_cache.view(torch.float8_e4m3fn) storage_head_dim = kv_cache.shape[-1] * 2 static_global_scale = float(_get_fp4_mla_global_scale(metadata, kv_cache.device).item()) - page_rows = metadata.paged_kv_indices.view( + page_rows = metadata.fp4_mla_state.paged_kv_indices.view( metadata.num_seqs, - metadata.fp4_mla_page_table_stride, + metadata.fp4_mla_state.page_table_stride, ) dequantized_pages = {} tokens = [] @@ -595,11 +605,11 @@ def _build_fp4_mla_attention_decode_case( seq_lens=seq_lens, page_size=page_size, ) - assert metadata.fp4_mla_v_scale_pool is not None + assert metadata.fp4_mla_state.v_scale_pool is not None persistent_pool_base = kv_cache_manager.get_mla_v_packed_pool_base() if persistent_pool_base is not None: persistent_pool_base.zero_() - metadata.fp4_mla_v_scale_pool.zero_() + metadata.fp4_mla_state.v_scale_pool.zero_() metadata.kv_lens_cuda_runtime = torch.tensor( context_seq_lens, @@ -608,13 +618,13 @@ def _build_fp4_mla_attention_decode_case( ) metadata.prompt_lens_cuda_runtime = metadata.kv_lens_cuda_runtime.clone() metadata.prompt_lens_cpu_runtime = torch.tensor(context_seq_lens, dtype=torch.int32) - metadata.batch_indices = torch.cat( + metadata.fp4_mla_state.batch_indices = torch.cat( [ torch.full((seq_len,), seq_idx, dtype=torch.int32, device=device) for seq_idx, seq_len in enumerate(context_seq_lens) ] ) - metadata.positions = torch.cat( + metadata.fp4_mla_state.positions = torch.cat( [torch.arange(seq_len, dtype=torch.int32, device=device) for seq_len in context_seq_lens] ) context_latent = ( @@ -632,8 +642,10 @@ def _build_fp4_mla_attention_decode_case( torch.cuda.synchronize() metadata.num_contexts = 0 - metadata.num_context_blocks = 0 - metadata.num_generation_blocks = len(seq_lens) * metadata.fp4_mla_page_table_stride + metadata.fp4_mla_state.num_context_blocks = 0 + metadata.fp4_mla_state.num_generation_blocks = ( + len(seq_lens) * metadata.fp4_mla_state.page_table_stride + ) metadata.kv_lens_cuda_runtime = torch.tensor(seq_lens, dtype=torch.int32, device=device) metadata.prompt_lens_cuda_runtime = torch.full( (len(seq_lens),), query_len_per_seq, dtype=torch.int32, device=device @@ -642,12 +654,12 @@ def _build_fp4_mla_attention_decode_case( (len(seq_lens),), query_len_per_seq, dtype=torch.int32 ) num_queries = len(seq_lens) * query_len_per_seq - metadata.batch_indices = torch.arange( + metadata.fp4_mla_state.batch_indices = torch.arange( len(seq_lens), dtype=torch.int32, device=device, ).repeat_interleave(query_len_per_seq) - metadata.positions = torch.cat( + metadata.fp4_mla_state.positions = torch.cat( [ torch.arange(context_len, seq_len, dtype=torch.int32, device=device) for context_len, seq_len in zip(context_seq_lens, seq_lens) @@ -701,16 +713,16 @@ def _build_fp4_mla_attention_decode_case( q_quant_input=q_quant_input, ) torch.cuda.synchronize() - assert metadata._fp4_mla_prequantized_q is not None - assert metadata._fp4_mla_prequantized_q_sf is not None - assert metadata._fp4_mla_q_batch_capacity == num_queries + assert metadata.fp4_mla_state.prequantized_q is not None + assert metadata.fp4_mla_state.prequantized_q_sf is not None + assert metadata.fp4_mla_state.q_batch_capacity == num_queries torch.testing.assert_close(q_rope_out, q_pe, rtol=0, atol=0) canonical_generation = _materialize_reference_cache_tokens( metadata, layer_idx=0, - batch_indices=metadata.batch_indices, - positions=metadata.positions, + batch_indices=metadata.fp4_mla_state.batch_indices, + positions=metadata.fp4_mla_state.positions, head_dim=head_dim, ) torch.testing.assert_close( @@ -722,7 +734,7 @@ def _build_fp4_mla_attention_decode_case( ) # Decode reference exercises the quantized cache without an HP overlay. - metadata.high_precision_kv_pool.zero_() + metadata.fp4_mla_state.hp_pool.zero_() torch.cuda.synchronize() return kv_cache_manager, metadata, q_quant_input, kv_lora_rank, qk_rope_head_dim @@ -745,7 +757,7 @@ def _fp4_mla_attention_decode_reference( dequant_k_residual = storage[..., head_dim : head_dim + FP4_MLA_K_RESIDUAL_DIM] num_heads = q_nope.shape[1] q_full = torch.cat((q_nope, q_pe), dim=-1).reshape(-1, head_dim) - global_scale = metadata._fp4_mla_q_global_scale + global_scale = metadata.fp4_mla_state.q_global_scale q_fp4, q_sf = torch.ops.trtllm.fp4_quantize_with_residual( q_full, global_scale, @@ -771,7 +783,7 @@ def _fp4_mla_attention_decode_reference( global_scale=FP4_MLA_P_GLOBAL_SCALE, ) - indptr = metadata.paged_kv_indptr_decode.cpu().tolist() + indptr = metadata.fp4_mla_state.paged_kv_indptr_decode.cpu().tolist() num_seqs = metadata.num_seqs - metadata.num_contexts indptr = indptr[: num_seqs + 1] kv_lens = ( @@ -787,7 +799,7 @@ def _fp4_mla_attention_decode_reference( for seq_idx in range(num_seqs): kv_len = kv_lens[seq_idx] page_count = indptr[seq_idx + 1] - indptr[seq_idx] - assert page_count == metadata.fp4_mla_page_table_stride + assert page_count == metadata.fp4_mla_state.page_table_stride assert page_count * metadata.page_size >= kv_len full_cache = dequant_cache[indptr[seq_idx] : indptr[seq_idx + 1]].reshape(-1, head_dim) assert full_cache.shape[0] >= kv_len @@ -880,9 +892,9 @@ def _assert_fp4_mla_attention_decode_accuracy( sm_scale=sm_scale, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, - prequantized_q=metadata._fp4_mla_prequantized_q, - prequantized_q_sf=metadata._fp4_mla_prequantized_q_sf, - q_batch_capacity=metadata._fp4_mla_q_batch_capacity, + prequantized_q=metadata.fp4_mla_state.prequantized_q, + prequantized_q_sf=metadata.fp4_mla_state.prequantized_q_sf, + q_batch_capacity=metadata.fp4_mla_state.q_batch_capacity, ) torch.cuda.synchronize() @@ -1069,9 +1081,9 @@ def test_fp4_mla_context_tail_uses_draft_slack_ring( ) torch.cuda.synchronize() - hp_page = int(metadata._fp4_mla_hp_page_indices[0].item()) - hp_ring = metadata.high_precision_kv_pool.view( - metadata.high_precision_kv_pool.shape[0], 1, 1, ring_size, head_dim + hp_page = int(metadata.fp4_mla_state.hp_page_indices[0].item()) + hp_ring = metadata.fp4_mla_state.hp_pool.view( + metadata.fp4_mla_state.hp_pool.shape[0], 1, 1, ring_size, head_dim ) torch.testing.assert_close( hp_ring[hp_page, 0, 0, expected_slot], From 17fe5052171775e99771853af23735b660fe15f9 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Mon, 14 Sep 2026 10:54:38 -0700 Subject: [PATCH 16/21] refactor: streamline FP4 MLA runtime helpers Split FP4 MLA launch helpers into focused modules and centralize CuTeDSL compilation without debug stdout filtering. Move validation into support selection with matching cache keys, consolidate state updates, and reuse FP8 context scratch with normal FMHA selection. Remove redundant guards and update existing test monkeypatch targets without adding test cases. Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 9 +- .../_torch/attention/backends/fmha/fp4_mla.py | 36 +- .../_torch/attention/backends/fmha/manager.py | 30 +- .../attention/backends/fp4_mla/__init__.py | 3961 +---------------- .../backends/fp4_mla/cache_manager.py | 4 +- .../backends/fp4_mla/cache_update.py | 1087 +++++ .../attention/backends/fp4_mla/config.py | 242 + .../backends/fp4_mla/cute_dsl_utils.py | 21 + .../attention/backends/fp4_mla/decode.py | 1130 +++++ .../backends/fp4_mla/fp4_mla_context.py | 12 +- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 58 +- ...p4_mla_cutedsl_mufu16_fused_v_transpose.py | 58 +- .../fp4_mla/fp4_mla_cutedsl_v_repack.py | 64 +- .../attention/backends/fp4_mla/layout.py | 257 ++ .../attention/backends/fp4_mla/metadata.py | 897 ++++ .../attention/backends/fp4_mla/state.py | 11 +- .../attention/backends/fp4_mla/v_cache.py | 409 ++ .../_torch/attention/backends/trtllm.py | 6 - .../_torch/pyexecutor/resource_manager.py | 7 +- .../unittest/_torch/attention/test_fp4_mla.py | 6 +- 20 files changed, 4278 insertions(+), 4027 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/config.py create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/cute_dsl_utils.py create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index cabf6be2101f..d283b3511eb1 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -431,8 +431,9 @@ The FMHA package is split by role: - `fmha/fp4_mla.py` implements FP4 MLA context and no-dequant decode. The shared FMHA availability check admits FP4 MLA only to libraries that declare `supports_fp4_mla`; it does not restrict the existing FP4 GQA path. - Selection caches model-invariant and cache-key-covered validation. Dynamic - sparse/sinks inputs and prepared-state checks still run on cache hits. + Selection caches validation in `_is_supported()`. Its cache key distinguishes + K/V input presence, sparse/sinks inputs, and prepared FP4 state, so changing + those conditions triggers validation again instead of reusing a valid entry. The core implementation requires dense TRTLLM MLA, BF16 absorption weights, fused RoPE with duplicated rotary tables, and KV Cache Manager V2. It uses FP8 context attention with an FP4 cache update and FP4 generation attention. @@ -447,6 +448,10 @@ The FMHA package is split by role: state reference and forwards prepare/MTP lifecycle updates to it. These buffers continue to use the metadata allocator for CUDA-graph address stability; the layer-local FP8 attention view is allocated lazily by FMHA. + The `fp4_mla` package exports entry points from focused `config`, `layout`, + `metadata`, `cache_update`, `v_cache`, and `decode` modules. The FP8 context + view has its own FMHA manager and uses normal capability-based selection; + its manager-owned scratch is allocated once and reused across layers. - `fmha/triton_custom_mask.py` implements the Triton custom-mask context phase. Custom-mask data applies to context requests; for mixed batches, `TrtllmAttention` can pair it with a later causal-generation provider through diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index 1774b591d35a..3ce18d0792d7 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -70,9 +70,9 @@ def _is_supported( *, phase: Optional[FmhaPhase] = None, ) -> bool: - # Masks/output formats are in the selection cache key; cache geometry - # and sparse compression are model invariants. Live inputs stay in forward. - del q, k, v, phase + # Input presence/readiness and mask/output formats are in the selection + # cache key; cache geometry and sparse compression are model invariants. + del q, phase if forward_args.output_sf is not None: raise NotImplementedError("FP4 MLA does not support quantized attention output.") if forward_args.attention_mask != PredefinedAttentionMask.CAUSAL: @@ -92,18 +92,6 @@ def _is_supported( raise RuntimeError("FP4 MLA requires a SELF-K-only KV cache.") if metadata.beam_width != 1: raise NotImplementedError("FP4 MLA does not support beam search.") - return True - - def forward( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - metadata: "TrtllmAttentionMetadata", - forward_args: AttentionForwardArgs, - ) -> None: - # These inputs/readiness conditions are not part of the FMHA cache - # key. Keep checking them even when selection reuses a cached library. if forward_args.attention_sinks is not None: raise NotImplementedError("FP4 MLA does not support attention sinks.") sparse_inputs = forward_args.sparse_runtime_params @@ -135,7 +123,7 @@ def forward( raise RuntimeError("FP4 MLA generation expects a fused query input.") else: raise NotImplementedError("FP4 MLA requires a context-only or generation-only call.") - super().forward(q, k, v, metadata, forward_args) + return True def run_mla_context(self, params: FmhaParams) -> None: attn = params.attn @@ -163,10 +151,6 @@ def run_mla_context(self, params: FmhaParams) -> None: raise RuntimeError("FP4 MLA context Q/K/V token counts do not match.") require_fp4_mla_fp8_context_support() - if metadata.is_cuda_graph: - raise NotImplementedError( - "FP4 MLA context does not support CUDA graphs with TRT-LLM FP8 FMHA." - ) num_tokens = q.shape[0] output = output.view(num_tokens, -1) @@ -174,7 +158,6 @@ def run_mla_context(self, params: FmhaParams) -> None: kv_lora_rank = attn.kv_lora_rank or 0 qk_rope_head_dim = attn.qk_rope_head_dim or 0 - attn._ensure_rope_table_size(metadata.max_seq_len) latent_cache = forward_args.latent_cache[:num_tokens] def update_fp4_cache() -> None: @@ -196,17 +179,18 @@ def update_fp4_cache() -> None: raise RuntimeError("FP8 MLA context scratch requires a KV cache manager.") scratch = getattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, None) scratch_head_dim = kv_lora_rank + qk_rope_head_dim - if not isinstance(scratch, _Fp8MlaContextScratch) or not scratch.matches( - metadata, - device=q.device, - head_dim=scratch_head_dim, - ): + if scratch is None: scratch = _Fp8MlaContextScratch.create( metadata, device=q.device, head_dim=scratch_head_dim, ) setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch) + else: + assert isinstance(scratch, _Fp8MlaContextScratch) + assert scratch.matches(metadata, device=q.device, head_dim=scratch_head_dim), ( + "FP8 MLA context scratch geometry must be shared by all layers of its KV manager." + ) fp8_attention = self._fp8_attention if fp8_attention is None: diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index c9afff66e1d8..0f2a014c53de 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -135,6 +135,12 @@ class _FmhaCacheKey(NamedTuple): attention_mask_type: AttentionMaskType use_spec_decoding: bool has_block_sparse_inputs: bool + # These support conditions can change without changing the query grid. + has_k_input: bool + has_v_input: bool + has_attention_sinks: bool + has_sparse_indices: bool + fp4_mla_state_ready: bool # LoRA can change the effective output from packed NVFP4 to unpacked BF16 # without changing the request shape. Keep those selection regimes apart. output_dtype: torch.dtype | None @@ -317,6 +323,9 @@ def _make_cache_key( q: torch.Tensor, metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, + *, + k: torch.Tensor | None = None, + v: torch.Tensor | None = None, ) -> _FmhaCacheKey: """Build the dynamic FMHA cache key for one attention instance. @@ -367,14 +376,29 @@ def _make_cache_key( generation_seq_len_q, _FMHA_CACHE_SEQ_LEN_Q_GRID ) - block_sparse_inputs = forward_args.sparse_runtime_params.block_sparse_inputs + sparse = forward_args.sparse_runtime_params + has_sparse_indices = ( + metadata.num_sparse_topk > 0 + or (sparse.sparse_kv_indices is not None and sparse.sparse_kv_indices.numel() > 0) + or (sparse.sparse_attn_indices is not None and sparse.sparse_attn_indices.numel() > 0) + ) + fp4_state = getattr(metadata, "fp4_mla_state", None) return _FmhaCacheKey( context_batch_size=context_batch_size, generation_batch_size=generation_batch_size, generation_seq_len_q=generation_seq_len_q, attention_mask_type=attention_mask_type, use_spec_decoding=metadata.use_spec_decoding, - has_block_sparse_inputs=block_sparse_inputs is not None, + has_block_sparse_inputs=sparse.block_sparse_inputs is not None, + has_k_input=k is not None, + has_v_input=v is not None, + has_attention_sinks=forward_args.attention_sinks is not None, + has_sparse_indices=has_sparse_indices, + fp4_mla_state_ready=( + fp4_state is not None + and fp4_state.hp_pool is not None + and fp4_state.v_scale_pool is not None + ), output_dtype=output_dtype, output_sf_dtype=output_sf_dtype, ) @@ -392,7 +416,7 @@ def select( if not _is_fmha_cache_enabled(): return self._select_uncached(attn, q, k, v, metadata, forward_args) - cache_key = self._make_cache_key(q, metadata, forward_args) + cache_key = self._make_cache_key(q, metadata, forward_args, k=k, v=v) fmha = self._cache.get(cache_key) if fmha is not None: if self._cache_sanity_check_enabled: diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py index 7c81954f85c1..ccfde0c4fb15 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -1,3799 +1,176 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared MLA FP4 KV-cache helpers. +"""FP4 MLA entry points; implementation lives in the focused sibling modules.""" -The high-precision (HP) BF16 KV pool is a standalone circular buffer used -alongside the paged FP4 KV pool when MLA models run with NVFP4 KV cache. -The pool retains one 16-token quantization tile plus speculative-rewind slack -per sequence. - -Used by the TRTLLM attention backend FP4 MLA FMHA path. -""" - -import importlib.util -import os -from typing import Any, Literal, Optional - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._utils import get_sm_version - -from .fp4_mla_kernels import ( - _fp4_mla_context_cache_update_kernel, - _fp4_mla_generation_fused_qk_rope_cache_update_kernel, +from .cache_update import ( + _get_fp4_mla_context_start_positions as _get_fp4_mla_context_start_positions, ) - -HP_BLOCK_SIZE: int = 16 -FP4_BLOCK_SIZE: int = 16 -FP4_MLA_TOKENS_PER_BLOCK: int = 128 -FP4_MLA_SCALE_ROW_GROUP: int = 128 -FP4_MLA_SCALE_COL_GROUP: int = 4 -FP4_MLA_P_GLOBAL_SCALE: float = 448.0 * 6.0 -FP4_MLA_Q_STATIC_AMAX: float = 400.0 -FP4_MLA_KV_STATIC_AMAX: float = 30.0 -FP4_MLA_Q_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_Q_STATIC_AMAX -FP4_MLA_KV_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_KV_STATIC_AMAX -# Max finite e4m3 magnitude for FP4 MLA block-scale clamping. -FP4_MLA_E4M3_MAX: float = 448.0 -FP4_MLA_Q_RESIDUAL_DIM: int = 64 -FP4_MLA_K_RESIDUAL_DIM: int = FP4_MLA_Q_RESIDUAL_DIM -FP4_MLA_Q_PREFIX_DIM: int = 512 -FP4_MLA_Q_PREFIX_BLOCK_DIM: int = 256 -FP4_MLA_Q1_PREFIX_BLOCK_DIM: int = 512 -FP4_MLA_Q_LOGICAL_DIM: int = FP4_MLA_Q_PREFIX_DIM + 2 * FP4_MLA_Q_RESIDUAL_DIM -FP4_MLA_Q_PACKED_DIM: int = FP4_MLA_Q_LOGICAL_DIM // 2 -FP4_MLA_Q_SF_GROUPS: int = FP4_MLA_Q_LOGICAL_DIM // FP4_BLOCK_SIZE -FP4_MLA_ATTENTION_BACKEND_ENV = "TRTLLM_FP4_MLA_ATTENTION_BACKEND" -FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV = "TRTLLM_FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE" -_FP4_MLA_CUTEDSL_BACKEND = "cutedsl" -_FP4_MLA_K_RESIDUAL_BACKENDS = ("triton", _FP4_MLA_CUTEDSL_BACKEND) -_FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH = 4 -_FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH = 16 -_FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH = 32 -_FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD = 256 -_FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD = 512 -_FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD = 640 -_FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD = 768 -_HPUpdatePhase = Literal["context", "generation"] -_FP4_MLA_TRITON_PRELOAD_KEYS = "_fp4_mla_triton_preload_keys" -_FP4_MLA_PAGE_TABLE_TILE_SIZE = 128 -_FP4_MLA_MAX_GRID_Z = 65_535 - - -# Environment helpers - - -def _env_enabled_default(name: str, default: bool) -> bool: - value = os.getenv(name) - if value is None or value == "": - return default - return value.lower() in ( - "1", - "true", - "yes", - "on", - ) - - -def _fp4_mla_cutedsl_fused_v_transpose_enabled() -> bool: - return _env_enabled_default(FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, False) - - -def _env_int(name: str) -> Optional[int]: - value = os.environ.get(name) - if value is None or value == "": - return None - return int(value) - - -def _fp4_mla_attention_backend() -> str: - backend = os.getenv(FP4_MLA_ATTENTION_BACKEND_ENV) - if backend: - return backend.lower() - return _FP4_MLA_CUTEDSL_BACKEND if get_sm_version() == 107 else "triton" - - -def _cutedsl_backend_available() -> bool: - try: - return all( - importlib.util.find_spec(module) is not None - for module in ("ctm", "cutlass", "cuda.bindings.driver") - ) - except ModuleNotFoundError: - return False - - -def _fp4_mla_cutedsl_kernel_module() -> Any: - if _fp4_mla_cutedsl_fused_v_transpose_enabled(): - from . import fp4_mla_cutedsl_mufu16_fused_v_transpose - - return fp4_mla_cutedsl_mufu16_fused_v_transpose - - from . import fp4_mla_cutedsl_mufu16 - - return fp4_mla_cutedsl_mufu16 - - -def _ceil_div(lhs: int, rhs: int) -> int: - return (lhs + rhs - 1) // rhs - - -def _fp4_mla_q1_kv_blocks_per_program(num_gen: int, v_head_dim: int) -> int: - """Select the Q1 KV work per program without changing the launch boundary.""" - large_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH * FP4_BLOCK_SIZE - if num_gen >= _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD and v_head_dim % large_batch_block_dim == 0: - return _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH - medium_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH * FP4_BLOCK_SIZE - if ( - num_gen >= _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD - and v_head_dim % medium_batch_block_dim == 0 - ): - return _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH - small_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH * FP4_BLOCK_SIZE - if v_head_dim % small_batch_block_dim == 0: - return _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH - return 1 - - -def _fp4_mla_q1_prefix_blocks_per_program( - num_gen: int, - q1_kv_blocks_per_program: int, -) -> int: - """Select rolled Q-prefix work only when the batch keeps it efficient.""" - max_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM - if ( - num_gen >= _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD - and q1_kv_blocks_per_program == _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH - ): - if num_gen >= _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD: - return min(4, max_prefix_blocks) - return min(2, max_prefix_blocks) - return 1 - - -def _fp4_mla_q1_preload_variants( - max_num_sequences: int, - v_head_dim: int, -) -> tuple[tuple[int, int], ...]: - """Return every Q1 tuning variant reachable by the configured batch limit.""" - batch_sizes = [1] - for threshold in ( - _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD, - _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD, - _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD, - _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD, - ): - if threshold <= max_num_sequences: - batch_sizes.append(threshold) - - variants = [] - for batch_size in batch_sizes: - kv_blocks = _fp4_mla_q1_kv_blocks_per_program(batch_size, v_head_dim) - prefix_blocks = _fp4_mla_q1_prefix_blocks_per_program( - batch_size, - kv_blocks, - ) - variant = (kv_blocks, prefix_blocks) - if variant not in variants: - variants.append(variant) - return tuple(variants) - - -def _fp4_mla_triton_preload_key_set(metadata: Any) -> set[tuple[object, ...]]: - """Return the engine-scoped set of Triton variants loaded during warmup.""" - owner = getattr(metadata, "kv_cache_manager", None) - if owner is None: - owner = metadata - preload_keys = getattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, None) - if preload_keys is None: - preload_keys = set() - setattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, preload_keys) - return preload_keys - - -@triton.jit -def _fp4_mla_store_sequence_append_metadata( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - sequence_idx, - num_tokens, - PREFIX_BLOCK: tl.constexpr, - TOKEN_BLOCK: tl.constexpr, -): - append_len = tl.load(append_lens_ptr + sequence_idx) - prefix_offsets = tl.arange(0, PREFIX_BLOCK) - token_start = append_len - append_len - for prefix_start in tl.range(0, sequence_idx, PREFIX_BLOCK): - preceding_sequences = prefix_start + prefix_offsets - preceding_lens = tl.load( - append_lens_ptr + preceding_sequences, - mask=preceding_sequences < sequence_idx, - other=0, - ) - token_start += tl.sum(preceding_lens) - - cached_len = tl.load(kv_lens_ptr + sequence_idx) - append_len - - token_offsets = tl.arange(0, TOKEN_BLOCK) - for block_start in tl.range(0, append_len, TOKEN_BLOCK): - local_offsets = block_start + token_offsets - token_mask = (local_offsets < append_len) & (token_start + local_offsets < num_tokens) - output_offsets = token_start + local_offsets - tl.store( - batch_indices_ptr + output_offsets, - sequence_idx, - mask=token_mask, - ) - tl.store( - positions_ptr + output_offsets, - cached_len + local_offsets, - mask=token_mask, - ) - - -@triton.jit( - do_not_specialize=[ - "num_tokens", - "num_contexts", - "num_generation_sequences", - ], - do_not_specialize_on_alignment=[ - "num_tokens", - "num_contexts", - "num_generation_sequences", - ], +from .cache_update import _prepare_fp4_mla_q_buffers as _prepare_fp4_mla_q_buffers +from .cache_update import ( + _scatter_fp4_mla_kv_cache_2d_context as _scatter_fp4_mla_kv_cache_2d_context, ) -def _fp4_mla_append_metadata_kernel( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - num_tokens, - num_contexts, - num_generation_sequences, - ONE_TOKEN_GENERATION: tl.constexpr, - PREFIX_BLOCK: tl.constexpr, - TOKEN_BLOCK: tl.constexpr, - GENERATION_BLOCK: tl.constexpr, -): - program_idx = tl.program_id(0) - if ONE_TOKEN_GENERATION: - if program_idx < num_contexts: - _fp4_mla_store_sequence_append_metadata( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - program_idx, - num_tokens, - PREFIX_BLOCK, - TOKEN_BLOCK, - ) - else: - generation_offsets = (program_idx - num_contexts) * GENERATION_BLOCK + tl.arange( - 0, GENERATION_BLOCK - ) - generation_mask = generation_offsets < num_generation_sequences - sequence_indices = num_contexts + generation_offsets - output_offsets = num_tokens - num_generation_sequences + generation_offsets - generation_mask = generation_mask & (output_offsets < num_tokens) - generation_positions = ( - tl.load( - kv_lens_ptr + sequence_indices, - mask=generation_mask, - other=1, - ) - - 1 - ) - tl.store( - batch_indices_ptr + output_offsets, - sequence_indices, - mask=generation_mask, - ) - tl.store( - positions_ptr + output_offsets, - generation_positions, - mask=generation_mask, - ) - else: - _fp4_mla_store_sequence_append_metadata( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - program_idx, - num_tokens, - PREFIX_BLOCK, - TOKEN_BLOCK, - ) - - -def populate_fp4_mla_append_metadata( - append_lens: torch.Tensor, - kv_lens: torch.Tensor, - batch_indices: torch.Tensor, - positions: torch.Tensor, - *, - num_tokens: int, - num_sequences: int, - num_contexts: int, - num_context_tokens: int, -) -> None: - """Populate FP4 MLA token-to-sequence metadata in one Triton launch. - - Mixed batches vectorize their one-token generation rows. Multi-token MTP - and fallback shapes use the generic per-sequence path in the same kernel. - """ - if num_sequences <= 0 or num_tokens <= 0: - return - if not 0 <= num_contexts <= num_sequences: - raise ValueError( - f"FP4 MLA num_contexts must be in [0, {num_sequences}], got {num_contexts}." - ) - if not 0 <= num_context_tokens <= num_tokens: - raise ValueError( - f"FP4 MLA num_context_tokens must be in [0, {num_tokens}], got {num_context_tokens}." - ) - - tensors = ( - append_lens, - kv_lens, - batch_indices, - positions, - ) - if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): - raise ValueError("FP4 MLA append metadata tensors must be contiguous and one-dimensional.") - if any(tensor.dtype != torch.int32 for tensor in tensors): - raise TypeError("FP4 MLA append metadata tensors must use int32.") - if any(not tensor.is_cuda for tensor in tensors): - raise ValueError("FP4 MLA append metadata tensors must be CUDA tensors.") - if any(tensor.device != append_lens.device for tensor in tensors[1:]): - raise ValueError("FP4 MLA append metadata tensors must be on the same device.") - sequence_tensors = (append_lens, kv_lens) - if any(tensor.numel() < num_sequences for tensor in sequence_tensors): - raise ValueError( - f"FP4 MLA sequence metadata tensors need at least {num_sequences} entries." - ) - token_tensors = (batch_indices, positions) - if any(tensor.numel() < num_tokens for tensor in token_tensors): - raise ValueError(f"FP4 MLA token metadata tensors need at least {num_tokens} entries.") - - num_generation_sequences = num_sequences - num_contexts - # Each scheduled generation sequence appends at least one token. Equality - # therefore identifies the common mixed batch with one token per decode - # row without reading the device append lengths back on the host. - one_token_generation = num_tokens == num_context_tokens + num_generation_sequences - generation_block = 128 - grid = num_sequences - if one_token_generation: - grid = num_contexts + triton.cdiv(num_generation_sequences, generation_block) - - _fp4_mla_append_metadata_kernel[(grid,)]( - append_lens, - kv_lens, - batch_indices, - positions, - num_tokens, - num_contexts, - num_generation_sequences, - ONE_TOKEN_GENERATION=one_token_generation, - PREFIX_BLOCK=128, - TOKEN_BLOCK=256, - GENERATION_BLOCK=generation_block, - num_warps=4, - ) - - -@triton.jit( - do_not_specialize=["num_gen", "generation_len"], - do_not_specialize_on_alignment=["num_gen", "generation_len"], +from .cache_update import ( + _scatter_fp4_mla_kv_cache_2d_generation as _scatter_fp4_mla_kv_cache_2d_generation, ) -def _fp4_mla_generation_lengths_kernel( - kv_lens_ptr, - prompt_lens_ptr, - corrected_kv_lens_ptr, - generation_lens_ptr, - num_gen, - generation_len, - BLOCK: tl.constexpr, -): - offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < num_gen - kv_lens = tl.load(kv_lens_ptr + offsets, mask=mask, other=0) - prompt_lens = tl.load(prompt_lens_ptr + offsets, mask=mask, other=0) - tl.store( - corrected_kv_lens_ptr + offsets, - kv_lens - prompt_lens + generation_len, - mask=mask, - ) - tl.store(generation_lens_ptr + offsets, generation_len, mask=mask) - - -def populate_fp4_mla_generation_lengths( - kv_lens: torch.Tensor, - prompt_lens: torch.Tensor, - corrected_kv_lens: torch.Tensor, - generation_lens: torch.Tensor, - *, - num_gen_tokens: int, - num_gen: int, -) -> None: - """Populate reusable FP4 MLA generation lengths in one Triton launch.""" - if num_gen <= 0 or num_gen_tokens % num_gen != 0: - raise ValueError( - "FP4 MLA generation lengths require a positive sequence count and " - f"uniform token count, got {num_gen_tokens} tokens for {num_gen} sequences." - ) - tensors = (kv_lens, prompt_lens, corrected_kv_lens, generation_lens) - if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): - raise ValueError( - "FP4 MLA generation length tensors must be contiguous and one-dimensional." - ) - if any(tensor.dtype != torch.int32 for tensor in tensors): - raise TypeError("FP4 MLA generation length tensors must use int32.") - if any(not tensor.is_cuda for tensor in tensors): - raise ValueError("FP4 MLA generation length tensors must be CUDA tensors.") - if any(tensor.device != kv_lens.device for tensor in tensors[1:]): - raise ValueError("FP4 MLA generation length tensors must be on the same device.") - if any(tensor.numel() < num_gen for tensor in tensors): - raise ValueError(f"FP4 MLA generation length tensors need at least {num_gen} entries.") - - block = 128 - _fp4_mla_generation_lengths_kernel[(triton.cdiv(num_gen, block),)]( - kv_lens, - prompt_lens, - corrected_kv_lens, - generation_lens, - num_gen, - num_gen_tokens // num_gen, - BLOCK=block, - num_warps=4, - ) - - -def _fp4_mla_page_table_spec(kv_cache_manager: Any) -> Any: - get_spec = getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None) - if not callable(get_spec): - raise RuntimeError("FP4 MLA requires Fp4MlaKVCacheManagerV2 page metadata.") - spec = get_spec() - for field_name in ( - "cache_pool_id", - "cache_page_index_scale", - "hp_pool_id", - "hp_page_index_scale", - ): - value = getattr(spec, field_name, None) - if not isinstance(value, int) or value < 0: - raise ValueError( - f"FP4 MLA page-table spec requires non-negative {field_name}, got {value}." - ) - if spec.cache_page_index_scale <= 0 or spec.hp_page_index_scale <= 0: - raise ValueError("FP4 MLA page-index scales must be positive.") - return spec - - -# Mixed batches frequently vary by one sequence. Keep per-forward dimensions -# out of Triton's specialization key and tile page rows at one fixed width so -# those shape changes cannot trigger JIT compilation on the critical path. -@triton.jit( - do_not_specialize=["num_sequences", "num_contexts", "max_pages"], - do_not_specialize_on_alignment=["num_sequences", "num_contexts", "max_pages"], +from .cache_update import _validate_fp4_mla_context_rope as _validate_fp4_mla_context_rope +from .cache_update import ( + _validate_fp4_mla_context_start_alignment as _validate_fp4_mla_context_start_alignment, ) -def _fp4_mla_materialize_page_table_kernel( - page_ids_ptr, - paged_kv_indptr_ptr, - paged_kv_indptr_decode_ptr, - block_offsets_ptr, - kv_lens_ptr, - generation_kv_lens_ptr, - block_offsets_stride, - num_sequences, - num_contexts, - max_pages, - PAGE_SIZE: tl.constexpr, - PAGE_INDEX_SCALE: tl.constexpr, - PAGE_TILE_SIZE: tl.constexpr, -): - sequence_idx = tl.program_id(0) - page_tile_idx = tl.program_id(1) - page_offsets = page_tile_idx * PAGE_TILE_SIZE + tl.arange(0, PAGE_TILE_SIZE) - generation_idx = sequence_idx - num_contexts - is_generation = generation_idx >= 0 - context_kv_len = tl.load(kv_lens_ptr + sequence_idx) - generation_kv_len = tl.load( - generation_kv_lens_ptr + generation_idx, - mask=is_generation, - other=0, - ) - kv_len = tl.maximum(tl.where(is_generation, generation_kv_len, context_kv_len), 0) - num_active_pages = tl.minimum( - (kv_len + PAGE_SIZE - 1) // PAGE_SIZE, - max_pages, - ) - active_page_mask = page_offsets < num_active_pages - encoded_page_offsets = tl.load( - block_offsets_ptr + sequence_idx * block_offsets_stride + page_offsets, - mask=active_page_mask, - other=-1, - ) - decoded_page_ids = tl.where( - encoded_page_offsets >= 0, - encoded_page_offsets // PAGE_INDEX_SCALE, - encoded_page_offsets, - ) - page_ids = tl.where(active_page_mask, decoded_page_ids, 0) - table_offset = sequence_idx * max_pages + page_offsets - # Fixed-stride indptrs expose the whole row. Initialize inactive slots so - # masked or prefetched page-table reads cannot observe stale page IDs. - tl.store( - page_ids_ptr + table_offset, - page_ids, - mask=page_offsets < max_pages, - ) - - first_lane = page_offsets == 0 - sequence_start = sequence_idx * max_pages - tl.store( - paged_kv_indptr_ptr + sequence_idx + page_offsets, - sequence_start, - mask=first_lane, - ) - tl.store( - paged_kv_indptr_decode_ptr + generation_idx + page_offsets, - generation_idx * max_pages, - mask=first_lane & is_generation, - ) - final_sequence = sequence_idx == num_sequences - 1 - table_end = num_sequences * max_pages - num_generation_sequences = num_sequences - num_contexts - tl.store( - paged_kv_indptr_ptr + num_sequences + page_offsets, - table_end, - mask=first_lane & final_sequence, - ) - tl.store( - paged_kv_indptr_decode_ptr + num_generation_sequences + page_offsets, - num_generation_sequences * max_pages, - mask=first_lane & final_sequence, - ) - - -def configure_fp4_mla_device_page_table( - metadata: Any, - kv_lens: Optional[torch.Tensor] = None, -) -> bool: - """Configure the fixed-stride, device-materialized page table. - - Context, generation, and fresh mixed batches receive the full block-offset - table on the GPU. The materialization kernel decodes V2 page indices and - refreshes rows from the final device KV lengths before cache update. - """ - metadata.fp4_mla_state.device_page_table = False - metadata.fp4_mla_state.device_page_table_valid = False - metadata.fp4_mla_state.page_table_stride = 0 - metadata.fp4_mla_state.context_repack_max_touched_pages = 1 - - kv_cache_manager = getattr(metadata, "kv_cache_manager", None) - num_contexts = int(getattr(metadata, "num_contexts", 0)) - num_sequences = int(getattr(metadata, "num_seqs", 0)) - num_generation_sequences = num_sequences - num_contexts - num_tokens = int(getattr(metadata, "num_tokens", 0)) - num_context_tokens = int(getattr(metadata, "num_ctx_tokens", 0)) - num_generation_tokens = num_tokens - num_context_tokens - block_offsets = getattr(metadata, "kv_cache_block_offsets", None) - page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) - paged_kv_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) - paged_kv_indptr_decode = getattr(metadata.fp4_mla_state, "paged_kv_indptr_decode", None) - max_page_capacity = int(getattr(kv_cache_manager, "max_blocks_per_seq", 0) or 0) - page_spec = _fp4_mla_page_table_spec(kv_cache_manager) - page_index_scale = int(page_spec.cache_page_index_scale) - - tensors = (block_offsets, page_ids, paged_kv_indptr, paged_kv_indptr_decode) - is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) - generation_only = num_contexts == 0 - fresh_mixed = ( - not is_cuda_graph - and num_contexts > 0 - and num_generation_sequences > 0 - and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 - ) - fresh_context_only = ( - not is_cuda_graph - and num_contexts > 0 - and num_generation_sequences == 0 - and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 - ) - has_valid_generation = num_generation_sequences == 0 or ( - num_generation_tokens >= num_generation_sequences - and num_generation_tokens % num_generation_sequences == 0 - ) - # NVFP4 exposes one data pool plus its paired block-scale pool. The - # materializer reads encoded data offsets from pool 0. - supported = ( - (generation_only or fresh_mixed or fresh_context_only) - and kv_cache_manager is not None - and has_valid_generation - and int(getattr(metadata, "beam_width", 1)) == 1 - and not bool(getattr(metadata, "is_spec_dec_tree", False)) - and not bool(getattr(metadata, "locality_domain_enabled", False)) - and not bool(getattr(metadata, "enable_helix", False)) - and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK - and max_page_capacity > 0 - and page_index_scale > 0 - and all(isinstance(tensor, torch.Tensor) for tensor in tensors) - and all(tensor.dtype == torch.int32 for tensor in tensors) - and all(tensor.is_cuda for tensor in tensors) - ) - if not supported: - return False - - hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) - max_pool_id = max(page_spec.cache_pool_id, page_spec.hp_pool_id) - if ( - not isinstance(hp_page_ids, torch.Tensor) - or hp_page_ids.dtype != torch.int32 - or not hp_page_ids.is_cuda - or block_offsets.shape[0] <= max_pool_id - ): - return False - metadata.fp4_mla_state.cache_pool_id = int(page_spec.cache_pool_id) - metadata.fp4_mla_state.cache_page_index_scale = int(page_spec.cache_page_index_scale) - metadata.fp4_mla_state.hp_pool_id = int(page_spec.hp_pool_id) - metadata.fp4_mla_state.hp_page_index_scale = int(page_spec.hp_page_index_scale) - - max_pages = max_page_capacity - host_kv_lens_available = ( - isinstance(kv_lens, torch.Tensor) - and kv_lens.device.type == "cpu" - and kv_lens.ndim == 1 - and kv_lens.numel() >= num_sequences - ) - if fresh_mixed and not host_kv_lens_available: - return False - if not is_cuda_graph and host_kv_lens_available: - generation_tokens_per_sequence = ( - num_generation_tokens // num_generation_sequences if num_generation_sequences > 0 else 0 - ) - # Eager execution can narrow the fixed row stride to the current - # batch. CUDA Graph metadata retains full configured capacity so a - # replay never changes tensor addresses or launch dimensions. - max_kv_len = int(kv_lens[:num_sequences].max().item()) + max( - 0, - generation_tokens_per_sequence - 1, - ) - max_pages = min( - max_page_capacity, - max(1, _ceil_div(max_kv_len, FP4_MLA_TOKENS_PER_BLOCK)), - ) - if num_contexts > 0: - max_context_len = int(kv_lens[:num_contexts].max().item()) - max_context_pages = _ceil_div( - max_context_len, - FP4_MLA_TOKENS_PER_BLOCK, - ) - metadata.fp4_mla_state.context_repack_max_touched_pages = min( - max_pages, - triton.next_power_of_2(max(1, max_context_pages)), - ) - - assert isinstance(block_offsets, torch.Tensor) - assert isinstance(page_ids, torch.Tensor) - assert isinstance(paged_kv_indptr, torch.Tensor) - assert isinstance(paged_kv_indptr_decode, torch.Tensor) - required_page_ids = num_sequences * max_pages - buffers_cover_table = ( - block_offsets.ndim == 4 - and block_offsets.shape[0] >= 1 - and block_offsets.shape[1] >= num_sequences - and block_offsets.shape[2] >= 1 - and block_offsets.shape[3] >= max_pages - and page_ids.ndim == 1 - and page_ids.numel() >= required_page_ids - and paged_kv_indptr.ndim == 1 - and paged_kv_indptr.numel() >= num_sequences + 1 - and paged_kv_indptr_decode.ndim == 1 - and paged_kv_indptr_decode.numel() >= num_generation_sequences + 1 - ) - if not buffers_cover_table: - return False - if metadata.fp4_mla_state.hp_page_indices.numel() < required_page_ids: - return False - - metadata.fp4_mla_state.device_page_table = True - metadata.fp4_mla_state.num_sequences = num_sequences - metadata.fp4_mla_state.page_table_stride = max_pages - metadata.fp4_mla_state.num_blocks = None - metadata.fp4_mla_state.num_context_blocks = num_contexts * max_pages - metadata.fp4_mla_state.num_generation_blocks = num_generation_sequences * max_pages - return True - - -def materialize_fp4_mla_device_page_table( - metadata: Any, - kv_lens: torch.Tensor, - generation_kv_lens: Optional[torch.Tensor] = None, -) -> None: - """Refresh the fixed-stride context and generation page table once per forward.""" - if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): - raise RuntimeError("FP4 MLA requires fixed-stride device page metadata.") - if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): - return - - num_contexts = int(metadata.num_contexts) - num_sequences = int(metadata.num_seqs) - num_generation_sequences = num_sequences - num_contexts - max_pages = int(metadata.fp4_mla_state.page_table_stride) - if num_sequences <= 0 or max_pages <= 0: - raise RuntimeError( - "FP4 MLA device page metadata requires positive sequence and page capacities." - ) - if ( - kv_lens.dtype != torch.int32 - or not kv_lens.is_cuda - or kv_lens.ndim != 1 - or kv_lens.stride(0) != 1 - or kv_lens.numel() < num_sequences - ): - raise ValueError( - "FP4 MLA device page metadata requires a contiguous CUDA int32 " - f"KV-length tensor with at least {num_sequences} entries." - ) - - if generation_kv_lens is None: - generation_kv_lens = kv_lens[num_contexts:num_sequences] - if ( - generation_kv_lens.dtype != torch.int32 - or not generation_kv_lens.is_cuda - or generation_kv_lens.ndim != 1 - or generation_kv_lens.stride(0) != 1 - or generation_kv_lens.numel() < num_generation_sequences - ): - raise ValueError( - "FP4 MLA device page metadata requires a contiguous CUDA int32 " - "generation KV-length tensor with at least " - f"{num_generation_sequences} entries." - ) - - cache_pool_id = int(getattr(metadata.fp4_mla_state, "cache_pool_id", 0)) - block_offsets = metadata.kv_cache_block_offsets[ - cache_pool_id, - :num_sequences, - 0, - :max_pages, - ] - page_ids = metadata.fp4_mla_state._paged_kv_indices[: num_sequences * max_pages] - page_index_scale = int(metadata.fp4_mla_state.cache_page_index_scale) - if page_index_scale <= 0: - raise RuntimeError("FP4 MLA device page metadata requires a positive page-index scale.") - grid = ( - num_sequences, - triton.cdiv(max_pages, _FP4_MLA_PAGE_TABLE_TILE_SIZE), - ) - _fp4_mla_materialize_page_table_kernel[grid]( - page_ids, - metadata.fp4_mla_state._paged_kv_indptr, - metadata.fp4_mla_state.paged_kv_indptr_decode, - block_offsets, - kv_lens, - generation_kv_lens, - block_offsets.stride(0), - num_sequences, - num_contexts, - max_pages, - PAGE_SIZE=metadata.page_size, - PAGE_INDEX_SCALE=page_index_scale, - PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, - num_warps=4, - ) - hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) - if not isinstance(hp_page_ids, torch.Tensor): - raise RuntimeError("FP4 MLA requires an HP page-table output tensor.") - hp_pool_id = int(metadata.fp4_mla_state.hp_pool_id) - hp_page_index_scale = int(metadata.fp4_mla_state.hp_page_index_scale) - hp_block_offsets = metadata.kv_cache_block_offsets[ - hp_pool_id, - :num_sequences, - 0, - :max_pages, - ] - _fp4_mla_materialize_page_table_kernel[grid]( - hp_page_ids, - metadata.fp4_mla_state._paged_kv_indptr, - metadata.fp4_mla_state.paged_kv_indptr_decode, - hp_block_offsets, - kv_lens, - generation_kv_lens, - hp_block_offsets.stride(0), - num_sequences, - num_contexts, - max_pages, - PAGE_SIZE=metadata.page_size, - PAGE_INDEX_SCALE=hp_page_index_scale, - PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, - num_warps=4, - ) - metadata.fp4_mla_state.device_page_table_valid = True - - -@triton.jit -def _cutedsl_swizzled_sf_offset(row_idx, col_idx, sf_cols: tl.constexpr): - padded_cols = ((sf_cols + 3) // 4) * 4 - return ( - col_idx % 4 - + (col_idx // 4) * (4 * 128) - + (row_idx % 32) * 16 - + ((row_idx % 128) // 32) * 4 - + (row_idx // 128) * (128 * padded_cols) - ) - - -@triton.jit -def _cutedsl_pad_q_and_sf_kernel( - q_padded_ptr, - q_ptr, - q_sf_padded_ptr, - q_sf_ptr, - num_heads, - output_heads: tl.constexpr, - packed_dim: tl.constexpr, - block_bytes: tl.constexpr, - sf_cols: tl.constexpr, - sf_cols_per_byte_block: tl.constexpr, -): - query_idx = tl.program_id(0) - byte_block = tl.program_id(1) - head_offsets = tl.arange(0, output_heads) - byte_offsets = byte_block * block_bytes + tl.arange(0, block_bytes) - head_mask = head_offsets < num_heads - byte_mask = byte_offsets < packed_dim - source_rows = query_idx * num_heads + head_offsets - destination_rows = query_idx * output_heads + head_offsets - values = tl.load( - q_ptr + source_rows[:, None] * packed_dim + byte_offsets[None, :], - mask=head_mask[:, None] & byte_mask[None, :], - other=0, - ) - tl.store( - q_padded_ptr + destination_rows[:, None] * packed_dim + byte_offsets[None, :], - values, - mask=byte_mask[None, :], - ) - sf_col_offsets = byte_block * sf_cols_per_byte_block + tl.arange(0, sf_cols_per_byte_block) - sf_col_mask = sf_col_offsets < sf_cols - source_offsets = _cutedsl_swizzled_sf_offset( - source_rows[:, None], sf_col_offsets[None, :], sf_cols - ) - destination_offsets = _cutedsl_swizzled_sf_offset( - destination_rows[:, None], sf_col_offsets[None, :], sf_cols - ) - sf_values = tl.load( - q_sf_ptr + source_offsets, - mask=head_mask[:, None] & sf_col_mask[None, :], - other=1.0, - ) - tl.store( - q_sf_padded_ptr + destination_offsets, - sf_values, - mask=sf_col_mask[None, :], - ) - - -_SM_COUNT_CACHE: dict[int, int] = {} - - -def _get_sm_count(device: torch.device) -> int: - """Return the SM (multiprocessor) count for ``device``, cached per index.""" - index = device.index if device.index is not None else torch.cuda.current_device() - count = _SM_COUNT_CACHE.get(index) - if count is None: - count = torch.cuda.get_device_properties(index).multi_processor_count - _SM_COUNT_CACHE[index] = count - return count - - -def _validate_fp4_mla_context_rope( - latent_cache: torch.Tensor, - rotary_cos_sin: torch.Tensor, - v_head_dim: int, -) -> int: - head_dim = latent_cache.shape[-1] - rope_dim = head_dim - v_head_dim - if rope_dim <= 0 or rope_dim % 2 != 0: - raise ValueError( - "FP4 MLA fused context K-RoPE requires a positive even RoPE dimension, " - f"got head_dim={head_dim}, v_head_dim={v_head_dim}." - ) - if rotary_cos_sin.device != latent_cache.device: - raise ValueError("FP4 MLA context latent cache and RoPE table must use the same device.") - if rotary_cos_sin.dtype != torch.float32: - raise TypeError( - f"FP4 MLA fused context K-RoPE requires a float32 table, got {rotary_cos_sin.dtype}." - ) - if not rotary_cos_sin.is_contiguous(): - raise ValueError("FP4 MLA fused context K-RoPE requires a contiguous RoPE table.") - table_row_size = rope_dim * 2 - if rotary_cos_sin.numel() < table_row_size or rotary_cos_sin.numel() % table_row_size != 0: - raise ValueError( - "FP4 MLA context RoPE table size must be a positive multiple of " - f"{table_row_size}, got {rotary_cos_sin.numel()}." - ) - return rope_dim - - -def _host_int_list_during_forward(value: Any, start: int, end: int) -> Optional[list[int]]: - if torch.cuda.is_current_stream_capturing(): - return None - return _host_int_list(value, start, end) - - -# FP4 MLA scale-layout helpers - - -def get_fp4_mla_v_scale_pool_size(v_head_dim: int, page_size: int) -> int: - """Return elements per page for the swizzled FP4 MLA V-scale pool. - - The PV matmul treats V as a RHS matrix shaped ``[v_head_dim, kv_tokens]``. - NVFP4 block scales therefore group along the token/K axis, not along the - latent dimension as the K-view cache does. The physical layout matches the - Triton block-scaled matmul scale layout: - ``[ceil(v_head_dim / 128), ceil(page_size / 16 / 4), 32, 16]``. - """ - return _get_fp4_mla_swizzled_scale_size(v_head_dim, page_size) - - -def _get_fp4_mla_swizzled_scale_size(rows: int, cols: int) -> int: - scale_cols = _ceil_div(cols, FP4_BLOCK_SIZE) - row_groups = _ceil_div(rows, FP4_MLA_SCALE_ROW_GROUP) - col_groups = _ceil_div(scale_cols, FP4_MLA_SCALE_COL_GROUP) - return row_groups * col_groups * 32 * 16 - - -def can_fuse_fp4_mla_q_quant( - metadata: Any, - q: torch.Tensor, - q_pe: torch.Tensor, - latent_cache: torch.Tensor, -) -> bool: - """Return whether generation can quantize Q in the fused cache update.""" - num_gen = metadata.num_seqs - metadata.num_contexts - return bool( - _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS - and get_sm_version() == 107 - and num_gen > 0 - and getattr(metadata, "kv_cache_manager", None) is not None - and q.shape[0] > 0 - and q.shape[0] % num_gen == 0 - and q.is_cuda - and q_pe.is_cuda - and latent_cache.is_cuda - and q.device == q_pe.device == latent_cache.device - and q.dtype == torch.bfloat16 - and q.is_contiguous() - and q.ndim == 3 - and 0 < q.shape[1] <= 128 - and q.shape[2] == FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM - and q_pe.dtype == torch.bfloat16 - and tuple(q_pe.shape) == (q.shape[0], q.shape[1], FP4_MLA_Q_RESIDUAL_DIM) - and latent_cache.dtype == torch.bfloat16 - and tuple(latent_cache.shape) == (q.shape[0], FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM) - and metadata.page_size == FP4_MLA_TOKENS_PER_BLOCK - ) - - -def _prepare_fp4_mla_q_buffers( - metadata: Any, - num_queries: int, - num_heads: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, int]: - """Return manager-owned, fixed-capacity packed-Q staging buffers.""" - if num_queries <= 0: - raise ValueError(f"FP4 MLA Q buffer preparation needs queries, got {num_queries}.") - owner = getattr(metadata, "kv_cache_manager", None) - if owner is None: - raise RuntimeError("Fused FP4 MLA Q quantization requires a KV cache manager.") - if num_heads <= 0 or num_heads > 128: - raise ValueError(f"FP4 MLA Q buffers require 1-128 local heads, got {num_heads}.") - buffers = getattr(owner, "_fp4_mla_q_buffers", None) - if buffers is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "Cannot create FP4 MLA Q buffers while capturing a CUDA graph. " - "Run a warmup forward first." - ) - buffers = {} - setattr(owner, "_fp4_mla_q_buffers", buffers) - - max_num_tokens = int(getattr(metadata, "max_num_tokens", num_queries) or num_queries) - max_num_sequences = int( - getattr(metadata, "max_num_sequences", None) - or getattr(metadata, "max_num_requests", num_queries) - or num_queries - ) - max_query_width = 1 + int(getattr(metadata, "max_total_draft_tokens", None) or 0) - capacity = min( - max_num_tokens, - max_num_sequences * max_query_width, - _FP4_MLA_MAX_GRID_Z, - ) - if num_queries > capacity: - raise ValueError( - f"FP4 MLA active queries exceed the configured Q capacity: {num_queries} > {capacity}." - ) - - device_index = device.index if device.index is not None else torch.cuda.current_device() - canonical_device = torch.device("cuda", device_index) - expected_q_shape = (capacity * num_heads, FP4_MLA_Q_PACKED_DIM) - expected_q_sf_shape = ( - _get_fp4_mla_swizzled_scale_size(capacity * num_heads, FP4_MLA_Q_LOGICAL_DIM), - ) - q_key = f"q_{num_heads}" - q_sf_key = f"q_sf_{num_heads}" - q_storage = buffers.get(q_key) - q_sf_storage = buffers.get(q_sf_key) - if q_storage is None and q_sf_storage is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "Cannot allocate FP4 MLA Q buffers while capturing a CUDA graph. " - "Run a warmup forward first." - ) - q_storage = torch.empty(expected_q_shape, dtype=torch.uint8, device=canonical_device) - q_sf_storage = torch.empty( - expected_q_sf_shape, - dtype=torch.float8_e4m3fn, - device=canonical_device, - ) - buffers[q_key] = q_storage - buffers[q_sf_key] = q_sf_storage - if ( - q_storage is None - or q_storage.dtype != torch.uint8 - or q_storage.device != canonical_device - or tuple(q_storage.shape) != expected_q_shape - or not q_storage.is_contiguous() - or q_sf_storage is None - or q_sf_storage.dtype != torch.float8_e4m3fn - or q_sf_storage.device != canonical_device - or tuple(q_sf_storage.shape) != expected_q_sf_shape - or not q_sf_storage.is_contiguous() - ): - raise RuntimeError("FP4 MLA packed-Q buffers do not match the configured capacity.") - return q_storage, q_sf_storage, capacity - - -def _get_fp4_mla_context_start_positions(metadata: Any, num_contexts: int) -> torch.Tensor: - kv_cache_params = getattr(metadata, "kv_cache_params", None) - cached_token_lens = getattr(kv_cache_params, "num_cached_tokens_per_seq", None) - if cached_token_lens is not None: - return torch.as_tensor(cached_token_lens[:num_contexts], dtype=torch.int64, device="cpu") - - return ( - ( - metadata.kv_lens_cuda_runtime[:num_contexts] - - metadata.prompt_lens_cuda_runtime[:num_contexts] - ) - .detach() - .cpu() - ) - - -def _validate_fp4_mla_context_start_alignment( - metadata: Any, - num_contexts: int, - *, - alignment: int = HP_BLOCK_SIZE, -) -> None: - context_start_positions = _get_fp4_mla_context_start_positions(metadata, num_contexts) - bad_start = (context_start_positions < 0) | ((context_start_positions % alignment) != 0) - if bool(torch.any(bad_start).item()): - starts = context_start_positions.detach().cpu().tolist() - raise ValueError( - "FP4 MLA shared-tile context update requires every context " - f"start position to be {alignment}-token aligned, got " - f"start positions {starts}." - ) - - -def get_fp4_mla_v_scale_pool_shape( - num_layers: int, - num_pages: int, - v_head_dim: int, - page_size: int, -) -> tuple[int, int, int, int, int, int]: - """Return the logical swizzled V-scale view shape. - - The leading dimensions are ``[layer, physical_page]``. The remaining - dimensions are the preshuffled ``[N // 128, K // 16 // 4, 32, 16]`` shape - consumed by Triton block-scaled matmul for the V/PV RHS operand. - """ - token_scale_cols = _ceil_div(page_size, FP4_BLOCK_SIZE) - return ( - num_layers, - num_pages, - _ceil_div(v_head_dim, FP4_MLA_SCALE_ROW_GROUP), - _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP), - 32, - 16, - ) - - -def get_fp4_mla_v_scale_pool_view( - metadata: Any, - *, - v_head_dim: int, -) -> torch.Tensor: - """View the auxiliary MLA V-scale pool in Triton's block-scaled layout.""" - pool = getattr(metadata.fp4_mla_state, "v_scale_pool", None) - if pool is None: - raise RuntimeError("FP4 MLA V scale pool is not allocated.") - - elems_per_page = get_fp4_mla_v_scale_pool_size(v_head_dim, metadata.page_size) - if pool.shape[-1] < elems_per_page: - raise RuntimeError( - f"FP4 MLA V scale pool page stride is too small: got " - f"{pool.shape[-1]}, need {elems_per_page}." - ) - - token_scale_cols = _ceil_div(metadata.page_size, FP4_BLOCK_SIZE) - col_groups = _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP) - shape = get_fp4_mla_v_scale_pool_shape( - pool.shape[0], pool.shape[1], v_head_dim, metadata.page_size - ) - strides = ( - pool.stride(0), - pool.stride(1), - col_groups * 32 * 16, - 32 * 16, - 16, - 1, - ) - return torch.as_strided(pool, size=shape, stride=strides) - - -# Python launch helpers - - -def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata.fp4_mla_state, "kv_global_scale", None) - if ( - not isinstance(global_scale, torch.Tensor) - or global_scale.device != device - or global_scale.dtype != torch.float32 - or global_scale.numel() != 1 - ): - raise RuntimeError("FP4 MLA requires a preallocated FP32 KV global-scale tensor.") - return global_scale - - -def _get_fp4_mla_q_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata.fp4_mla_state, "q_global_scale", None) - if ( - not isinstance(global_scale, torch.Tensor) - or global_scale.device != device - or global_scale.dtype != torch.float32 - or global_scale.numel() != 1 - ): - raise RuntimeError("FP4 MLA requires a preallocated FP32 Q global-scale tensor.") - return global_scale - - -def _get_fp4_mla_kv_cache_tensors( - metadata: Any, layer_idx: int -) -> tuple[torch.Tensor, torch.Tensor]: - return metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) - - -def _get_fp4_mla_hp_pool_layout( - metadata: Any, - pool: torch.Tensor, -) -> tuple[int, int]: - """Return the manager-owned HP ring size and per-token head dimension.""" - manager = getattr(metadata, "kv_cache_manager", None) - if manager is None or not hasattr(manager, "fp4_mla_hp_pool_size"): - raise ValueError("FP4 MLA requires a V2 manager-owned HP ring.") - hp_pool_size = manager.fp4_mla_hp_pool_size - if ( - hp_pool_size < HP_BLOCK_SIZE - or pool.ndim != 4 - or pool.shape[2] < 1 - or pool.shape[-1] % hp_pool_size != 0 - ): - raise ValueError( - "FP4 MLA high-precision KV pool does not match its configured " - f"ring: shape={tuple(pool.shape)}, ring_size={hp_pool_size}." - ) - return hp_pool_size, pool.shape[-1] // hp_pool_size - - -def _validate_fp4_mla_hp_generation_width( - hp_pool_size: int, - generation_len: int, -) -> None: - """Ensure one target plus rewindable drafts fit without clobbering the live tail.""" - max_rewind_len = hp_pool_size - HP_BLOCK_SIZE - if generation_len <= 0 or generation_len - 1 > max_rewind_len: - raise RuntimeError( - "FP4 MLA generation exceeds the HP ring's rewind slack: " - f"generation={generation_len}, max_rewind={max_rewind_len}." - ) - - -def _validate_fp4_mla_kv_storage_shape( - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - *, - head_dim: int, - backend: str, -) -> int: - """Validate the backend-specific physical KV and scale strides.""" - residual_dim = FP4_MLA_K_RESIDUAL_DIM if backend in _FP4_MLA_K_RESIDUAL_BACKENDS else 0 - expected_storage_head_dim = head_dim + residual_dim - storage_head_dim = kv_cache.shape[-1] * 2 - if storage_head_dim != expected_storage_head_dim: - raise RuntimeError( - "FP4 MLA KV cache storage head dimension does not match the selected backend: " - f"got {storage_head_dim}, expected {expected_storage_head_dim}. Recreate the engine " - f"after setting {FP4_MLA_ATTENTION_BACKEND_ENV}." - ) - - expected_scale_columns = expected_storage_head_dim // FP4_BLOCK_SIZE - if sf_cache.shape[-1] != expected_scale_columns: - raise RuntimeError( - "FP4 MLA KV cache scale storage does not match the contiguous data layout: " - f"got {sf_cache.shape[-1]} columns, expected {expected_scale_columns}." - ) - return storage_head_dim - - -def _scatter_fp4_mla_kv_cache_2d_context( - metadata: Any, - latent_cache: torch.Tensor, - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - v_sf: torch.Tensor, - global_scale: torch.Tensor, - rotary_cos_sin: Optional[torch.Tensor], - *, - token_offset: int, - local_layer: int, - v_head_dim: int, - head_dim: int, - num_tokens: int, - num_dim_blocks: int, - sf_per_token: int, - sf_per_page: int, - v_packed_base: Optional[torch.Tensor] = None, - v_page_offset: int = 0, -) -> bool: - num_contexts = metadata.num_contexts - if num_contexts > 0: - prompt_lens_cpu = metadata.prompt_lens_cpu_runtime[:num_contexts] - ctx_token_count = int(prompt_lens_cpu.sum().item()) - if num_tokens != ctx_token_count: - raise RuntimeError( - f"FP4 MLA 2D context scatter needs {ctx_token_count} context tokens, got " - f"{num_tokens}." - ) - _validate_fp4_mla_context_start_alignment(metadata, num_contexts, alignment=FP4_BLOCK_SIZE) - - apply_k_rope = rotary_cos_sin is not None - rope_dim = ( - _validate_fp4_mla_context_rope(latent_cache, rotary_cos_sin, v_head_dim) - if rotary_cos_sin is not None - else 0 - ) - rotary_cos_sin_ptr = rotary_cos_sin if rotary_cos_sin is not None else latent_cache - - hp_pool = getattr(metadata.fp4_mla_state, "hp_pool", None) - if not isinstance(hp_pool, torch.Tensor): - raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") - if hp_pool.device != latent_cache.device: - raise ValueError("FP4 MLA latent cache and high-precision pool must share a device.") - if hp_pool.dtype != torch.bfloat16: - raise TypeError(f"FP4 MLA high-precision KV pool must use BF16, got {hp_pool.dtype}.") - hp_pool_size, pool_head_dim = _get_fp4_mla_hp_pool_layout(metadata, hp_pool) - if pool_head_dim < head_dim: - raise RuntimeError( - f"FP4 MLA HP pool head dimension is too small: got " - f"{pool_head_dim}, need at least {head_dim}." - ) - if local_layer < 0 or local_layer >= hp_pool.shape[1]: - raise ValueError( - f"FP4 MLA local layer {local_layer} is outside the HP pool's {hp_pool.shape[1]} layers." - ) - if hp_pool.stride(-1) != 1: - raise ValueError("FP4 MLA high-precision KV pool must be contiguous in head_dim.") - hp_page_ids = metadata.fp4_mla_state.hp_page_indices - if not isinstance(hp_page_ids, torch.Tensor): - raise RuntimeError("FP4 MLA context cache update requires HP page metadata.") - store_hp_tail = num_contexts > 0 - num_hp_pages = hp_pool.shape[0] - pool_s0 = hp_pool.stride(0) - pool_s1 = hp_pool.stride(1) - - write_v_packed = v_packed_base is not None - v_packed_output = v_packed_base if write_v_packed else kv_cache - v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 - v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 - - _fp4_mla_context_cache_update_kernel[ - ( - num_tokens, - num_dim_blocks, - ) - ]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - latent_cache, - global_scale, - rotary_cos_sin_ptr, - hp_pool, - hp_page_ids, - metadata.fp4_mla_state.batch_indices, - metadata.fp4_mla_state.positions, - metadata.fp4_mla_state.paged_kv_indices, - metadata.fp4_mla_state.paged_kv_indptr, - metadata.fp4_mla_state.paged_kv_indices.shape[0], - metadata.fp4_mla_state.paged_kv_indptr.shape[0], - metadata.fp4_mla_state.batch_indices.shape[0], - v_sf.shape[1], - v_sf.shape[0], - num_contexts, - num_hp_pages, - token_offset, - num_tokens, - local_layer, - v_page_offset if write_v_packed else 0, - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - latent_cache.stride(0), - latent_cache.stride(1), - v_sf.stride(0), - v_sf.stride(1), - v_packed_s0, - v_packed_s1, - pool_s0, - pool_s1, - HEAD_D=head_dim, - V_HEAD_D=v_head_dim, - HP_BLOCK=FP4_BLOCK_SIZE, - HP_POOL_SIZE=hp_pool_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=sf_per_page, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - STORE_K_RESIDUAL=(_fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS), - ROPE_DIM=rope_dim, - APPLY_K_ROPE=apply_k_rope, - POOL_HEAD_D=pool_head_dim, - STORE_HP_TAIL=store_hp_tail, - WRITE_V_PACKED=write_v_packed, - ) - return store_hp_tail - - -def _fp4_mla_generation_num_blocks_device(metadata: Any) -> torch.Tensor: - """Device-side scalar view holding the generation page-table capacity. - - ``paged_kv_indptr_decode[num_gen]`` is the fixed-stride generation-table - endpoint. Device kernels combine it with the live KV lengths, so inactive - slots are never consumed. - """ - num_gen = metadata.num_seqs - metadata.num_contexts - return metadata.fp4_mla_state.paged_kv_indptr_decode[num_gen : num_gen + 1] - - -def _fp4_mla_uniform_generation_lengths( - metadata: Any, num_gen_tokens: int, num_gen: int -) -> tuple[torch.Tensor, torch.Tensor]: - """Return preallocated CUDA generation KV and append lengths.""" - if num_gen <= 0 or num_gen_tokens % num_gen != 0: - raise RuntimeError("FP4 MLA generation requires a non-empty uniform request batch.") - - num_contexts = metadata.num_contexts - num_seqs = metadata.num_seqs - kv_lens_gen = metadata.kv_lens_cuda_runtime[num_contexts:num_seqs] - prompt_lens_gen = metadata.prompt_lens_cuda_runtime[num_contexts:num_seqs] - corrected_kv_lens = getattr(metadata.fp4_mla_state, "generation_kv_lens", None) - generation_lens = getattr(metadata.fp4_mla_state, "generation_append_lens", None) - tensors = ( - kv_lens_gen, - prompt_lens_gen, - corrected_kv_lens, - generation_lens, - ) - if ( - not all(isinstance(tensor, torch.Tensor) for tensor in tensors) - or corrected_kv_lens.numel() < num_gen - or generation_lens.numel() < num_gen - or not all(tensor.is_cuda for tensor in tensors) - ): - raise RuntimeError("FP4 MLA generation lengths require preallocated CUDA buffers.") - - record_for_capture = bool( - getattr(metadata, "is_cuda_graph", False) - and torch.cuda.is_current_stream_capturing() - and not getattr(metadata.fp4_mla_state, "generation_lengths_capture_recorded", False) - ) - precomputed = ( - not record_for_capture - and metadata.fp4_mla_state.generation_lengths_num_tokens == num_gen_tokens - and metadata.fp4_mla_state.generation_lengths_num_seqs == num_gen - and metadata.fp4_mla_state.generation_lengths_num_contexts == num_contexts - ) - if not precomputed: - populate_fp4_mla_generation_lengths( - kv_lens_gen, - prompt_lens_gen, - corrected_kv_lens[:num_gen], - generation_lens[:num_gen], - num_gen_tokens=num_gen_tokens, - num_gen=num_gen, - ) - metadata.fp4_mla_state.generation_lengths_num_tokens = num_gen_tokens - metadata.fp4_mla_state.generation_lengths_num_seqs = num_gen - metadata.fp4_mla_state.generation_lengths_num_contexts = num_contexts - if record_for_capture: - metadata.fp4_mla_state.generation_lengths_capture_recorded = True - return corrected_kv_lens[:num_gen], generation_lens[:num_gen] - - -def _materialize_fp4_mla_device_page_table_for_forward( - metadata: Any, - generation_kv_lens: Optional[torch.Tensor] = None, -) -> None: - """Materialize all fixed-stride rows from final per-forward device lengths.""" - if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): - raise RuntimeError("FP4 MLA cache update requires fixed-stride device page metadata.") - if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): - return - - num_contexts = int(metadata.num_contexts) - num_sequences = int(metadata.num_seqs) - num_generation_sequences = num_sequences - num_contexts - if generation_kv_lens is None: - if num_generation_sequences > 0: - num_generation_tokens = int(metadata.num_tokens) - int(metadata.num_ctx_tokens) - generation_kv_lens, _ = _fp4_mla_uniform_generation_lengths( - metadata, - num_generation_tokens, - num_generation_sequences, - ) - else: - generation_kv_lens = metadata.kv_lens_cuda_runtime[num_contexts:num_sequences] - materialize_fp4_mla_device_page_table( - metadata, - metadata.kv_lens_cuda_runtime[:num_sequences], - generation_kv_lens, - ) - - -def _scatter_fp4_mla_kv_cache_2d_generation( - metadata: Any, - latent_cache: torch.Tensor, - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - v_sf: torch.Tensor, - global_scale: torch.Tensor, - *, - token_offset: int, - local_layer: int, - v_head_dim: int, - head_dim: int, - num_tokens: int, - num_dim_blocks: int, - sf_per_token: int, - sf_per_page: int, - rotary_cos_sin: torch.Tensor, - q_pe: torch.Tensor, - q_rope_out: torch.Tensor, - q_quant_input: torch.Tensor, - q_fp4_out: torch.Tensor, - q_sf_out: torch.Tensor, - v_packed_base: Optional[torch.Tensor], - v_page_offset: int, -) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - num_contexts = metadata.num_contexts - num_seqs = metadata.num_seqs - num_gen = num_seqs - num_contexts - if num_gen <= 0: - return - if num_tokens < num_gen: - raise RuntimeError( - f"FP4 MLA 2D generation scatter needs at least {num_gen} generation " - f"tokens, got {num_tokens}." - ) - if num_tokens % num_gen != 0: - raise NotImplementedError( - "FP4 MLA no-dequant generation scatter requires a uniform linear MTP " - f"generation length, got {num_tokens} tokens for {num_gen} sequences." - ) - # The prompt_lens/kv_lens runtime aliases can lag at the decode anchor - # (seq_lens == 1) under CUDA graph / one-engine MTP while each generation - # sequence really appends num_tokens // num_gen tokens this step. Recover the - # true per-sequence lengths for the no-dequant kernel below (a no-op when the - # aliases already match). - kv_lens_gen, gen_lens_gen = _fp4_mla_uniform_generation_lengths(metadata, num_tokens, num_gen) - _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens_gen) - - pool = getattr(metadata.fp4_mla_state, "hp_pool", None) - if pool is None: - raise RuntimeError("FP4 MLA 2D generation scatter requires the HP KV pool.") - try: - hp_pool_size, hp_head_dim = _get_fp4_mla_hp_pool_layout(metadata, pool) - except ValueError as error: - raise RuntimeError(str(error)) from error - if hp_head_dim < head_dim: - raise RuntimeError( - f"FP4 MLA 2D generation scatter needs at least {head_dim} HP channels, got " - f"{hp_head_dim}." - ) - hp_page_ids = _fp4_mla_generation_hp_page_ids(metadata, num_gen) - if not isinstance(hp_page_ids, torch.Tensor): - raise RuntimeError("FP4 MLA generation requires HP page metadata.") - num_hp_pages = pool.shape[0] - - max_gen_len = num_tokens // num_gen - _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) - max_rewind_len = hp_pool_size - HP_BLOCK_SIZE - page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) - rope_dim = head_dim - v_head_dim - block_q_heads = 32 - q1_kv_blocks_per_program = 1 - # Grouped Q1 kernels reuse K's packed codes for V. Keep one dimension - # block per program until their warp-specialized V quantizer is split out. - if max_gen_len == 1: - q1_kv_blocks_per_program = _fp4_mla_q1_kv_blocks_per_program(num_gen, v_head_dim) - q_prefix_block_dim = ( - FP4_MLA_Q1_PREFIX_BLOCK_DIM if max_gen_len == 1 else FP4_MLA_Q_PREFIX_BLOCK_DIM - ) - q_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // q_prefix_block_dim - q_prefix_blocks_per_program = _fp4_mla_q1_prefix_blocks_per_program( - num_gen, - q1_kv_blocks_per_program, - ) - q_work_blocks = q_prefix_blocks // q_prefix_blocks_per_program + 1 - if latent_cache.dtype != torch.bfloat16: - raise TypeError( - "Fused FP4 MLA Q/K RoPE and cache storage requires BF16 latent KV, " - f"got {latent_cache.dtype}." - ) - if rope_dim <= 0 or rope_dim % 2 != 0: - raise ValueError( - "Fused FP4 MLA K RoPE requires a positive even K tail, " - f"got head_dim={head_dim} v_head_dim={v_head_dim}." - ) - if ( - rotary_cos_sin is None - or rotary_cos_sin.device != latent_cache.device - or rotary_cos_sin.dtype != torch.float32 - or rotary_cos_sin.numel() % (rope_dim * 2) != 0 - ): - raise ValueError( - "Fused FP4 MLA K RoPE requires a same-device FP32 rotary " - f"table with rows of {rope_dim * 2} values." - ) - if ( - q_pe is None - or q_rope_out is None - or q_pe.dtype != torch.bfloat16 - or q_rope_out.dtype != torch.bfloat16 - or q_pe.device != latent_cache.device - or q_rope_out.device != latent_cache.device - or q_pe.ndim != 3 - or q_rope_out.shape != q_pe.shape - or q_pe.shape[0] != num_tokens - or q_pe.shape[1] <= 0 - or q_pe.shape[2] != rope_dim - ): - raise ValueError( - "Fused FP4 MLA Q RoPE requires same-device BF16 q_pe and " - f"q_rope_out tensors shaped [tokens, heads, {rope_dim}]." - ) - num_q_heads = q_pe.shape[1] - q_head_blocks = _ceil_div(num_q_heads, block_q_heads) - max_gen_tiles = _ceil_div(max_gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) - rotary_table = rotary_cos_sin - q_global_scale = _get_fp4_mla_q_global_scale(metadata, latent_cache.device) - q_pe_input = q_pe - q_rope_output = q_rope_out - q_full_input = q_quant_input - q_fp4_output = q_fp4_out - q_sf_output = q_sf_out - write_v_packed = v_packed_base is not None - v_packed_output = v_packed_base if write_v_packed else kv_cache - v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 - v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 - kv_work_blocks = ( - v_head_dim // FP4_BLOCK_SIZE // q1_kv_blocks_per_program + 1 - if max_gen_len == 1 - else num_dim_blocks - ) - launch_grid = ( - num_gen, - max( - kv_work_blocks, - max_gen_len * q_head_blocks * q_work_blocks, - ), - ) - store_k_residual = _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS - - def launch_generation_update( - grid: tuple[int, ...], - *, - page_ids_len: int, - indptr_len: int, - max_gen_tiles_variant: int, - q_prefix_block_dim_variant: int, - q_prefix_blocks_variant: int, - q_prefix_blocks_per_program_variant: int, - q1_kv_blocks_per_program_variant: int, - ) -> None: - q_work_blocks_variant = q_prefix_blocks_variant // q_prefix_blocks_per_program_variant + 1 - _fp4_mla_generation_fused_qk_rope_cache_update_kernel[grid]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - pool, - latent_cache, - global_scale, - q_global_scale, - rotary_table, - q_pe_input, - q_rope_output, - q_full_input, - q_fp4_output, - q_sf_output, - kv_lens_gen, - gen_lens_gen, - page_ids, - hp_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, - page_ids_len, - hp_page_ids.numel(), - indptr_len, - v_sf.shape[1], - num_hp_pages, - v_sf.shape[0], - local_layer, - v_page_offset if write_v_packed else 0, - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - pool.stride(0), - pool.stride(1), - v_sf.stride(0), - v_sf.stride(1), - v_packed_s0, - v_packed_s1, - q_pe_input.stride(0), - q_pe_input.stride(1) if q_pe_input.ndim > 1 else 0, - q_pe_input.stride(2) if q_pe_input.ndim > 2 else 0, - q_rope_output.stride(0), - q_rope_output.stride(1) if q_rope_output.ndim > 1 else 0, - q_rope_output.stride(2) if q_rope_output.ndim > 2 else 0, - HEAD_D=hp_head_dim, - V_HEAD_D=v_head_dim, - HP_BLOCK=FP4_BLOCK_SIZE, - HP_POOL_SIZE=hp_pool_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=sf_per_page, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - STORE_K_RESIDUAL=store_k_residual, - FUSE_ROPE_CACHE_STORE=True, - WRITE_V_PACKED=write_v_packed, - MAX_GEN_TILES=max_gen_tiles_variant, - ROPE_DIM=rope_dim, - ROPE_PAIR_BLOCK=triton.next_power_of_2(rope_dim // 2), - NUM_DIM_BLOCKS=num_dim_blocks, - NUM_Q_HEADS=num_q_heads, - Q_HEAD_BLOCKS=max(q_head_blocks, 1), - BLOCK_Q_HEADS=block_q_heads, - Q_PREFIX_D=FP4_MLA_Q_PREFIX_DIM, - Q_PREFIX_BLOCK_D=q_prefix_block_dim_variant, - Q_PREFIX_BLOCKS=q_prefix_blocks_variant, - Q_PREFIX_BLOCKS_PER_PROGRAM=q_prefix_blocks_per_program_variant, - Q_WORK_BLOCKS=q_work_blocks_variant, - Q_SF_COLS=FP4_MLA_Q_SF_GROUPS, - WRITE_Q=True, - Q1_KV_BLOCKS_PER_PROGRAM=q1_kv_blocks_per_program_variant, - maxnreg=56, - ) - - # Triton compiles and loads a CUDA module on first launch. Use runtime-zero - # work here so every reachable static tuning variant is resident before - # warmup hands the engine to serving. - if getattr(metadata, "is_warmup", False) and not torch.cuda.is_current_stream_capturing(): - configured_generation_len = int(getattr(metadata, "max_total_draft_tokens", 0) or 0) + 1 - max_preload_generation_len = max(max_gen_len, configured_generation_len) - if max_preload_generation_len - 1 > max_rewind_len: - raise NotImplementedError( - "FP4 MLA finite Triton preload exceeds the HP ring's rewind slack: " - f"max_rewind={max_rewind_len}, generation=" - f"{max_preload_generation_len}." - ) - max_num_sequences = int(getattr(metadata, "max_num_sequences", num_seqs) or num_seqs) - q1_variants = _fp4_mla_q1_preload_variants( - max_num_sequences, - v_head_dim, - ) - multi_token_tiles = tuple( - sorted( - { - _ceil_div(gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) - for gen_len in range(2, max_preload_generation_len + 1) - } - ) - ) - preload_key = ( - "generation-cache-update", - str(latent_cache.device), - hp_pool_size, - hp_head_dim, - v_head_dim, - num_q_heads, - metadata.page_size, - write_v_packed, - store_k_residual, - tuple(q1_variants), - multi_token_tiles, - tuple(kv_cache.stride()), - tuple(sf_cache.stride()), - tuple(pool.stride()), - tuple(v_sf.stride()), - tuple(v_packed_output.stride()), - tuple(q_pe_input.stride()), - tuple(q_rope_output.stride()), - str(kv_cache.dtype), - str(latent_cache.dtype), - str(q_pe_input.dtype), - str(q_fp4_output.dtype), - str(q_sf_output.dtype), - ) - preload_keys = _fp4_mla_triton_preload_key_set(metadata) - if preload_key not in preload_keys: - context_page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) - if not isinstance(context_page_ids, torch.Tensor): - context_page_ids = page_ids - context_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) - if not isinstance(context_indptr, torch.Tensor): - context_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode - context_batch_indices = getattr(metadata.fp4_mla_state, "batch_indices", None) - if not isinstance(context_batch_indices, torch.Tensor): - context_batch_indices = hp_page_ids - context_positions = getattr(metadata.fp4_mla_state, "positions", None) - if not isinstance(context_positions, torch.Tensor): - context_positions = hp_page_ids - _fp4_mla_context_cache_update_kernel[(1, num_dim_blocks)]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - latent_cache, - global_scale, - rotary_table, - pool, - hp_page_ids, - context_batch_indices, - context_positions, - context_page_ids, - context_indptr, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - latent_cache.stride(0), - latent_cache.stride(1), - v_sf.stride(0), - v_sf.stride(1), - v_packed_s0, - v_packed_s1, - pool.stride(0), - pool.stride(1), - HEAD_D=head_dim, - V_HEAD_D=v_head_dim, - HP_BLOCK=FP4_BLOCK_SIZE, - HP_POOL_SIZE=hp_pool_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=sf_per_page, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - STORE_K_RESIDUAL=store_k_residual, - ROPE_DIM=rope_dim, - APPLY_K_ROPE=True, - POOL_HEAD_D=hp_head_dim, - STORE_HP_TAIL=True, - WRITE_V_PACKED=write_v_packed, - ) - q1_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM - for q1_kv_blocks, q1_prefix_blocks_per_program in q1_variants: - launch_generation_update( - (1, 1), - page_ids_len=0, - indptr_len=0, - max_gen_tiles_variant=1, - q_prefix_block_dim_variant=FP4_MLA_Q1_PREFIX_BLOCK_DIM, - q_prefix_blocks_variant=q1_prefix_blocks, - q_prefix_blocks_per_program_variant=q1_prefix_blocks_per_program, - q1_kv_blocks_per_program_variant=q1_kv_blocks, - ) - multi_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q_PREFIX_BLOCK_DIM - for multi_token_tile in multi_token_tiles: - launch_generation_update( - (1, 1), - page_ids_len=0, - indptr_len=0, - max_gen_tiles_variant=multi_token_tile, - q_prefix_block_dim_variant=FP4_MLA_Q_PREFIX_BLOCK_DIM, - q_prefix_blocks_variant=multi_prefix_blocks, - q_prefix_blocks_per_program_variant=1, - q1_kv_blocks_per_program_variant=1, - ) - torch.cuda.synchronize(latent_cache.device) - preload_keys.add(preload_key) - - launch_generation_update( - launch_grid, - page_ids_len=page_ids.shape[0], - indptr_len=metadata.fp4_mla_state.paged_kv_indptr_decode.shape[0], - max_gen_tiles_variant=max(max_gen_tiles, 1), - q_prefix_block_dim_variant=q_prefix_block_dim, - q_prefix_blocks_variant=q_prefix_blocks, - q_prefix_blocks_per_program_variant=q_prefix_blocks_per_program, - q1_kv_blocks_per_program_variant=q1_kv_blocks_per_program, - ) - return kv_lens_gen, gen_lens_gen, page_ids - - -# Public cache update and decode entry points - - -def scatter_fp4_mla_kv_cache( - metadata: Any, - latent_cache: torch.Tensor, - layer_idx: int, - *, - token_offset: int, - phase: _HPUpdatePhase, - local_layer: int, - v_head_dim: int, - rotary_cos_sin: Optional[torch.Tensor] = None, - q_pe: Optional[torch.Tensor] = None, - q_rope_out: Optional[torch.Tensor] = None, - q_quant_input: Optional[torch.Tensor] = None, -) -> bool: - """Quantize MLA latent tokens and scatter them into the paged FP4 cache. - - Contract: this helper scatters exactly ``latent_cache.shape[0]`` tokens, - reading index metadata at ``batch_indices[token_offset : token_offset + N]`` - and ``positions[token_offset : token_offset + N]``. Callers must pass a - latent_cache pre-sliced to the current phase (context or generation) so - that ``shape[0]`` matches the number of index entries they intend to - consume. ``MLA.forward_impl`` (tensorrt_llm/_torch/modules/attention.py) - slices ``latent_cache[:num_ctx_tokens]`` for context and - ``latent_cache[num_ctx_tokens:]`` for generation before dispatching. - - Callers must pass ``phase``, ``local_layer``, and ``v_head_dim``. Context - scatter writes the final FP4 tile representation directly. Dimensions - below ``v_head_dim`` share one 16-token by 16-dim FP4 tile between K - and V, with the scale written into K's token-major and V's dim-major - layouts. Tail K-only dimensions use K's per-token 1D scales. For - exclusively owned CuTeDSL pages, context scatter also writes the - persistent packed-V sidecar. - Context scatter can rotate the K tail directly from the unassembled latent - tensor. Generation scatter rewrites each touched 16-token tile by reading - old tokens from the HP pool and new tokens from ``latent_cache``. The - static-scale generation - specialization can also rotate Q and new K tails while updating the HP pool. - The context kernel also stores the final incomplete tile in the BF16 HP - pool. When ``q_quant_input`` is supplied, the generation kernel also emits - backend-ready residual FP4 Q. The return value reports whether the current - phase updated the HP pool. - """ - if phase == "generation": - metadata.fp4_mla_state.prequantized_q = None - metadata.fp4_mla_state.prequantized_q_sf = None - metadata.fp4_mla_state.q_batch_capacity = None - if latent_cache.numel() == 0: - raise ValueError("FP4 MLA cache scatter requires at least one latent token.") - - latent_cache = latent_cache.reshape(latent_cache.shape[0], -1).contiguous() - num_tokens = latent_cache.shape[0] - head_dim = latent_cache.shape[-1] - if head_dim % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA KV head_dim must be divisible by {FP4_BLOCK_SIZE}, got {head_dim}." - ) - indices_len = metadata.fp4_mla_state.batch_indices.shape[0] - positions_len = metadata.fp4_mla_state.positions.shape[0] - if token_offset + num_tokens > indices_len or token_offset + num_tokens > positions_len: - raise RuntimeError( - f"FP4 MLA scatter would read batch_indices[{token_offset}:" - f"{token_offset + num_tokens}] / positions[{token_offset}:" - f"{token_offset + num_tokens}], but only {indices_len} / " - f"{positions_len} entries are available. This indicates " - "latent_cache was not pre-sliced to the current phase's token " - "range (see MLA.forward_impl)." - ) - - _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) - - backend = _fp4_mla_attention_backend() - global_scale = _get_fp4_mla_global_scale(metadata, latent_cache.device) - kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) - storage_head_dim = _validate_fp4_mla_kv_storage_shape( - kv_cache, - sf_cache, - head_dim=head_dim, - backend=backend, - ) - sf_per_token = storage_head_dim // FP4_BLOCK_SIZE - - if phase not in ("context", "generation"): - raise ValueError("FP4 MLA scatter requires phase='context' or 'generation'.") - if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: - raise RuntimeError("FP4 MLA scatter requires the auxiliary V scale pool.") - if metadata.page_size % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA scatter requires page_size divisible by " - f"{FP4_BLOCK_SIZE}, got {metadata.page_size}." - ) - if v_head_dim > head_dim: - raise ValueError(f"FP4 MLA v_head_dim={v_head_dim} cannot exceed head_dim={head_dim}.") - if head_dim - v_head_dim != FP4_MLA_K_RESIDUAL_DIM: - raise ValueError( - "FP4 MLA K residual quantization requires the K-only tail to match " - f"the {FP4_MLA_K_RESIDUAL_DIM}-channel residual, got " - f"head_dim={head_dim} v_head_dim={v_head_dim}." - ) - if v_head_dim % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA v_head_dim must be divisible by {FP4_BLOCK_SIZE}, got {v_head_dim}." - ) - - sf_cache = sf_cache.view(torch.float8_e4m3fn) - v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=v_head_dim) - num_dim_blocks = triton.cdiv(head_dim, FP4_BLOCK_SIZE) - sf_per_page = metadata.page_size // FP4_BLOCK_SIZE - - generation_state = None - generation_inputs = (rotary_cos_sin, q_pe, q_rope_out, q_quant_input) - q_fp4_out = None - q_sf_out = None - if phase == "context": - if any(arg is not None for arg in (q_pe, q_rope_out, q_quant_input)): - raise ValueError("FP4 MLA context cache update does not accept generation Q tensors.") - hp_pool_updated = False - else: - if not all(arg is not None for arg in generation_inputs): - raise ValueError( - "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " - "and q_quant_input for fused RoPE, cache update, and Q quantization." - ) - if not can_fuse_fp4_mla_q_quant(metadata, q_quant_input, q_pe, latent_cache): - raise ValueError( - "Fused FP4 MLA Q quantization received an unsupported shape, dtype, " - "scale mode, or backend." - ) - q_fp4_out, q_sf_out, q_batch_capacity = _prepare_fp4_mla_q_buffers( - metadata, - num_tokens, - q_quant_input.shape[1], - q_quant_input.device, - ) - metadata.fp4_mla_state.prequantized_q = q_fp4_out - metadata.fp4_mla_state.prequantized_q_sf = q_sf_out - metadata.fp4_mla_state.q_batch_capacity = q_batch_capacity - hp_pool_updated = True - cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND - fused_v_transpose = cutedsl_backend and _fp4_mla_cutedsl_fused_v_transpose_enabled() - kv_cache_manager = getattr(metadata, "kv_cache_manager", None) - persistent_v_packed = None - v_packed_base = None - v_page_offset = 0 - direct_v_packed_write = False - if cutedsl_backend and not fused_v_transpose: - persistent_v_packed = _get_cutedsl_persistent_v_packed_cache( - metadata, - local_layer, - kv_cache, - v_head_dim=v_head_dim, - page_size=metadata.page_size, - block_v=FP4_MLA_SCALE_ROW_GROUP, - ) - # Reused/imported pages may not carry this process-local sidecar. Write - # packed V directly only when the cache pages are exclusively owned. - # The fused generation kernel handles uniform linear-MTP batches and - # updates every 16-token tile touched by the verification window. - num_gen = metadata.num_seqs - metadata.num_contexts - block_reuse = getattr(kv_cache_manager, "enable_block_reuse", True) - direct_context_v_packed_write = ( - phase == "context" - and metadata.num_contexts > 0 - and num_tokens > 0 - and block_reuse is False - ) - direct_generation_v_packed_write = ( - phase == "generation" - and num_gen > 0 - and num_tokens >= num_gen - and num_tokens % num_gen == 0 - and block_reuse is False - ) - direct_v_packed_write = direct_context_v_packed_write or direct_generation_v_packed_write - if direct_v_packed_write: - v_packed_base = _get_fp4_mla_v_packed_pool_base(metadata) - get_v_page_offset = getattr(kv_cache_manager, "get_mla_v_packed_page_offset", None) - v_page_offset = ( - int(get_v_page_offset(local_layer)) - if callable(get_v_page_offset) - else local_layer * kv_cache.shape[0] - ) - expected_row_width = metadata.page_size // 2 - required_base_rows = (v_page_offset + kv_cache.shape[0]) * v_head_dim - if ( - not isinstance(v_packed_base, torch.Tensor) - or v_packed_base.dtype != torch.uint8 - or v_packed_base.device != kv_cache.device - or v_packed_base.ndim != 2 - or v_packed_base.shape[0] < required_base_rows - or v_packed_base.shape[1] != expected_row_width - or not v_packed_base.is_contiguous() - ): - raise RuntimeError( - "FP4 MLA direct V-packed cache update requires the " - "stable full-pool base to be a contiguous uint8 tensor " - f"with at least {required_base_rows} rows and " - f"{expected_row_width} columns on {kv_cache.device}." - ) - expected_layer_ptr = v_packed_base.data_ptr() + ( - v_page_offset * v_head_dim * expected_row_width - ) - if expected_layer_ptr != persistent_v_packed.data_ptr(): - raise RuntimeError( - "FP4 MLA V-packed layer view does not match its stable " - "full-pool base and page offset." - ) - v_pack_num_valid = None - if phase == "context": - _materialize_fp4_mla_device_page_table_for_forward(metadata) - hp_pool_updated = _scatter_fp4_mla_kv_cache_2d_context( - metadata, - latent_cache, - kv_cache, - sf_cache, - v_sf, - global_scale, - rotary_cos_sin, - token_offset=token_offset, - local_layer=local_layer, - v_head_dim=v_head_dim, - head_dim=head_dim, - num_tokens=num_tokens, - num_dim_blocks=num_dim_blocks, - sf_per_token=sf_per_token, - sf_per_page=sf_per_page, - v_packed_base=v_packed_base, - v_page_offset=v_page_offset, - ) - v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices - else: - generation_state = _scatter_fp4_mla_kv_cache_2d_generation( - metadata, - latent_cache, - kv_cache, - sf_cache, - v_sf, - global_scale, - token_offset=token_offset, - local_layer=local_layer, - v_head_dim=v_head_dim, - head_dim=head_dim, - num_tokens=num_tokens, - num_dim_blocks=num_dim_blocks, - sf_per_token=sf_per_token, - sf_per_page=sf_per_page, - rotary_cos_sin=rotary_cos_sin, - q_pe=q_pe, - q_rope_out=q_rope_out, - q_quant_input=q_quant_input, - q_fp4_out=q_fp4_out, - q_sf_out=q_sf_out, - v_packed_base=v_packed_base, - v_page_offset=v_page_offset, - ) - v_pack_page_ids = _fp4_mla_generation_page_ids( - metadata, metadata.num_seqs - metadata.num_contexts - ) - if getattr(metadata, "is_cuda_graph", False): - # Frozen launch grids cannot follow the per-replay page count; - # the repack kernels stride over this device-side count instead. - v_pack_num_valid = _fp4_mla_generation_num_blocks_device(metadata) - cutedsl_repack_page_indptr = None - cutedsl_repack_kv_lens = None - cutedsl_repack_generation_lens = None - cutedsl_repack_max_touched_pages = 1 - if cutedsl_backend and not fused_v_transpose and not direct_v_packed_write: - if phase == "context": - num_contexts = metadata.num_contexts - cutedsl_v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices[ - : metadata.fp4_mla_state.num_context_blocks - ] - cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr[: num_contexts + 1] - cutedsl_repack_kv_lens = metadata.kv_lens_cuda_runtime[:num_contexts] - cutedsl_repack_generation_lens = metadata.prompt_lens_cuda_runtime[:num_contexts] - cutedsl_repack_max_touched_pages = int( - metadata.fp4_mla_state.context_repack_max_touched_pages - ) - elif generation_state is not None: - kv_lens_gen, gen_lens_gen, generation_page_ids = generation_state - num_gen = kv_lens_gen.numel() - cutedsl_v_pack_page_ids = generation_page_ids - cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode - cutedsl_repack_kv_lens = kv_lens_gen - cutedsl_repack_generation_lens = gen_lens_gen - cutedsl_repack_max_touched_pages = _ceil_div( - num_tokens // num_gen + metadata.page_size - 1, - metadata.page_size, - ) - else: - cutedsl_v_pack_page_ids = v_pack_page_ids - _repack_cutedsl_v_packed_cache( - persistent_v_packed, - kv_cache, - cutedsl_v_pack_page_ids, - v_head_dim=v_head_dim, - page_size=metadata.page_size, - block_v=FP4_MLA_SCALE_ROW_GROUP, - page_indptr=cutedsl_repack_page_indptr, - kv_lens=cutedsl_repack_kv_lens, - generation_lens=cutedsl_repack_generation_lens, - max_touched_pages=cutedsl_repack_max_touched_pages, - ) - _maybe_update_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - v_pack_page_ids, - num_queries=num_tokens, - v_head_dim=v_head_dim, - page_size=metadata.page_size, - local_layer=local_layer, - v_sf=v_sf[local_layer], - num_valid_pages=v_pack_num_valid, - ) - return hp_pool_updated - - -def _validate_fp4_mla_cache_shape(page_size: int, head_dim: int) -> None: - if page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - f"FP4 MLA KV cache requires tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK} " - f"for swizzled block scales, got {page_size}." - ) - - sf_per_token = head_dim // FP4_BLOCK_SIZE - if head_dim % FP4_BLOCK_SIZE != 0 or sf_per_token % 4 != 0: - raise ValueError( - f"FP4 MLA KV head_dim must produce a scale column count divisible by 4; " - f"got head_dim={head_dim}, scale_columns={sf_per_token}." - ) - - -def _validate_fp4_mla_attention_q_shape(head_dim: int, q_residual_dim: int) -> None: - if q_residual_dim % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA Q residual_dim must be divisible by {FP4_BLOCK_SIZE}, got {q_residual_dim}." - ) - if q_residual_dim <= 0 or q_residual_dim > head_dim: - raise ValueError( - f"FP4 MLA Q residual_dim must be in (0, head_dim], got " - f"residual_dim={q_residual_dim}, head_dim={head_dim}." - ) - - q_head_dim = head_dim + q_residual_dim - q_sf_per_token = q_head_dim // FP4_BLOCK_SIZE - if q_head_dim % FP4_BLOCK_SIZE != 0 or q_sf_per_token % FP4_MLA_SCALE_COL_GROUP != 0: - raise ValueError( - f"FP4 MLA residual Q must produce a scale column count divisible " - f"by {FP4_MLA_SCALE_COL_GROUP}; got q_head_dim={q_head_dim}, " - f"scale_columns={q_sf_per_token}." - ) - - -def _ensure_workspace_tensor( - metadata: Any, - attr_name: str, - shape: tuple[int, ...], - *, - dtype: torch.dtype, - device: torch.device, -) -> torch.Tensor: - workspaces = metadata.fp4_mla_state.workspaces - tensor = workspaces.get(attr_name) - needs_alloc = ( - tensor is None - or tensor.dtype != dtype - or tensor.device != device - or len(tensor.shape) != len(shape) - or any(tensor.shape[idx] < dim for idx, dim in enumerate(shape)) - ) - if needs_alloc: - if torch.cuda.is_current_stream_capturing(): - raise ValueError( - f"Cannot allocate {attr_name} while capturing a CUDA graph. " - "Run a warmup prepare/forward first." - ) - tensor = torch.empty(shape, dtype=dtype, device=device) - workspaces[attr_name] = tensor - - slices = tuple(slice(0, dim) for dim in shape) - return tensor[slices] - - -def _shared_v_pack_storage_enabled() -> bool: - return os.getenv("TRTLLM_FP4_MLA_SHARE_V_PACK_STORAGE", "1").lower() not in ( - "0", - "false", - "no", - "off", - ) - - -def _select_triton_block_v(num_queries: int, *, prefer_prepacked_v: bool = False) -> int: - env_block_v = _env_int("TRTLLM_FP4_MLA_BLOCK_V") - if env_block_v is not None: - return env_block_v - if prefer_prepacked_v: - return 128 - return 32 if num_queries <= 32 else 128 - - -def _v_packed_shape( - kv_cache: torch.Tensor, - v_head_dim: int, - page_size: int, - block_v: int, -) -> tuple[int, int]: - return (kv_cache.shape[0] * _ceil_div(v_head_dim, block_v) * block_v, page_size // 2) - - -def _get_fp4_mla_v_packed_pool(metadata: Any, local_layer: int) -> Optional[torch.Tensor]: - return metadata.kv_cache_manager.get_mla_v_packed_pool(local_layer) - - -def _get_fp4_mla_v_packed_pool_base(metadata: Any) -> Optional[torch.Tensor]: - return metadata.kv_cache_manager.get_mla_v_packed_pool_base() - - -def _get_fp4_mla_v_scale_pool_base(metadata: Any) -> Optional[torch.Tensor]: - return metadata.kv_cache_manager.get_mla_v_scale_pool_base() - - -def _get_cutedsl_persistent_v_packed_cache( - metadata: Any, - local_layer: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - block_v: int, -) -> torch.Tensor: - v_packed = _get_fp4_mla_v_packed_pool(metadata, local_layer) - if v_packed is None: - raise RuntimeError( - "CuTeDSL FP4 MLA requires the manager-owned persistent V-packed " - "pool; the scratch full-repack fallback has been removed." - ) - expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) - if ( - not isinstance(v_packed, torch.Tensor) - or v_packed.dtype != torch.uint8 - or v_packed.device != kv_cache.device - or tuple(v_packed.shape) != expected_shape - or not v_packed.is_contiguous() - ): - raise RuntimeError( - "FP4 MLA persistent V-packed pool must be a contiguous uint8 tensor " - f"with shape {expected_shape} on {kv_cache.device}; got " - f"{type(v_packed).__name__}, " - f"shape={getattr(v_packed, 'shape', None)}, " - f"dtype={getattr(v_packed, 'dtype', None)}, " - f"device={getattr(v_packed, 'device', None)}." - ) - return v_packed - - -def _repack_cutedsl_v_packed_cache( - v_packed: torch.Tensor, - kv_cache: torch.Tensor, - page_ids: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - block_v: int, - page_indptr: Optional[torch.Tensor] = None, - kv_lens: Optional[torch.Tensor] = None, - generation_lens: Optional[torch.Tensor] = None, - max_touched_pages: int = 1, -) -> None: - if page_ids.numel() == 0: - return - from .fp4_mla_cutedsl_v_repack import fp4_mla_repack_v_cache - - fp4_mla_repack_v_cache( - v_packed, - kv_cache, - page_ids, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - page_indptr=page_indptr, - kv_lens=kv_lens, - generation_lens=generation_lens, - max_touched_pages=max_touched_pages, - ) - - -def _v_packed_cache_tag( - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> tuple[Any, ...]: - v_sf_tag = ( - None - if v_sf is None - else ( - int(v_sf.data_ptr()), - str(v_sf.device), - str(v_sf.dtype), - tuple(int(dim) for dim in v_sf.shape), - tuple(int(stride) for stride in v_sf.stride()), - ) - ) - page_ids_tag = ( - None - if page_ids is None - else ( - int(page_ids.data_ptr()), - str(page_ids.device), - str(page_ids.dtype), - tuple(int(dim) for dim in page_ids.shape), - tuple(int(stride) for stride in page_ids.stride()), - ) - ) - return ( - int(layer_idx), - None if local_layer is None else int(local_layer), - int(kv_cache.data_ptr()), - str(kv_cache.device), - str(kv_cache.dtype), - tuple(int(dim) for dim in kv_cache.shape), - tuple(int(stride) for stride in kv_cache.stride()), - int(v_head_dim), - int(page_size), - int(block_v), - v_sf_tag, - page_ids_tag, - ) - - -def _triton_prepack_v_enabled() -> bool: - if _fp4_mla_attention_backend() != "triton": - return False - default = _env_enabled_default("TRTLLM_FP4_MLA_PREPACK_V", True) - return _env_enabled_default("TRTLLM_FP4_MLA_TRITON_PREPACK_V", default) - - -def _triton_can_prepack_v(v_head_dim: int, page_size: int, block_v: int) -> bool: - return ( - _triton_prepack_v_enabled() - and hasattr(tl, "make_tensor_descriptor") - and block_v in (32, 128) - and v_head_dim % block_v == 0 - and page_size == FP4_MLA_TOKENS_PER_BLOCK - ) - - -def _triton_v_packed_attr(layer_idx: int) -> str: - if _shared_v_pack_storage_enabled(): - return "_fp4_mla_triton_attention_v_packed_buf" - return f"_fp4_mla_triton_attention_v_packed_buf_l{layer_idx}" - - -def _triton_v_packed_valid_attr(layer_idx: int) -> str: - return f"_fp4_mla_triton_attention_v_packed_valid_l{layer_idx}" - - -def _triton_shared_v_packed_valid_attr() -> str: - return "_fp4_mla_triton_attention_v_packed_valid_tag" - - -def _triton_v_packed_cache_tag( - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> tuple[Any, ...]: - return ( - "triton", - _v_packed_cache_tag( - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ), - ) - - -def _set_triton_v_packed_cache_valid( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> None: - valid_attr = ( - _triton_shared_v_packed_valid_attr() - if _shared_v_pack_storage_enabled() - else _triton_v_packed_valid_attr(layer_idx) - ) - metadata.fp4_mla_state.v_packed_cache_tags[valid_attr] = _triton_v_packed_cache_tag( - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ) - - -def _is_triton_v_packed_cache_valid( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> bool: - valid_attr = ( - _triton_shared_v_packed_valid_attr() - if _shared_v_pack_storage_enabled() - else _triton_v_packed_valid_attr(layer_idx) - ) - return metadata.fp4_mla_state.v_packed_cache_tags.get(valid_attr) == _triton_v_packed_cache_tag( - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ) - - -def _get_triton_v_packed_cache( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> Optional[torch.Tensor]: - if not _triton_can_prepack_v(v_head_dim, page_size, block_v): - return None - if not _is_triton_v_packed_cache_valid( - metadata, - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ): - return None - v_packed = metadata.fp4_mla_state.workspaces.get(_triton_v_packed_attr(layer_idx)) - expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) - if ( - v_packed is None - or v_packed.dtype != torch.uint8 - or v_packed.device != kv_cache.device - or len(v_packed.shape) != 2 - or v_packed.shape[0] < expected_shape[0] - or v_packed.shape[1] < expected_shape[1] - ): - return None - return v_packed[: expected_shape[0], : expected_shape[1]] - - -def _update_triton_v_packed_cache( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - page_ids: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - block_v: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - num_valid_pages: Optional[torch.Tensor] = None, -) -> Optional[torch.Tensor]: - if not _triton_can_prepack_v(v_head_dim, page_size, block_v): - return None - if page_ids.numel() == 0: - return None - from .fp4_mla_triton import fp4_mla_repack_v_cache_triton - - def _tma_alloc(size: int, alignment: int, stream): - return torch.empty(size, device=kv_cache.device, dtype=torch.int8) - - triton.set_allocator(_tma_alloc) - attr_name = _triton_v_packed_attr(layer_idx) - v_packed = _ensure_workspace_tensor( - metadata, - attr_name, - _v_packed_shape(kv_cache, v_head_dim, page_size, block_v), - dtype=torch.uint8, - device=kv_cache.device, - ) - fp4_mla_repack_v_cache_triton( - v_packed, - kv_cache, - page_ids, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - num_valid_pages=num_valid_pages, - ) - _set_triton_v_packed_cache_valid( - metadata, - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ) - return v_packed - - -def _maybe_update_triton_v_packed_cache( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - page_ids: torch.Tensor, - *, - num_queries: int, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - num_valid_pages: Optional[torch.Tensor] = None, -) -> None: - block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) - _update_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - page_ids, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - num_valid_pages=num_valid_pages, - ) - - -def _max_generation_pages(metadata: Any) -> int: - num_gen = metadata.num_seqs - metadata.num_contexts - if num_gen <= 0: - return 0 - if not getattr(metadata.fp4_mla_state, "device_page_table", False): - raise RuntimeError("FP4 MLA generation requires fixed-stride device page metadata.") - max_pages = int(metadata.fp4_mla_state.page_table_stride) - if max_pages <= 0: - raise RuntimeError("FP4 MLA device page-table stride must be positive.") - return max_pages - - -def _fp4_mla_generation_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: - """Return the fixed-stride generation page-table view.""" - expected_num_gen = metadata.num_seqs - metadata.num_contexts - if num_gen_seqs != expected_num_gen: - raise RuntimeError( - "FP4 MLA generation sequence count does not match metadata: " - f"{num_gen_seqs} != {expected_num_gen}." - ) - max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) - start = metadata.num_contexts * max_pages - end = start + num_gen_seqs * max_pages - if ( - not isinstance(page_ids, torch.Tensor) - or page_ids.ndim != 1 - or page_ids.dtype != torch.int32 - or not page_ids.is_contiguous() - or page_ids.numel() < end - ): - raise RuntimeError("FP4 MLA fixed-stride generation page-table backing is invalid.") - return page_ids[start:end] - - -def _fp4_mla_generation_hp_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: - """Return generation rows from the fixed-stride V2 HP page table.""" - expected_num_gen = metadata.num_seqs - metadata.num_contexts - if num_gen_seqs != expected_num_gen: - raise RuntimeError( - "FP4 MLA generation sequence count does not match metadata: " - f"{num_gen_seqs} != {expected_num_gen}." - ) - max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) - start = metadata.num_contexts * max_pages - end = start + num_gen_seqs * max_pages - if ( - not isinstance(page_ids, torch.Tensor) - or page_ids.ndim != 1 - or page_ids.dtype != torch.int32 - or not page_ids.is_contiguous() - or page_ids.numel() < end - ): - raise RuntimeError("FP4 MLA fixed-stride generation HP page-table backing is invalid.") - return page_ids[start:end] - - -def _host_int_list(value: Any, start: int, end: int) -> Optional[list[int]]: - if value is None: - return None - if isinstance(value, torch.Tensor): - if value.is_cuda: - return None - return [int(item) for item in value[start:end].tolist()] - try: - return [int(item) for item in value[start:end]] - except (TypeError, ValueError): - return None - - -def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> bool: - if getattr(metadata, "is_cuda_graph", False): - return False - - start = metadata.num_contexts - end = metadata.num_seqs - block_counts = _host_int_list(getattr(metadata.fp4_mla_state, "num_blocks", None), start, end) - if block_counts is not None and ( - not block_counts or min(block_counts) != max_pages or max(block_counts) != max_pages - ): - return False - - kv_lens_cuda = getattr(metadata, "kv_lens_cuda_runtime", None) - if isinstance(kv_lens_cuda, torch.Tensor): - cache_key = ( - start, - end, - max_pages, - page_size, - tuple(block_counts) if block_counts is not None else None, - kv_lens_cuda.data_ptr(), - ) - cache = getattr(metadata.fp4_mla_state, "full_pages_cache", None) - if cache is not None and cache[0] == cache_key: - return bool(cache[1]) - kv_lens = [int(item) for item in kv_lens_cuda[start:end].detach().cpu().tolist()] - result = bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size - setattr(metadata.fp4_mla_state, "full_pages_cache", (cache_key, result)) - return result - - kv_cache_params = getattr(metadata, "kv_cache_params", None) - cached_token_lens = _host_int_list( - getattr(kv_cache_params, "num_cached_tokens_per_seq", None), - start, - end, - ) - seq_lens_kv = _host_int_list(getattr(metadata, "seq_lens_kv", None), start, end) - if cached_token_lens is not None and seq_lens_kv is not None: - if len(cached_token_lens) != len(seq_lens_kv): - return False - kv_lens = [ - cached_len + seq_len for cached_len, seq_len in zip(cached_token_lens, seq_lens_kv) - ] - elif kv_cache_params is None: - kv_lens = _host_int_list(getattr(metadata, "prompt_lens_cpu_runtime", None), start, end) - else: - return False - - return bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size - - -def _get_linear_mtp_query_len_per_seq( - metadata: Any, - *, - num_queries: int, - num_gen_seqs: int, -) -> int: - """Return the uniform generation query length required by linear MTP. - - Derives the length from the real query-token count (``num_queries``, taken - from the q shape) and the generation sequence count, which are reliable in - every representation. The host ``prompt_lens``/``seq_lens`` mirror can lag at - the decode anchor (== 1) under CUDA graph / one-engine MTP, so it is only - consulted to produce a precise diagnostic when the counts do not divide - evenly (a genuinely non-uniform batch, which the no-dequant path does not - support). - """ - if num_gen_seqs <= 0: - return 1 - - if num_queries % num_gen_seqs == 0: - return num_queries // num_gen_seqs - - start = metadata.num_contexts - end = metadata.num_seqs - query_lens = _host_int_list_during_forward( - getattr(metadata, "prompt_lens_cpu_runtime", None), start, end - ) - if query_lens is None: - query_lens = _host_int_list_during_forward(getattr(metadata, "seq_lens", None), start, end) - raise NotImplementedError( - "FP4 MLA no-dequant attention requires a uniform linear MTP generation " - f"query length; got {num_queries} query tokens for {num_gen_seqs} " - f"sequences (per-sequence lengths {query_lens})." - ) - - -def _run_triton_attention_decode( - *, - metadata: Any, - layer_idx: int, - local_layer: int, - q_fp4: torch.Tensor, - q_sf: torch.Tensor, - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - v_sf: torch.Tensor, - global_scale: torch.Tensor, - src_page_ids: torch.Tensor, - kv_lens: torch.Tensor, - p_fp4: torch.Tensor, - p_sf: torch.Tensor, - max_scores: torch.Tensor, - denom: torch.Tensor, - output: torch.Tensor, - num_queries: int, - num_heads: int, - head_dim: int, - kv_lora_rank: int, - q_residual_dim: int, - query_len_per_seq: int, - max_pages: int, - sm_scale: float, - q_global_scale: torch.Tensor, -) -> None: - """Dispatch the ``triton`` FP4 MLA decode pipeline. - - Mirrors the four-stage layout used by ``fp4_mla_cutile.py`` - (page-stats with packed P -> reduce-stats -> prob-scale -> PV) but - routes through the self-contained kernels in - ``fp4_mla_triton.py``. Threads through the constexpr assume flags, - TMA descriptors, occupancy/num-warps launch meta, and pipelined PV loop. - """ - from .fp4_mla_triton import ( - _fp4_mla_attention_group_reduce_stats_kernel as _attn_group_reduce_stats_kernel, - ) - from .fp4_mla_triton import _fp4_mla_attention_page_stats_kernel as _attn_page_stats_kernel - from .fp4_mla_triton import _fp4_mla_attention_prob_scale_kernel as _attn_prob_scale_kernel - from .fp4_mla_triton import _fp4_mla_attention_pv_kernel as _attn_pv_kernel - from .fp4_mla_triton import ( - _fp4_mla_attention_pv_prepacked_v_kernel as _attn_pv_prepacked_v_kernel, - ) - from .fp4_mla_triton import _fp4_mla_attention_pv_reduce_kernel as _attn_pv_reduce_kernel - from .fp4_mla_triton import _fp4_mla_attention_reduce_stats_kernel as _attn_reduce_stats_kernel - - block_h = 128 - block_t = metadata.page_size - # Adaptive BLOCK_V: the fallback PV path uses a finer V split at small batch - # on B200 (~148 SMs). PV grid = num_queries * num_head_blocks(1) * - # (kv_lora_rank / BLOCK_V). We want >= ~2*num_SMs programs so that >1 CTA - # lands per SM and hides the L1TEX scoreboard stalls. Empirically (sweep): - # bs<=32 -> BLOCK_V=32; bs>=64 -> BLOCK_V=128. - # (BLOCK_V=16 is rejected by the V TMA descriptor min-stride requirement.) - # With prepacked V, BLOCK_V=128 avoids reloading the same P tile four times - # and matches the cutile prepacked-V tile shape. - block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) - q_storage_head_dim = head_dim + q_residual_dim - # The virtual GEMM tail evaluates QK + Q_r K + Q K_r in one reduction. - # Q and Q_r still occupy the 640-channel interleaved physical Q buffer; - # the final Q term reuses Q's main tail groups while K_r comes from the - # contiguous 64-channel tail of the primary paged KV cache. - q_head_dim = head_dim + q_residual_dim + FP4_MLA_K_RESIDUAL_DIM - # BLOCK_K = 512 aligns the K-window with the 512-channel non-residual prefix. - block_k = 512 - full_block_end = (q_head_dim // block_k) * block_k - tail_k = q_head_dim - full_block_end - tail_block_k = 1 << (tail_k - 1).bit_length() if tail_k > 0 else block_k - q_sf_per_token = q_storage_head_dim // FP4_BLOCK_SIZE - k_sf_per_token = (head_dim + FP4_MLA_K_RESIDUAL_DIM) // FP4_BLOCK_SIZE - sf_per_page = metadata.page_size // FP4_BLOCK_SIZE - num_head_blocks = triton.cdiv(num_heads, block_h) - - assume_full_heads = num_heads % block_h == 0 - assume_full_v = kv_lora_rank % block_v == 0 - # Match the cutile path: only mark pages "full" when we can prove every - # generation sequence has the same number of cached tokens AND - # query_len_per_seq == 1 (so the kv_len adjustment is a no-op). - assume_full_pages = ( - _infer_assume_full_pages(metadata, max_pages, metadata.page_size) and query_len_per_seq == 1 - ) - # Leave validity checks on. Matches cutile's default and is correctness- - # safe. The perfect-shape PV fast path (tl.ext.make_view + load_view_tko) - # remains gated off — when measured on the TileIR backend (ENABLE_TILE=1) - # it was net-slower on the bench, so the cost of enabling it isn't worth - # the win on the FP4 MLA shapes we care about. - assume_valid_pages = False - num_gen_seqs = num_queries // query_len_per_seq - if ( - not assume_valid_pages - and assume_full_pages - and src_page_ids.numel() == num_gen_seqs * max_pages - ): - assume_valid_pages = True - # cutile checks only `make_tensor_descriptor`; on the nvt backend the - # presence of TMA descriptors implies `tl.ext.make_view` is available too. - use_tma_data_load = hasattr(triton.language, "make_tensor_descriptor") - - # Install the device-side scratch allocator on every call. Triton stores - # the allocator in a ContextVar (triton.runtime._allocation), so a single - # process-wide install is not visible from worker threads / asyncio tasks - # that run with a different Context — the kernel launch would then hit the - # default NullAllocator and raise. Matches the cutile path. - if use_tma_data_load: - - def _tma_alloc(size: int, alignment: int, stream): - return torch.empty(size, device=q_fp4.device, dtype=torch.int8) - - triton.set_allocator(_tma_alloc) - - # cutile-equivalent launch meta. occupancy=2 lets two CTAs land per SM - # which improves wave-tail efficiency at the bs=32 hot point. - # NOTE: num_stages=2 (instead of the Triton 3.6 default of 3) sidesteps - # the TritonGPUAutomaticWarpSpecialization + NVWSInsertTmemAref pass that - # ICEs on the page_stats kernel under Triton 3.6.0 / sm_100. - launch_meta = {"occupancy": 2} - # The matmul kernels (page-stats QK and PV) are register-limited: at the - # Triton default of num_warps=4 the [BLOCK_H, BLOCK_T] epilogue spills the - # register file down to ~2 CTAs/SM (12.5% occupancy), so there are too few - # warps to hide the QK/PV load latency (ncu: ~0.3 eligible warps/scheduler). - # Spreading the tile epilogue over num_warps=8 halves the per-thread - # register need and roughly doubles resident warps. Matches the cutile - # ("nvt") backend, which launches page-stats at num_warps=8. Both are - # overridable for tuning. - sm_count = _get_sm_count(q_fp4.device) - # page-stats num_warps: the full-pages fast path (uniform q_len==1 decode) - # benefits from num_warps=8 (more warps hide the QK load latency); the - # masked path (q_len>1 / ragged lengths) carries extra per-thread state and - # measured markedly faster at num_warps=4 (e.g. bs256 q_len4: 131->95ms). - page_stats_num_warps = _env_int("TRTLLM_FP4_MLA_PAGE_STATS_NUM_WARPS") - if page_stats_num_warps is None: - page_stats_num_warps = 8 if assume_full_pages else 4 - page_stats_launch_meta = {"occupancy": 2, "num_warps": page_stats_num_warps} - # PV benefits from num_warps=8 across shapes measured. - pv_num_warps = _env_int("TRTLLM_FP4_MLA_PV_NUM_WARPS") or 8 - pv_launch_meta = {"occupancy": 2, "num_warps": pv_num_warps} - # PV loop pipelining. With TMA loads, num_stages>=2 lets the next page's - # loads overlap with the current MMA via mbarrier. The PV report shows - # long_scoreboard=4.5 cycles avg on V loads at PV_LOOP_STAGES=2; bumping the - # depth pays off when the grid is small enough that occupancy can absorb - # the extra in-flight tile state — i.e. medium batch / large max_pages. - # Larger pipelines hurt at small batch (more live state, fewer dim blocks). - if num_queries <= 16 or max_pages <= 4: - pv_loop_stages = 2 - else: - pv_loop_stages = 3 - - # Page-stats kernel: per (query, head_block, page) program, does QK, - # softmax stats, and packs probs into FP4 with the per-page local-max - # scaling trick. The page-max correction is applied later by - # prob_scale_kernel via p_sf in-place rescaling. - page_stats_shape = (num_queries, max_pages, num_heads) - page_max = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_page_max_buf", - page_stats_shape, - dtype=torch.float32, - device=q_fp4.device, - ) - page_sum = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_page_sum_buf", - page_stats_shape, - dtype=torch.float32, - device=q_fp4.device, - ) - - pack_prob_in_page_stats = True - _attn_page_stats_kernel[(num_queries, num_head_blocks, max_pages)]( - page_max, - page_sum, - p_fp4, - p_sf, - q_fp4, - q_sf, - kv_cache, - sf_cache, - global_scale, - q_global_scale, - src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - q_fp4.stride(0), - q_fp4.stride(1), - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - page_max.stride(0), - page_max.stride(1), - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - q_fp4.shape[0], - sm_scale, - NUM_HEADS=num_heads, - Q_HEAD_D=q_head_dim, - Q_STORAGE_HEAD_D=q_storage_head_dim, - K_HEAD_D=head_dim, - Q_RESIDUAL_D=q_residual_dim, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - Q_SF_PER_TOKEN=q_sf_per_token, - K_SF_PER_TOKEN=k_sf_per_token, - SF_PER_PAGE=sf_per_page, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - BLOCK_H=block_h, - BLOCK_T=block_t, - BLOCK_K=block_k, - FULL_BLOCK_END=full_block_end, - TAIL_BLOCK_K=tail_block_k, - USE_TMA_DATA_LOAD=use_tma_data_load, - PACK_PROBS=pack_prob_in_page_stats, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_VALID_PAGES=assume_valid_pages, - **page_stats_launch_meta, - ) - # Two-level softmax-stats reduction. The single-level reduce launched only - # (num_queries * num_head_blocks) CTAs, each serially walking all max_pages - # twice -- at small batch that handful of CTAs left the GPU almost idle and - # the reduce cost more than the QK matmul. Level 1 parallelizes the page - # reduction across a page-group axis (online-softmax partials, pipelined); - # level 2 reuses the existing reduce kernel to fold the few groups into the - # global (max, denom). When the (query, head) grid already fills the GPU the - # group count collapses to 1 and this degenerates to the original reduce. - seqhead_ctas = num_queries * num_head_blocks - # Aim for ~3 waves of level-1 CTAs so page loads have enough memory-level - # parallelism to hide latency, while keeping the group count small enough - # that the level-2 combine loop stays short. - target_l1_ctas = 3 * sm_count - num_reduce_groups = _ceil_div(target_l1_ctas, max(seqhead_ctas, 1)) - num_reduce_groups = max(1, min(num_reduce_groups, max_pages, 64)) - # The grouped (two-level) reduce needs an auxiliary workspace, and - # _ensure_workspace_tensor can only (re)allocate it outside CUDA graph - # capture. If a warmup forward did not already size that workspace (e.g. the - # warmup batch took the single-level path), fall back to the single-level - # reduce during capture so we never allocate mid-capture. The single-level - # reduce is numerically identical (it just launches fewer CTAs). - if num_reduce_groups > 1 and torch.cuda.is_current_stream_capturing(): - gmax = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_max_buf") - gsum = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_sum_buf") - groups_ready = ( - gmax is not None - and gsum is not None - and gmax.shape[0] >= num_queries - and gmax.shape[1] >= num_reduce_groups - and gmax.shape[2] >= num_heads - and gsum.shape[0] >= num_queries - and gsum.shape[1] >= num_reduce_groups - and gsum.shape[2] >= num_heads - ) - if not groups_ready: - num_reduce_groups = 1 - if num_reduce_groups <= 1: - _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( - max_scores, - denom, - page_max, - page_sum, - max_pages, - max_scores.stride(0), - page_max.stride(0), - page_max.stride(1), - NUM_HEADS=num_heads, - MAX_PAGES=max_pages, - BLOCK_H=block_h, - **launch_meta, - ) - else: - group_pages = _ceil_div(max_pages, num_reduce_groups) - num_reduce_groups = _ceil_div(max_pages, group_pages) - group_max = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_group_max_buf", - (num_queries, num_reduce_groups, num_heads), - dtype=torch.float32, - device=q_fp4.device, - ) - group_sum = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_group_sum_buf", - (num_queries, num_reduce_groups, num_heads), - dtype=torch.float32, - device=q_fp4.device, - ) - _attn_group_reduce_stats_kernel[(num_queries, num_head_blocks, num_reduce_groups)]( - group_max, - group_sum, - page_max, - page_sum, - max_pages, - group_max.stride(0), - group_max.stride(1), - page_max.stride(0), - page_max.stride(1), - NUM_HEADS=num_heads, - GROUP_PAGES=group_pages, - BLOCK_H=block_h, - PIPELINE_STAGES=min(group_pages, 4), - **launch_meta, - ) - _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( - max_scores, - denom, - group_max, - group_sum, - num_reduce_groups, - max_scores.stride(0), - group_max.stride(0), - group_max.stride(1), - NUM_HEADS=num_heads, - MAX_PAGES=num_reduce_groups, - BLOCK_H=block_h, - **launch_meta, - ) - _attn_prob_scale_kernel[(num_queries, num_head_blocks, max_pages)]( - p_sf, - max_scores, - denom, - page_max, - metadata.fp4_mla_state.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - max_scores.stride(0), - page_max.stride(0), - page_max.stride(1), - NUM_HEADS=num_heads, - PAGE_SIZE=metadata.page_size, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - BLOCK_H=block_h, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_VALID_PAGES=assume_valid_pages, - **launch_meta, - ) - num_dim_blocks = triton.cdiv(kv_lora_rank, block_v) - v_packed = _get_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - v_head_dim=kv_lora_rank, - page_size=metadata.page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=src_page_ids, - ) - if ( - v_packed is None - and _triton_can_prepack_v(kv_lora_rank, metadata.page_size, block_v) - and not torch.cuda.is_current_stream_capturing() - ): - v_packed = _update_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - src_page_ids, - v_head_dim=kv_lora_rank, - page_size=metadata.page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - ) - use_triton_v_packed_cache = v_packed is not None - - # PV page split: partition the page range across additional programs and - # reduce in a follow-up kernel. ncu showed PV at waves/SM=0.49 for bs=32 — - # PV is L1-bandwidth bound, so raising in-flight CTAs is the lever. - # BLOCK_V is bounded below by the 16-byte TMA descriptor min-stride. - # PV page split: ncu shows that with the current shape (bs=32, max_pages=256) - # the PV kernel is L1-cache-throughput bound (long_scoreboard=4.5 cycles - # avg, L1 global LD hit-rate <40%). Increasing the program count via page - # splitting reduced waves/SM idle time but did NOT improve wall-time at - # current shapes — the per-CTA L1 thrash is the limit. Gate the split off - # by default; re-enable only for very small grids where occupancy is the - # bottleneck rather than per-CTA L1 pressure. - page_split = 1 - base_grid = num_queries * num_head_blocks * num_dim_blocks - if max_pages >= 16 and base_grid < 148: - for p in (8, 4, 2): - if max_pages % p == 0 and max_pages // p >= 16 and base_grid * p <= 148 * 4: - page_split = p - break - # The page-split PV path needs a partial-output workspace, which - # _ensure_workspace_tensor can only (re)allocate outside CUDA graph capture. - # Fall back to the unsplit PV (numerically identical) during capture unless a - # warmup forward already sized that workspace, so capture never allocates. - if page_split > 1 and torch.cuda.is_current_stream_capturing(): - pbuf = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_pv_partial_buf") - partial_ready = ( - pbuf is not None - and pbuf.shape[0] >= num_queries - and pbuf.shape[1] >= page_split - and pbuf.shape[2] >= num_heads - and pbuf.shape[3] >= kv_lora_rank - ) - if not partial_ready: - page_split = 1 - if page_split > 1: - pages_per_split = max_pages // page_split - partial_out = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_pv_partial_buf", - (num_queries, page_split, num_heads, kv_lora_rank), - dtype=torch.float32, - device=q_fp4.device, - ) - if use_triton_v_packed_cache: - _attn_pv_prepacked_v_kernel[ - (num_queries, num_head_blocks, num_dim_blocks * page_split) - ]( - output, - p_fp4, - p_sf, - v_packed, - v_sf, - global_scale, - src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - PAGE_SPLIT=page_split, - PAGES_PER_SPLIT=pages_per_split, - PARTIAL_OUT=True, - partial_out_ptr=partial_out, - partial_s0=partial_out.stride(0), - partial_s1=partial_out.stride(1), - partial_s2=partial_out.stride(2), - partial_s3=partial_out.stride(3), - **pv_launch_meta, - ) - else: - _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks * page_split)]( - output, - p_fp4, - p_sf, - kv_cache, - kv_cache, - v_sf, - global_scale, - src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, - USE_PREPACKED_V=False, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - PAGE_SPLIT=page_split, - PAGES_PER_SPLIT=pages_per_split, - PARTIAL_OUT=True, - partial_out_ptr=partial_out, - partial_s0=partial_out.stride(0), - partial_s1=partial_out.stride(1), - partial_s2=partial_out.stride(2), - partial_s3=partial_out.stride(3), - **pv_launch_meta, - ) - _attn_pv_reduce_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( - output, - partial_out, - global_scale, - output.stride(0), - output.stride(1), - output.stride(2), - partial_out.stride(0), - partial_out.stride(1), - partial_out.stride(2), - partial_out.stride(3), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SPLIT=page_split, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_V=assume_full_v, - **launch_meta, - ) - else: - if use_triton_v_packed_cache: - _attn_pv_prepacked_v_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( - output, - p_fp4, - p_sf, - v_packed, - v_sf, - global_scale, - src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - **pv_launch_meta, - ) - else: - _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( - output, - p_fp4, - p_sf, - kv_cache, - kv_cache, - v_sf, - global_scale, - src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, - USE_PREPACKED_V=False, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - **pv_launch_meta, - ) - - -def run_fp4_mla_attention_decode( - metadata: Any, - layer_idx: int, - local_layer: int, - q: torch.Tensor, - output: torch.Tensor, - *, - sm_scale: float, - kv_lora_rank: int, - qk_rope_head_dim: int, - prequantized_q: torch.Tensor, - prequantized_q_sf: torch.Tensor, - q_batch_capacity: int, -) -> None: - """Run MLA decode with FP4 QK and FP4 PV tensor-core matmuls. - - Q is supplied in its assembled ``[latent, RoPE]`` layout and quantized to - FP4 directly. QK reads ``[KV-nope, K-RoPE, K-RoPE-residual]`` contiguously - from the primary cache with swizzled block scales. Softmax probabilities - are quantized to FP4 per page, and PV repacks V nibbles from the shared KV - cache while reading the auxiliary V-view scale pool. No BF16 dequantized - KV workspace is materialized on this path. Callers must supply the packed Q - and scales produced by the fused generation cache update. - """ - head_dim = kv_lora_rank + qk_rope_head_dim - if qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM: - raise ValueError( - "FP4 MLA K residual attention requires " - f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got {qk_rope_head_dim}." - ) - _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) - if metadata.page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - f"FP4 MLA attention decode requires page_size={FP4_MLA_TOKENS_PER_BLOCK}, " - f"got {metadata.page_size}." - ) - - if q.ndim != 3 or q.shape[-1] != head_dim: - raise ValueError( - "FP4 MLA attention Q must have shape " - f"[tokens, heads, {head_dim}], got {tuple(q.shape)}." - ) - if not q.is_contiguous(): - raise ValueError("FP4 MLA attention Q must be contiguous.") - - num_queries = q.shape[0] - if num_queries == 0: - raise ValueError("FP4 MLA attention decode requires at least one query token.") - num_gen_seqs = metadata.num_seqs - metadata.num_contexts - query_len_per_seq = _get_linear_mtp_query_len_per_seq( - metadata, - num_queries=num_queries, - num_gen_seqs=num_gen_seqs, - ) - - num_heads = q.shape[1] - if output.shape[:2] != (num_queries, num_heads): - raise ValueError("FP4 MLA attention output batch dimensions do not match.") - - backend = _fp4_mla_attention_backend() - if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: - raise RuntimeError( - "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." - ) - - global_scale = _get_fp4_mla_global_scale(metadata, q.device) - q_residual_dim = FP4_MLA_Q_RESIDUAL_DIM - _validate_fp4_mla_attention_q_shape(head_dim, q_residual_dim) - - if prequantized_q is None or prequantized_q_sf is None or q_batch_capacity is None: - raise RuntimeError( - "FP4 MLA decode requires Q prequantized by the fused generation cache update." - ) - - capacity = int(q_batch_capacity) - expected_q_shape = (capacity * num_heads, FP4_MLA_Q_PACKED_DIM) - expected_q_sf_shape = ( - _get_fp4_mla_swizzled_scale_size( - capacity * num_heads, - FP4_MLA_Q_LOGICAL_DIM, - ), - ) - if ( - q.dtype != torch.bfloat16 - or capacity <= 0 - or num_queries > capacity - or tuple(prequantized_q.shape) != expected_q_shape - or prequantized_q.dtype != torch.uint8 - or not prequantized_q.is_contiguous() - or tuple(prequantized_q_sf.shape) != expected_q_sf_shape - or prequantized_q_sf.dtype != torch.float8_e4m3fn - or not prequantized_q_sf.is_contiguous() - ): - raise ValueError("FP4 MLA prequantized Q does not satisfy the fused-Q contract.") - active_q_rows = num_queries * num_heads - active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( - active_q_rows, - FP4_MLA_Q_LOGICAL_DIM, - ) - q_fp4 = prequantized_q[:active_q_rows] - q_sf = prequantized_q_sf[:active_q_sf_bytes] - q_global_scale = _get_fp4_mla_q_global_scale(metadata, q.device) - - kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) - _validate_fp4_mla_kv_storage_shape( - kv_cache, - sf_cache, - head_dim=head_dim, - backend=backend, - ) - sf_cache = sf_cache.view(torch.float8_e4m3fn) - - v_sf_pool = metadata.fp4_mla_state.v_scale_pool - v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=kv_lora_rank)[local_layer].view( - torch.float8_e4m3fn - ) - # The kv_lens runtime alias can lag at the decode anchor (seq_lens == 1) under - # CUDA graph / one-engine MTP; recover the true total per sequence so the - # per-query causal masking sees the full 1 + draft_len window (no-op when the - # alias already matches). - kv_lens, _ = _fp4_mla_uniform_generation_lengths(metadata, num_queries, num_gen_seqs) - _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens) - src_page_ids = _fp4_mla_generation_page_ids(metadata, num_gen_seqs) - max_pages = _max_generation_pages(metadata) - if max_pages == 0: - raise RuntimeError("FP4 MLA attention decode requires generation cache pages.") - if backend == _FP4_MLA_CUTEDSL_BACKEND: - if get_sm_version() != 107: - raise RuntimeError( - "FP4 MLA cutedsl attention backend requires Rubin SM107; " - f"current architecture is SM{get_sm_version()}." - ) - if not _cutedsl_backend_available(): - raise RuntimeError( - "FP4 MLA cutedsl attention backend requires the Rubin CTM and " - "CuTeDSL runtime packages." - ) - if not 0 < num_heads <= 128 or kv_lora_rank != 512: - raise ValueError( - "FP4 MLA cutedsl attention requires 1-128 local heads and " - f"kv_lora_rank=512, got num_heads={num_heads}, " - f"kv_lora_rank={kv_lora_rank}." - ) - - cutedsl_kernel = _fp4_mla_cutedsl_kernel_module() - QK_LOGICAL_DIM = cutedsl_kernel.QK_LOGICAL_DIM - QK_SF_GROUPS = cutedsl_kernel.QK_SF_GROUPS - SMEM_P4_V_N_PER_CTA = cutedsl_kernel.SMEM_P4_V_N_PER_CTA - run_trtllm_fp4_mla_decode_page_native_from_raw = ( - cutedsl_kernel.run_trtllm_fp4_mla_decode_page_native_from_raw - ) - - physical_heads = 128 - kernel_q = q_fp4 - kernel_q_sf = q_sf - if num_heads < physical_heads: - kernel_q_storage, kernel_q_sf_storage, q_batch_capacity = _prepare_fp4_mla_q_buffers( - metadata, - num_queries, - physical_heads, - q.device, - ) - active_q_rows = num_queries * physical_heads - active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( - active_q_rows, - QK_LOGICAL_DIM, - ) - kernel_q = kernel_q_storage[:active_q_rows] - kernel_q_sf = kernel_q_sf_storage[:active_q_sf_bytes] - _cutedsl_pad_q_and_sf_kernel[(num_queries, _ceil_div(QK_LOGICAL_DIM // 2, 64))]( - kernel_q, - q_fp4, - kernel_q_sf, - q_sf, - num_heads, - output_heads=physical_heads, - packed_dim=QK_LOGICAL_DIM // 2, - block_bytes=64, - sf_cols=QK_SF_GROUPS, - sf_cols_per_byte_block=8, - ) - - fused_v_transpose = _fp4_mla_cutedsl_fused_v_transpose_enabled() - if fused_v_transpose: - # The fusion kernel reads V from the canonical KV cache and uses - # the current layer V scales directly. Keep a None placeholder so - # the mufu16 and fused-V launchers share one Python call site. - core_v_packed = None - core_v_sf = v_sf - v_page_offset = 0 - else: - v_packed = _get_cutedsl_persistent_v_packed_cache( - metadata, - local_layer, - kv_cache, - v_head_dim=kv_lora_rank, - page_size=metadata.page_size, - block_v=SMEM_P4_V_N_PER_CTA, - ) - core_v_packed = _get_fp4_mla_v_packed_pool_base(metadata) - if core_v_packed is None: - raise RuntimeError("Persistent FP4 MLA V packing requires a stable full-pool base.") - get_v_page_offset = getattr( - metadata.kv_cache_manager, "get_mla_v_packed_page_offset", None - ) - v_page_offset = ( - int(get_v_page_offset(local_layer)) - if callable(get_v_page_offset) - else local_layer * kv_cache.shape[0] - ) - page_bytes = kv_lora_rank * (metadata.page_size // 2) - expected_layer_ptr = core_v_packed.data_ptr() + v_page_offset * page_bytes - if expected_layer_ptr != v_packed.data_ptr(): - raise RuntimeError( - "Persistent FP4 MLA V-packed layer view does not match its " - "full-pool base and page offset." - ) - - v_sf_pool_base = _get_fp4_mla_v_scale_pool_base(metadata) - if v_sf_pool_base is None: - v_sf_pool_base = v_sf_pool.flatten(0, 1).view(torch.uint8) - if v_sf_pool_base.data_ptr() != v_sf_pool.data_ptr(): - raise RuntimeError( - "Persistent FP4 MLA V-scale pool must flatten without a copy." - ) - if ( - not isinstance(v_sf_pool_base, torch.Tensor) - or v_sf_pool_base.dtype != torch.uint8 - or v_sf_pool_base.device != v_sf.device - or v_sf_pool_base.ndim != 2 - or v_sf_pool_base.shape[1] != v_sf_pool.shape[-1] - or not v_sf_pool_base.is_contiguous() - ): - raise RuntimeError( - "Persistent FP4 MLA V-scale pool base must be a contiguous " - "two-dimensional uint8 tensor with the configured page stride." - ) - core_v_sf = v_sf_pool_base.view(torch.float8_e4m3fn) - get_v_sf_page_offset = getattr( - metadata.kv_cache_manager, "get_mla_v_scale_page_offset", None - ) - v_sf_page_offset = ( - int(get_v_sf_page_offset(local_layer)) - if callable(get_v_sf_page_offset) - else local_layer * kv_cache.shape[0] - ) - if v_sf_page_offset != v_page_offset: - raise RuntimeError( - "Persistent FP4 MLA V-packed and V-scale pools require " - "matching encoded layer offsets." - ) - expected_v_sf_ptr = ( - core_v_sf.data_ptr() - + v_sf_page_offset * v_sf_pool.stride(1) * v_sf_pool.element_size() - ) - if expected_v_sf_ptr != v_sf.data_ptr(): - raise RuntimeError( - "Persistent FP4 MLA V-scale layer view does not match its " - "full-pool base and page offset." - ) - - kernel_output = output - if num_heads < physical_heads: - kernel_output = _ensure_workspace_tensor( - metadata, - "_fp4_mla_cutedsl_output_buf", - (num_queries, physical_heads, kv_lora_rank), - dtype=output.dtype, - device=output.device, - ) - - run_trtllm_fp4_mla_decode_page_native_from_raw( - kernel_q, - kernel_q_sf, - kv_cache, - sf_cache, - core_v_packed, - core_v_sf, - global_scale, - src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], - kv_lens, - kernel_output, - max_kv_len=max_pages * metadata.page_size, - sm_scale=float(sm_scale), - num_heads=physical_heads, - q_global_scale=q_global_scale, - page_size=metadata.page_size, - query_len_per_seq=query_len_per_seq, - v_page_offset=v_page_offset, - q_batch_capacity=q_batch_capacity, - partition_runtime_valid_k=bool(getattr(metadata, "is_cuda_graph", False)), - ) - if kernel_output is not output: - output.copy_(kernel_output[:, :num_heads]) - return - - total_p_rows = num_queries * max_pages * num_heads - p_fp4 = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_p_buf", - (max(total_p_rows, 1), metadata.page_size // 2), - dtype=torch.uint8, - device=q.device, - )[:total_p_rows] - p_sf = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_p_sf_buf", - (max(_get_fp4_mla_swizzled_scale_size(total_p_rows, metadata.page_size), 1),), - dtype=torch.float8_e4m3fn, - device=q.device, - ) - stats_shape = (num_queries, num_heads) - max_scores = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_max_buf", - stats_shape, - dtype=torch.float32, - device=q.device, - ) - denom = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_denom_buf", - stats_shape, - dtype=torch.float32, - device=q.device, - ) - - if backend != "triton": - raise ValueError( - f"Unsupported FP4 MLA attention backend '{backend}'. " - f"Set {FP4_MLA_ATTENTION_BACKEND_ENV} to 'triton' or " - f"'{_FP4_MLA_CUTEDSL_BACKEND}'." - ) - - # Self-contained public-Triton path: TMA-loaded QK + fused page-stats pack, - # reduce-stats, prob-scale, and PV with an optional prepacked V cache. - _run_triton_attention_decode( - metadata=metadata, - layer_idx=layer_idx, - local_layer=local_layer, - q_fp4=q_fp4, - q_sf=q_sf.contiguous().view(-1), - kv_cache=kv_cache, - sf_cache=sf_cache, - v_sf=v_sf, - global_scale=global_scale, - src_page_ids=src_page_ids, - kv_lens=kv_lens, - p_fp4=p_fp4, - p_sf=p_sf, - max_scores=max_scores, - denom=denom, - output=output, - num_queries=num_queries, - num_heads=num_heads, - head_dim=head_dim, - kv_lora_rank=kv_lora_rank, - q_residual_dim=q_residual_dim, - query_len_per_seq=query_len_per_seq, - max_pages=max_pages, - sm_scale=float(sm_scale), - q_global_scale=q_global_scale, - ) +from .cache_update import can_fuse_fp4_mla_q_quant, scatter_fp4_mla_kv_cache +from .config import _FP4_MLA_CUTEDSL_BACKEND as _FP4_MLA_CUTEDSL_BACKEND +from .config import _FP4_MLA_K_RESIDUAL_BACKENDS as _FP4_MLA_K_RESIDUAL_BACKENDS +from .config import _FP4_MLA_MAX_GRID_Z as _FP4_MLA_MAX_GRID_Z +from .config import _FP4_MLA_PAGE_TABLE_TILE_SIZE as _FP4_MLA_PAGE_TABLE_TILE_SIZE +from .config import _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH as _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH +from .config import _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH as _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH +from .config import _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH as _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH +from .config import _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD as _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD +from .config import _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD as _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD +from .config import ( + _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD as _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD, +) +from .config import ( + _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD as _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD, +) +from .config import _FP4_MLA_TRITON_PRELOAD_KEYS as _FP4_MLA_TRITON_PRELOAD_KEYS +from .config import ( + FP4_BLOCK_SIZE, + FP4_MLA_ATTENTION_BACKEND_ENV, + FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, + FP4_MLA_E4M3_MAX, + FP4_MLA_K_RESIDUAL_DIM, + FP4_MLA_KV_GLOBAL_SCALE, + FP4_MLA_KV_STATIC_AMAX, + FP4_MLA_P_GLOBAL_SCALE, + FP4_MLA_Q1_PREFIX_BLOCK_DIM, + FP4_MLA_Q_GLOBAL_SCALE, + FP4_MLA_Q_LOGICAL_DIM, + FP4_MLA_Q_PACKED_DIM, + FP4_MLA_Q_PREFIX_BLOCK_DIM, + FP4_MLA_Q_PREFIX_DIM, + FP4_MLA_Q_RESIDUAL_DIM, + FP4_MLA_Q_SF_GROUPS, + FP4_MLA_Q_STATIC_AMAX, + FP4_MLA_SCALE_COL_GROUP, + FP4_MLA_SCALE_ROW_GROUP, + FP4_MLA_TOKENS_PER_BLOCK, + HP_BLOCK_SIZE, +) +from .config import _ceil_div as _ceil_div +from .config import _cutedsl_backend_available as _cutedsl_backend_available +from .config import _env_enabled_default as _env_enabled_default +from .config import _env_int as _env_int +from .config import _fp4_mla_attention_backend as _fp4_mla_attention_backend +from .config import ( + _fp4_mla_cutedsl_fused_v_transpose_enabled as _fp4_mla_cutedsl_fused_v_transpose_enabled, +) +from .config import _fp4_mla_cutedsl_kernel_module as _fp4_mla_cutedsl_kernel_module +from .config import _fp4_mla_q1_kv_blocks_per_program as _fp4_mla_q1_kv_blocks_per_program +from .config import _fp4_mla_q1_prefix_blocks_per_program as _fp4_mla_q1_prefix_blocks_per_program +from .config import _fp4_mla_q1_preload_variants as _fp4_mla_q1_preload_variants +from .config import _fp4_mla_triton_preload_key_set as _fp4_mla_triton_preload_key_set +from .config import _HPUpdatePhase as _HPUpdatePhase +from .decode import _SM_COUNT_CACHE as _SM_COUNT_CACHE +from .decode import _cutedsl_pad_q_and_sf_kernel as _cutedsl_pad_q_and_sf_kernel +from .decode import _cutedsl_swizzled_sf_offset as _cutedsl_swizzled_sf_offset +from .decode import _get_sm_count as _get_sm_count +from .decode import _run_triton_attention_decode as _run_triton_attention_decode +from .decode import run_fp4_mla_attention_decode +from .layout import _ensure_workspace_tensor as _ensure_workspace_tensor +from .layout import _get_fp4_mla_global_scale as _get_fp4_mla_global_scale +from .layout import _get_fp4_mla_hp_pool_layout as _get_fp4_mla_hp_pool_layout +from .layout import _get_fp4_mla_kv_cache_tensors as _get_fp4_mla_kv_cache_tensors +from .layout import _get_fp4_mla_q_global_scale as _get_fp4_mla_q_global_scale +from .layout import _get_fp4_mla_swizzled_scale_size as _get_fp4_mla_swizzled_scale_size +from .layout import _validate_fp4_mla_attention_q_shape as _validate_fp4_mla_attention_q_shape +from .layout import _validate_fp4_mla_cache_shape as _validate_fp4_mla_cache_shape +from .layout import _validate_fp4_mla_hp_generation_width as _validate_fp4_mla_hp_generation_width +from .layout import _validate_fp4_mla_kv_storage_shape as _validate_fp4_mla_kv_storage_shape +from .layout import ( + get_fp4_mla_v_scale_pool_shape, + get_fp4_mla_v_scale_pool_size, + get_fp4_mla_v_scale_pool_view, +) +from .metadata import _fp4_mla_append_metadata_kernel as _fp4_mla_append_metadata_kernel +from .metadata import _fp4_mla_generation_hp_page_ids as _fp4_mla_generation_hp_page_ids +from .metadata import _fp4_mla_generation_lengths_kernel as _fp4_mla_generation_lengths_kernel +from .metadata import _fp4_mla_generation_num_blocks_device as _fp4_mla_generation_num_blocks_device +from .metadata import _fp4_mla_generation_page_ids as _fp4_mla_generation_page_ids +from .metadata import ( + _fp4_mla_materialize_page_table_kernel as _fp4_mla_materialize_page_table_kernel, +) +from .metadata import _fp4_mla_page_table_spec as _fp4_mla_page_table_spec +from .metadata import ( + _fp4_mla_store_sequence_append_metadata as _fp4_mla_store_sequence_append_metadata, +) +from .metadata import _fp4_mla_uniform_generation_lengths as _fp4_mla_uniform_generation_lengths +from .metadata import _get_linear_mtp_query_len_per_seq as _get_linear_mtp_query_len_per_seq +from .metadata import _host_int_list as _host_int_list +from .metadata import _host_int_list_during_forward as _host_int_list_during_forward +from .metadata import _infer_assume_full_pages as _infer_assume_full_pages +from .metadata import ( + _materialize_fp4_mla_device_page_table_for_forward as _materialize_fp4_mla_device_page_table_for_forward, +) +from .metadata import _max_generation_pages as _max_generation_pages +from .metadata import ( + configure_fp4_mla_device_page_table, + materialize_fp4_mla_device_page_table, + populate_fp4_mla_append_metadata, + populate_fp4_mla_generation_lengths, +) +from .v_cache import ( + _get_cutedsl_persistent_v_packed_cache as _get_cutedsl_persistent_v_packed_cache, +) +from .v_cache import _get_fp4_mla_v_packed_pool as _get_fp4_mla_v_packed_pool +from .v_cache import _get_fp4_mla_v_packed_pool_base as _get_fp4_mla_v_packed_pool_base +from .v_cache import _get_fp4_mla_v_scale_pool_base as _get_fp4_mla_v_scale_pool_base +from .v_cache import _get_triton_v_packed_cache as _get_triton_v_packed_cache +from .v_cache import _is_triton_v_packed_cache_valid as _is_triton_v_packed_cache_valid +from .v_cache import _maybe_update_triton_v_packed_cache as _maybe_update_triton_v_packed_cache +from .v_cache import _repack_cutedsl_v_packed_cache as _repack_cutedsl_v_packed_cache +from .v_cache import _select_triton_block_v as _select_triton_block_v +from .v_cache import _set_triton_v_packed_cache_valid as _set_triton_v_packed_cache_valid +from .v_cache import _shared_v_pack_storage_enabled as _shared_v_pack_storage_enabled +from .v_cache import _triton_can_prepack_v as _triton_can_prepack_v +from .v_cache import _triton_prepack_v_enabled as _triton_prepack_v_enabled +from .v_cache import _triton_shared_v_packed_valid_attr as _triton_shared_v_packed_valid_attr +from .v_cache import _triton_v_packed_attr as _triton_v_packed_attr +from .v_cache import _triton_v_packed_cache_tag as _triton_v_packed_cache_tag +from .v_cache import _triton_v_packed_valid_attr as _triton_v_packed_valid_attr +from .v_cache import _update_triton_v_packed_cache as _update_triton_v_packed_cache +from .v_cache import _v_packed_cache_tag as _v_packed_cache_tag +from .v_cache import _v_packed_shape as _v_packed_shape + +__all__ = [ + "FP4_BLOCK_SIZE", + "FP4_MLA_ATTENTION_BACKEND_ENV", + "FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV", + "FP4_MLA_E4M3_MAX", + "FP4_MLA_KV_GLOBAL_SCALE", + "FP4_MLA_KV_STATIC_AMAX", + "FP4_MLA_K_RESIDUAL_DIM", + "FP4_MLA_P_GLOBAL_SCALE", + "FP4_MLA_Q1_PREFIX_BLOCK_DIM", + "FP4_MLA_Q_GLOBAL_SCALE", + "FP4_MLA_Q_LOGICAL_DIM", + "FP4_MLA_Q_PACKED_DIM", + "FP4_MLA_Q_PREFIX_BLOCK_DIM", + "FP4_MLA_Q_PREFIX_DIM", + "FP4_MLA_Q_RESIDUAL_DIM", + "FP4_MLA_Q_SF_GROUPS", + "FP4_MLA_Q_STATIC_AMAX", + "FP4_MLA_SCALE_COL_GROUP", + "FP4_MLA_SCALE_ROW_GROUP", + "FP4_MLA_TOKENS_PER_BLOCK", + "HP_BLOCK_SIZE", + "can_fuse_fp4_mla_q_quant", + "configure_fp4_mla_device_page_table", + "get_fp4_mla_v_scale_pool_shape", + "get_fp4_mla_v_scale_pool_size", + "get_fp4_mla_v_scale_pool_view", + "materialize_fp4_mla_device_page_table", + "populate_fp4_mla_append_metadata", + "populate_fp4_mla_generation_lengths", + "run_fp4_mla_attention_decode", + "scatter_fp4_mla_kv_cache", +] diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index e5e8f59ed92f..e83da387f9a1 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -26,7 +26,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig -from . import ( +from .config import ( _FP4_MLA_CUTEDSL_BACKEND, _FP4_MLA_K_RESIDUAL_BACKENDS, FP4_BLOCK_SIZE, @@ -35,8 +35,8 @@ HP_BLOCK_SIZE, _fp4_mla_attention_backend, _fp4_mla_cutedsl_fused_v_transpose_enabled, - get_fp4_mla_v_scale_pool_size, ) +from .layout import get_fp4_mla_v_scale_pool_size @dataclass(frozen=True) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py new file mode 100644 index 000000000000..0ff9a6442e7d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py @@ -0,0 +1,1087 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP4 MLA context/generation cache scatter and fused Q/RoPE quantization.""" + +from typing import Any, Optional + +import torch +import triton + +from tensorrt_llm._utils import get_sm_version + +from .config import ( + _FP4_MLA_CUTEDSL_BACKEND, + _FP4_MLA_K_RESIDUAL_BACKENDS, + _FP4_MLA_MAX_GRID_Z, + FP4_BLOCK_SIZE, + FP4_MLA_K_RESIDUAL_DIM, + FP4_MLA_Q1_PREFIX_BLOCK_DIM, + FP4_MLA_Q_LOGICAL_DIM, + FP4_MLA_Q_PACKED_DIM, + FP4_MLA_Q_PREFIX_BLOCK_DIM, + FP4_MLA_Q_PREFIX_DIM, + FP4_MLA_Q_RESIDUAL_DIM, + FP4_MLA_Q_SF_GROUPS, + FP4_MLA_SCALE_ROW_GROUP, + FP4_MLA_TOKENS_PER_BLOCK, + HP_BLOCK_SIZE, + _ceil_div, + _fp4_mla_attention_backend, + _fp4_mla_cutedsl_fused_v_transpose_enabled, + _fp4_mla_q1_kv_blocks_per_program, + _fp4_mla_q1_prefix_blocks_per_program, + _fp4_mla_q1_preload_variants, + _fp4_mla_triton_preload_key_set, + _HPUpdatePhase, +) +from .fp4_mla_kernels import ( + _fp4_mla_context_cache_update_kernel, + _fp4_mla_generation_fused_qk_rope_cache_update_kernel, +) +from .layout import ( + _get_fp4_mla_global_scale, + _get_fp4_mla_hp_pool_layout, + _get_fp4_mla_kv_cache_tensors, + _get_fp4_mla_q_global_scale, + _get_fp4_mla_swizzled_scale_size, + _validate_fp4_mla_cache_shape, + _validate_fp4_mla_hp_generation_width, + _validate_fp4_mla_kv_storage_shape, + get_fp4_mla_v_scale_pool_view, +) +from .metadata import ( + _fp4_mla_generation_hp_page_ids, + _fp4_mla_generation_num_blocks_device, + _fp4_mla_generation_page_ids, + _fp4_mla_uniform_generation_lengths, + _materialize_fp4_mla_device_page_table_for_forward, +) +from .v_cache import ( + _get_cutedsl_persistent_v_packed_cache, + _get_fp4_mla_v_packed_pool_base, + _maybe_update_triton_v_packed_cache, + _repack_cutedsl_v_packed_cache, +) + + +def _validate_fp4_mla_context_rope( + latent_cache: torch.Tensor, + rotary_cos_sin: torch.Tensor, + v_head_dim: int, +) -> int: + head_dim = latent_cache.shape[-1] + rope_dim = head_dim - v_head_dim + if rope_dim <= 0 or rope_dim % 2 != 0: + raise ValueError( + "FP4 MLA fused context K-RoPE requires a positive even RoPE dimension, " + f"got head_dim={head_dim}, v_head_dim={v_head_dim}." + ) + if rotary_cos_sin.device != latent_cache.device: + raise ValueError("FP4 MLA context latent cache and RoPE table must use the same device.") + if rotary_cos_sin.dtype != torch.float32: + raise TypeError( + f"FP4 MLA fused context K-RoPE requires a float32 table, got {rotary_cos_sin.dtype}." + ) + if not rotary_cos_sin.is_contiguous(): + raise ValueError("FP4 MLA fused context K-RoPE requires a contiguous RoPE table.") + table_row_size = rope_dim * 2 + if rotary_cos_sin.numel() < table_row_size or rotary_cos_sin.numel() % table_row_size != 0: + raise ValueError( + "FP4 MLA context RoPE table size must be a positive multiple of " + f"{table_row_size}, got {rotary_cos_sin.numel()}." + ) + return rope_dim + + +def can_fuse_fp4_mla_q_quant( + metadata: Any, + q: torch.Tensor, + q_pe: torch.Tensor, + latent_cache: torch.Tensor, +) -> bool: + """Return whether generation can quantize Q in the fused cache update.""" + num_gen = metadata.num_seqs - metadata.num_contexts + return bool( + _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS + and get_sm_version() == 107 + and num_gen > 0 + and getattr(metadata, "kv_cache_manager", None) is not None + and q.shape[0] > 0 + and q.shape[0] % num_gen == 0 + and q.is_cuda + and q_pe.is_cuda + and latent_cache.is_cuda + and q.device == q_pe.device == latent_cache.device + and q.dtype == torch.bfloat16 + and q.is_contiguous() + and q.ndim == 3 + and 0 < q.shape[1] <= 128 + and q.shape[2] == FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM + and q_pe.dtype == torch.bfloat16 + and tuple(q_pe.shape) == (q.shape[0], q.shape[1], FP4_MLA_Q_RESIDUAL_DIM) + and latent_cache.dtype == torch.bfloat16 + and tuple(latent_cache.shape) == (q.shape[0], FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM) + and metadata.page_size == FP4_MLA_TOKENS_PER_BLOCK + ) + + +def _prepare_fp4_mla_q_buffers( + metadata: Any, + num_queries: int, + num_heads: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Return manager-owned, fixed-capacity packed-Q staging buffers.""" + if num_queries <= 0: + raise ValueError(f"FP4 MLA Q buffer preparation needs queries, got {num_queries}.") + owner = getattr(metadata, "kv_cache_manager", None) + if owner is None: + raise RuntimeError("Fused FP4 MLA Q quantization requires a KV cache manager.") + if num_heads <= 0 or num_heads > 128: + raise ValueError(f"FP4 MLA Q buffers require 1-128 local heads, got {num_heads}.") + buffers = getattr(owner, "_fp4_mla_q_buffers", None) + if buffers is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Cannot create FP4 MLA Q buffers while capturing a CUDA graph. " + "Run a warmup forward first." + ) + buffers = {} + setattr(owner, "_fp4_mla_q_buffers", buffers) + + max_num_tokens = int(getattr(metadata, "max_num_tokens", num_queries) or num_queries) + max_num_sequences = int( + getattr(metadata, "max_num_sequences", None) + or getattr(metadata, "max_num_requests", num_queries) + or num_queries + ) + max_query_width = 1 + int(getattr(metadata, "max_total_draft_tokens", None) or 0) + capacity = min( + max_num_tokens, + max_num_sequences * max_query_width, + _FP4_MLA_MAX_GRID_Z, + ) + if num_queries > capacity: + raise ValueError( + f"FP4 MLA active queries exceed the configured Q capacity: {num_queries} > {capacity}." + ) + + device_index = device.index if device.index is not None else torch.cuda.current_device() + canonical_device = torch.device("cuda", device_index) + expected_q_shape = (capacity * num_heads, FP4_MLA_Q_PACKED_DIM) + expected_q_sf_shape = ( + _get_fp4_mla_swizzled_scale_size(capacity * num_heads, FP4_MLA_Q_LOGICAL_DIM), + ) + q_key = f"q_{num_heads}" + q_sf_key = f"q_sf_{num_heads}" + q_storage = buffers.get(q_key) + q_sf_storage = buffers.get(q_sf_key) + if q_storage is None and q_sf_storage is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Cannot allocate FP4 MLA Q buffers while capturing a CUDA graph. " + "Run a warmup forward first." + ) + q_storage = torch.empty(expected_q_shape, dtype=torch.uint8, device=canonical_device) + q_sf_storage = torch.empty( + expected_q_sf_shape, + dtype=torch.float8_e4m3fn, + device=canonical_device, + ) + buffers[q_key] = q_storage + buffers[q_sf_key] = q_sf_storage + if ( + q_storage is None + or q_storage.dtype != torch.uint8 + or q_storage.device != canonical_device + or tuple(q_storage.shape) != expected_q_shape + or not q_storage.is_contiguous() + or q_sf_storage is None + or q_sf_storage.dtype != torch.float8_e4m3fn + or q_sf_storage.device != canonical_device + or tuple(q_sf_storage.shape) != expected_q_sf_shape + or not q_sf_storage.is_contiguous() + ): + raise RuntimeError("FP4 MLA packed-Q buffers do not match the configured capacity.") + return q_storage, q_sf_storage, capacity + + +def _get_fp4_mla_context_start_positions(metadata: Any, num_contexts: int) -> torch.Tensor: + kv_cache_params = getattr(metadata, "kv_cache_params", None) + cached_token_lens = getattr(kv_cache_params, "num_cached_tokens_per_seq", None) + if cached_token_lens is not None: + return torch.as_tensor(cached_token_lens[:num_contexts], dtype=torch.int64, device="cpu") + + return ( + ( + metadata.kv_lens_cuda_runtime[:num_contexts] + - metadata.prompt_lens_cuda_runtime[:num_contexts] + ) + .detach() + .cpu() + ) + + +def _validate_fp4_mla_context_start_alignment( + metadata: Any, + num_contexts: int, + *, + alignment: int = HP_BLOCK_SIZE, +) -> None: + context_start_positions = _get_fp4_mla_context_start_positions(metadata, num_contexts) + bad_start = (context_start_positions < 0) | ((context_start_positions % alignment) != 0) + if bool(torch.any(bad_start).item()): + starts = context_start_positions.detach().cpu().tolist() + raise ValueError( + "FP4 MLA shared-tile context update requires every context " + f"start position to be {alignment}-token aligned, got " + f"start positions {starts}." + ) + + +def _scatter_fp4_mla_kv_cache_2d_context( + metadata: Any, + latent_cache: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + rotary_cos_sin: Optional[torch.Tensor], + *, + token_offset: int, + local_layer: int, + v_head_dim: int, + head_dim: int, + num_tokens: int, + num_dim_blocks: int, + sf_per_token: int, + sf_per_page: int, + v_packed_base: Optional[torch.Tensor] = None, + v_page_offset: int = 0, +) -> bool: + num_contexts = metadata.num_contexts + if num_contexts > 0: + prompt_lens_cpu = metadata.prompt_lens_cpu_runtime[:num_contexts] + ctx_token_count = int(prompt_lens_cpu.sum().item()) + if num_tokens != ctx_token_count: + raise RuntimeError( + f"FP4 MLA 2D context scatter needs {ctx_token_count} context tokens, got " + f"{num_tokens}." + ) + _validate_fp4_mla_context_start_alignment(metadata, num_contexts, alignment=FP4_BLOCK_SIZE) + + apply_k_rope = rotary_cos_sin is not None + rope_dim = ( + _validate_fp4_mla_context_rope(latent_cache, rotary_cos_sin, v_head_dim) + if rotary_cos_sin is not None + else 0 + ) + rotary_cos_sin_ptr = rotary_cos_sin if rotary_cos_sin is not None else latent_cache + + hp_pool = getattr(metadata.fp4_mla_state, "hp_pool", None) + if not isinstance(hp_pool, torch.Tensor): + raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") + if hp_pool.device != latent_cache.device: + raise ValueError("FP4 MLA latent cache and high-precision pool must share a device.") + if hp_pool.dtype != torch.bfloat16: + raise TypeError(f"FP4 MLA high-precision KV pool must use BF16, got {hp_pool.dtype}.") + hp_pool_size, pool_head_dim = _get_fp4_mla_hp_pool_layout(metadata, hp_pool) + if pool_head_dim < head_dim: + raise RuntimeError( + f"FP4 MLA HP pool head dimension is too small: got " + f"{pool_head_dim}, need at least {head_dim}." + ) + if local_layer < 0 or local_layer >= hp_pool.shape[1]: + raise ValueError( + f"FP4 MLA local layer {local_layer} is outside the HP pool's {hp_pool.shape[1]} layers." + ) + if hp_pool.stride(-1) != 1: + raise ValueError("FP4 MLA high-precision KV pool must be contiguous in head_dim.") + hp_page_ids = metadata.fp4_mla_state.hp_page_indices + if not isinstance(hp_page_ids, torch.Tensor): + raise RuntimeError("FP4 MLA context cache update requires HP page metadata.") + store_hp_tail = num_contexts > 0 + num_hp_pages = hp_pool.shape[0] + pool_s0 = hp_pool.stride(0) + pool_s1 = hp_pool.stride(1) + + write_v_packed = v_packed_base is not None + v_packed_output = v_packed_base if write_v_packed else kv_cache + v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 + v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 + + _fp4_mla_context_cache_update_kernel[ + ( + num_tokens, + num_dim_blocks, + ) + ]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + latent_cache, + global_scale, + rotary_cos_sin_ptr, + hp_pool, + hp_page_ids, + metadata.fp4_mla_state.batch_indices, + metadata.fp4_mla_state.positions, + metadata.fp4_mla_state.paged_kv_indices, + metadata.fp4_mla_state.paged_kv_indptr, + metadata.fp4_mla_state.paged_kv_indices.shape[0], + metadata.fp4_mla_state.paged_kv_indptr.shape[0], + metadata.fp4_mla_state.batch_indices.shape[0], + v_sf.shape[1], + v_sf.shape[0], + num_contexts, + num_hp_pages, + token_offset, + num_tokens, + local_layer, + v_page_offset if write_v_packed else 0, + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + latent_cache.stride(0), + latent_cache.stride(1), + v_sf.stride(0), + v_sf.stride(1), + v_packed_s0, + v_packed_s1, + pool_s0, + pool_s1, + HEAD_D=head_dim, + V_HEAD_D=v_head_dim, + HP_BLOCK=FP4_BLOCK_SIZE, + HP_POOL_SIZE=hp_pool_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=sf_per_page, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + STORE_K_RESIDUAL=(_fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS), + ROPE_DIM=rope_dim, + APPLY_K_ROPE=apply_k_rope, + POOL_HEAD_D=pool_head_dim, + STORE_HP_TAIL=store_hp_tail, + WRITE_V_PACKED=write_v_packed, + ) + return store_hp_tail + + +def _scatter_fp4_mla_kv_cache_2d_generation( + metadata: Any, + latent_cache: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + *, + token_offset: int, + local_layer: int, + v_head_dim: int, + head_dim: int, + num_tokens: int, + num_dim_blocks: int, + sf_per_token: int, + sf_per_page: int, + rotary_cos_sin: torch.Tensor, + q_pe: torch.Tensor, + q_rope_out: torch.Tensor, + q_quant_input: torch.Tensor, + q_fp4_out: torch.Tensor, + q_sf_out: torch.Tensor, + v_packed_base: Optional[torch.Tensor], + v_page_offset: int, +) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + num_contexts = metadata.num_contexts + num_seqs = metadata.num_seqs + num_gen = num_seqs - num_contexts + if num_gen <= 0: + return + if num_tokens < num_gen: + raise RuntimeError( + f"FP4 MLA 2D generation scatter needs at least {num_gen} generation " + f"tokens, got {num_tokens}." + ) + if num_tokens % num_gen != 0: + raise NotImplementedError( + "FP4 MLA no-dequant generation scatter requires a uniform linear MTP " + f"generation length, got {num_tokens} tokens for {num_gen} sequences." + ) + # The prompt_lens/kv_lens runtime aliases can lag at the decode anchor + # (seq_lens == 1) under CUDA graph / one-engine MTP while each generation + # sequence really appends num_tokens // num_gen tokens this step. Recover the + # true per-sequence lengths for the no-dequant kernel below (a no-op when the + # aliases already match). + kv_lens_gen, gen_lens_gen = _fp4_mla_uniform_generation_lengths(metadata, num_tokens, num_gen) + _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens_gen) + + pool = getattr(metadata.fp4_mla_state, "hp_pool", None) + if pool is None: + raise RuntimeError("FP4 MLA 2D generation scatter requires the HP KV pool.") + try: + hp_pool_size, hp_head_dim = _get_fp4_mla_hp_pool_layout(metadata, pool) + except ValueError as error: + raise RuntimeError(str(error)) from error + if hp_head_dim < head_dim: + raise RuntimeError( + f"FP4 MLA 2D generation scatter needs at least {head_dim} HP channels, got " + f"{hp_head_dim}." + ) + hp_page_ids = _fp4_mla_generation_hp_page_ids(metadata, num_gen) + if not isinstance(hp_page_ids, torch.Tensor): + raise RuntimeError("FP4 MLA generation requires HP page metadata.") + num_hp_pages = pool.shape[0] + + max_gen_len = num_tokens // num_gen + _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) + max_rewind_len = hp_pool_size - HP_BLOCK_SIZE + page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) + rope_dim = head_dim - v_head_dim + block_q_heads = 32 + q1_kv_blocks_per_program = 1 + # Grouped Q1 kernels reuse K's packed codes for V. Keep one dimension + # block per program until their warp-specialized V quantizer is split out. + if max_gen_len == 1: + q1_kv_blocks_per_program = _fp4_mla_q1_kv_blocks_per_program(num_gen, v_head_dim) + q_prefix_block_dim = ( + FP4_MLA_Q1_PREFIX_BLOCK_DIM if max_gen_len == 1 else FP4_MLA_Q_PREFIX_BLOCK_DIM + ) + q_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // q_prefix_block_dim + q_prefix_blocks_per_program = _fp4_mla_q1_prefix_blocks_per_program( + num_gen, + q1_kv_blocks_per_program, + ) + q_work_blocks = q_prefix_blocks // q_prefix_blocks_per_program + 1 + if latent_cache.dtype != torch.bfloat16: + raise TypeError( + "Fused FP4 MLA Q/K RoPE and cache storage requires BF16 latent KV, " + f"got {latent_cache.dtype}." + ) + if rope_dim <= 0 or rope_dim % 2 != 0: + raise ValueError( + "Fused FP4 MLA K RoPE requires a positive even K tail, " + f"got head_dim={head_dim} v_head_dim={v_head_dim}." + ) + if ( + rotary_cos_sin is None + or rotary_cos_sin.device != latent_cache.device + or rotary_cos_sin.dtype != torch.float32 + or rotary_cos_sin.numel() % (rope_dim * 2) != 0 + ): + raise ValueError( + "Fused FP4 MLA K RoPE requires a same-device FP32 rotary " + f"table with rows of {rope_dim * 2} values." + ) + if ( + q_pe is None + or q_rope_out is None + or q_pe.dtype != torch.bfloat16 + or q_rope_out.dtype != torch.bfloat16 + or q_pe.device != latent_cache.device + or q_rope_out.device != latent_cache.device + or q_pe.ndim != 3 + or q_rope_out.shape != q_pe.shape + or q_pe.shape[0] != num_tokens + or q_pe.shape[1] <= 0 + or q_pe.shape[2] != rope_dim + ): + raise ValueError( + "Fused FP4 MLA Q RoPE requires same-device BF16 q_pe and " + f"q_rope_out tensors shaped [tokens, heads, {rope_dim}]." + ) + num_q_heads = q_pe.shape[1] + q_head_blocks = _ceil_div(num_q_heads, block_q_heads) + max_gen_tiles = _ceil_div(max_gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) + rotary_table = rotary_cos_sin + q_global_scale = _get_fp4_mla_q_global_scale(metadata, latent_cache.device) + q_pe_input = q_pe + q_rope_output = q_rope_out + q_full_input = q_quant_input + q_fp4_output = q_fp4_out + q_sf_output = q_sf_out + write_v_packed = v_packed_base is not None + v_packed_output = v_packed_base if write_v_packed else kv_cache + v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 + v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 + kv_work_blocks = ( + v_head_dim // FP4_BLOCK_SIZE // q1_kv_blocks_per_program + 1 + if max_gen_len == 1 + else num_dim_blocks + ) + launch_grid = ( + num_gen, + max( + kv_work_blocks, + max_gen_len * q_head_blocks * q_work_blocks, + ), + ) + store_k_residual = _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS + + def launch_generation_update( + grid: tuple[int, ...], + *, + page_ids_len: int, + indptr_len: int, + max_gen_tiles_variant: int, + q_prefix_block_dim_variant: int, + q_prefix_blocks_variant: int, + q_prefix_blocks_per_program_variant: int, + q1_kv_blocks_per_program_variant: int, + ) -> None: + q_work_blocks_variant = q_prefix_blocks_variant // q_prefix_blocks_per_program_variant + 1 + _fp4_mla_generation_fused_qk_rope_cache_update_kernel[grid]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + pool, + latent_cache, + global_scale, + q_global_scale, + rotary_table, + q_pe_input, + q_rope_output, + q_full_input, + q_fp4_output, + q_sf_output, + kv_lens_gen, + gen_lens_gen, + page_ids, + hp_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode, + page_ids_len, + hp_page_ids.numel(), + indptr_len, + v_sf.shape[1], + num_hp_pages, + v_sf.shape[0], + local_layer, + v_page_offset if write_v_packed else 0, + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + pool.stride(0), + pool.stride(1), + v_sf.stride(0), + v_sf.stride(1), + v_packed_s0, + v_packed_s1, + q_pe_input.stride(0), + q_pe_input.stride(1) if q_pe_input.ndim > 1 else 0, + q_pe_input.stride(2) if q_pe_input.ndim > 2 else 0, + q_rope_output.stride(0), + q_rope_output.stride(1) if q_rope_output.ndim > 1 else 0, + q_rope_output.stride(2) if q_rope_output.ndim > 2 else 0, + HEAD_D=hp_head_dim, + V_HEAD_D=v_head_dim, + HP_BLOCK=FP4_BLOCK_SIZE, + HP_POOL_SIZE=hp_pool_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=sf_per_page, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + STORE_K_RESIDUAL=store_k_residual, + FUSE_ROPE_CACHE_STORE=True, + WRITE_V_PACKED=write_v_packed, + MAX_GEN_TILES=max_gen_tiles_variant, + ROPE_DIM=rope_dim, + ROPE_PAIR_BLOCK=triton.next_power_of_2(rope_dim // 2), + NUM_DIM_BLOCKS=num_dim_blocks, + NUM_Q_HEADS=num_q_heads, + Q_HEAD_BLOCKS=max(q_head_blocks, 1), + BLOCK_Q_HEADS=block_q_heads, + Q_PREFIX_D=FP4_MLA_Q_PREFIX_DIM, + Q_PREFIX_BLOCK_D=q_prefix_block_dim_variant, + Q_PREFIX_BLOCKS=q_prefix_blocks_variant, + Q_PREFIX_BLOCKS_PER_PROGRAM=q_prefix_blocks_per_program_variant, + Q_WORK_BLOCKS=q_work_blocks_variant, + Q_SF_COLS=FP4_MLA_Q_SF_GROUPS, + WRITE_Q=True, + Q1_KV_BLOCKS_PER_PROGRAM=q1_kv_blocks_per_program_variant, + maxnreg=56, + ) + + # Triton compiles and loads a CUDA module on first launch. Use runtime-zero + # work here so every reachable static tuning variant is resident before + # warmup hands the engine to serving. + if getattr(metadata, "is_warmup", False) and not torch.cuda.is_current_stream_capturing(): + configured_generation_len = int(getattr(metadata, "max_total_draft_tokens", 0) or 0) + 1 + max_preload_generation_len = max(max_gen_len, configured_generation_len) + if max_preload_generation_len - 1 > max_rewind_len: + raise NotImplementedError( + "FP4 MLA finite Triton preload exceeds the HP ring's rewind slack: " + f"max_rewind={max_rewind_len}, generation=" + f"{max_preload_generation_len}." + ) + max_num_sequences = int(getattr(metadata, "max_num_sequences", num_seqs) or num_seqs) + q1_variants = _fp4_mla_q1_preload_variants( + max_num_sequences, + v_head_dim, + ) + multi_token_tiles = tuple( + sorted( + { + _ceil_div(gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) + for gen_len in range(2, max_preload_generation_len + 1) + } + ) + ) + preload_key = ( + "generation-cache-update", + str(latent_cache.device), + hp_pool_size, + hp_head_dim, + v_head_dim, + num_q_heads, + metadata.page_size, + write_v_packed, + store_k_residual, + tuple(q1_variants), + multi_token_tiles, + tuple(kv_cache.stride()), + tuple(sf_cache.stride()), + tuple(pool.stride()), + tuple(v_sf.stride()), + tuple(v_packed_output.stride()), + tuple(q_pe_input.stride()), + tuple(q_rope_output.stride()), + str(kv_cache.dtype), + str(latent_cache.dtype), + str(q_pe_input.dtype), + str(q_fp4_output.dtype), + str(q_sf_output.dtype), + ) + preload_keys = _fp4_mla_triton_preload_key_set(metadata) + if preload_key not in preload_keys: + context_page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) + if not isinstance(context_page_ids, torch.Tensor): + context_page_ids = page_ids + context_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) + if not isinstance(context_indptr, torch.Tensor): + context_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode + context_batch_indices = getattr(metadata.fp4_mla_state, "batch_indices", None) + if not isinstance(context_batch_indices, torch.Tensor): + context_batch_indices = hp_page_ids + context_positions = getattr(metadata.fp4_mla_state, "positions", None) + if not isinstance(context_positions, torch.Tensor): + context_positions = hp_page_ids + _fp4_mla_context_cache_update_kernel[(1, num_dim_blocks)]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + latent_cache, + global_scale, + rotary_table, + pool, + hp_page_ids, + context_batch_indices, + context_positions, + context_page_ids, + context_indptr, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + latent_cache.stride(0), + latent_cache.stride(1), + v_sf.stride(0), + v_sf.stride(1), + v_packed_s0, + v_packed_s1, + pool.stride(0), + pool.stride(1), + HEAD_D=head_dim, + V_HEAD_D=v_head_dim, + HP_BLOCK=FP4_BLOCK_SIZE, + HP_POOL_SIZE=hp_pool_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=sf_per_page, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + STORE_K_RESIDUAL=store_k_residual, + ROPE_DIM=rope_dim, + APPLY_K_ROPE=True, + POOL_HEAD_D=hp_head_dim, + STORE_HP_TAIL=True, + WRITE_V_PACKED=write_v_packed, + ) + q1_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM + for q1_kv_blocks, q1_prefix_blocks_per_program in q1_variants: + launch_generation_update( + (1, 1), + page_ids_len=0, + indptr_len=0, + max_gen_tiles_variant=1, + q_prefix_block_dim_variant=FP4_MLA_Q1_PREFIX_BLOCK_DIM, + q_prefix_blocks_variant=q1_prefix_blocks, + q_prefix_blocks_per_program_variant=q1_prefix_blocks_per_program, + q1_kv_blocks_per_program_variant=q1_kv_blocks, + ) + multi_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q_PREFIX_BLOCK_DIM + for multi_token_tile in multi_token_tiles: + launch_generation_update( + (1, 1), + page_ids_len=0, + indptr_len=0, + max_gen_tiles_variant=multi_token_tile, + q_prefix_block_dim_variant=FP4_MLA_Q_PREFIX_BLOCK_DIM, + q_prefix_blocks_variant=multi_prefix_blocks, + q_prefix_blocks_per_program_variant=1, + q1_kv_blocks_per_program_variant=1, + ) + torch.cuda.synchronize(latent_cache.device) + preload_keys.add(preload_key) + + launch_generation_update( + launch_grid, + page_ids_len=page_ids.shape[0], + indptr_len=metadata.fp4_mla_state.paged_kv_indptr_decode.shape[0], + max_gen_tiles_variant=max(max_gen_tiles, 1), + q_prefix_block_dim_variant=q_prefix_block_dim, + q_prefix_blocks_variant=q_prefix_blocks, + q_prefix_blocks_per_program_variant=q_prefix_blocks_per_program, + q1_kv_blocks_per_program_variant=q1_kv_blocks_per_program, + ) + return kv_lens_gen, gen_lens_gen, page_ids + + +# Public cache update and decode entry points + + +def scatter_fp4_mla_kv_cache( + metadata: Any, + latent_cache: torch.Tensor, + layer_idx: int, + *, + token_offset: int, + phase: _HPUpdatePhase, + local_layer: int, + v_head_dim: int, + rotary_cos_sin: Optional[torch.Tensor] = None, + q_pe: Optional[torch.Tensor] = None, + q_rope_out: Optional[torch.Tensor] = None, + q_quant_input: Optional[torch.Tensor] = None, +) -> bool: + """Quantize MLA latent tokens and scatter them into the paged FP4 cache. + + Contract: this helper scatters exactly ``latent_cache.shape[0]`` tokens, + reading index metadata at ``batch_indices[token_offset : token_offset + N]`` + and ``positions[token_offset : token_offset + N]``. Callers must pass a + latent_cache pre-sliced to the current phase (context or generation) so + that ``shape[0]`` matches the number of index entries they intend to + consume. ``MLA.forward_impl`` (tensorrt_llm/_torch/modules/attention.py) + slices ``latent_cache[:num_ctx_tokens]`` for context and + ``latent_cache[num_ctx_tokens:]`` for generation before dispatching. + + Callers must pass ``phase``, ``local_layer``, and ``v_head_dim``. Context + scatter writes the final FP4 tile representation directly. Dimensions + below ``v_head_dim`` share one 16-token by 16-dim FP4 tile between K + and V, with the scale written into K's token-major and V's dim-major + layouts. Tail K-only dimensions use K's per-token 1D scales. For + exclusively owned CuTeDSL pages, context scatter also writes the + persistent packed-V sidecar. + Context scatter can rotate the K tail directly from the unassembled latent + tensor. Generation scatter rewrites each touched 16-token tile by reading + old tokens from the HP pool and new tokens from ``latent_cache``. The + static-scale generation + specialization can also rotate Q and new K tails while updating the HP pool. + The context kernel also stores the final incomplete tile in the BF16 HP + pool. When ``q_quant_input`` is supplied, the generation kernel also emits + backend-ready residual FP4 Q. The return value reports whether the current + phase updated the HP pool. + """ + if phase == "generation": + metadata.fp4_mla_state.generation_cache_scattered = False + metadata.fp4_mla_state.prequantized_q = None + metadata.fp4_mla_state.prequantized_q_sf = None + metadata.fp4_mla_state.q_batch_capacity = None + if latent_cache.numel() == 0: + raise ValueError("FP4 MLA cache scatter requires at least one latent token.") + + latent_cache = latent_cache.reshape(latent_cache.shape[0], -1).contiguous() + num_tokens = latent_cache.shape[0] + head_dim = latent_cache.shape[-1] + if head_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA KV head_dim must be divisible by {FP4_BLOCK_SIZE}, got {head_dim}." + ) + indices_len = metadata.fp4_mla_state.batch_indices.shape[0] + positions_len = metadata.fp4_mla_state.positions.shape[0] + if token_offset + num_tokens > indices_len or token_offset + num_tokens > positions_len: + raise RuntimeError( + f"FP4 MLA scatter would read batch_indices[{token_offset}:" + f"{token_offset + num_tokens}] / positions[{token_offset}:" + f"{token_offset + num_tokens}], but only {indices_len} / " + f"{positions_len} entries are available. This indicates " + "latent_cache was not pre-sliced to the current phase's token " + "range (see MLA.forward_impl)." + ) + + _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) + + backend = _fp4_mla_attention_backend() + global_scale = _get_fp4_mla_global_scale(metadata, latent_cache.device) + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + storage_head_dim = _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=backend, + ) + sf_per_token = storage_head_dim // FP4_BLOCK_SIZE + + if phase not in ("context", "generation"): + raise ValueError("FP4 MLA scatter requires phase='context' or 'generation'.") + if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: + raise RuntimeError("FP4 MLA scatter requires the auxiliary V scale pool.") + if metadata.page_size % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA scatter requires page_size divisible by " + f"{FP4_BLOCK_SIZE}, got {metadata.page_size}." + ) + if v_head_dim > head_dim: + raise ValueError(f"FP4 MLA v_head_dim={v_head_dim} cannot exceed head_dim={head_dim}.") + if head_dim - v_head_dim != FP4_MLA_K_RESIDUAL_DIM: + raise ValueError( + "FP4 MLA K residual quantization requires the K-only tail to match " + f"the {FP4_MLA_K_RESIDUAL_DIM}-channel residual, got " + f"head_dim={head_dim} v_head_dim={v_head_dim}." + ) + if v_head_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA v_head_dim must be divisible by {FP4_BLOCK_SIZE}, got {v_head_dim}." + ) + + sf_cache = sf_cache.view(torch.float8_e4m3fn) + v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=v_head_dim) + num_dim_blocks = triton.cdiv(head_dim, FP4_BLOCK_SIZE) + sf_per_page = metadata.page_size // FP4_BLOCK_SIZE + + generation_state = None + generation_inputs = (rotary_cos_sin, q_pe, q_rope_out, q_quant_input) + q_fp4_out = None + q_sf_out = None + if phase == "context": + if any(arg is not None for arg in (q_pe, q_rope_out, q_quant_input)): + raise ValueError("FP4 MLA context cache update does not accept generation Q tensors.") + hp_pool_updated = False + else: + if not all(arg is not None for arg in generation_inputs): + raise ValueError( + "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " + "and q_quant_input for fused RoPE, cache update, and Q quantization." + ) + if not can_fuse_fp4_mla_q_quant(metadata, q_quant_input, q_pe, latent_cache): + raise ValueError( + "Fused FP4 MLA Q quantization received an unsupported shape, dtype, " + "scale mode, or backend." + ) + q_fp4_out, q_sf_out, q_batch_capacity = _prepare_fp4_mla_q_buffers( + metadata, + num_tokens, + q_quant_input.shape[1], + q_quant_input.device, + ) + metadata.fp4_mla_state.prequantized_q = q_fp4_out + metadata.fp4_mla_state.prequantized_q_sf = q_sf_out + metadata.fp4_mla_state.q_batch_capacity = q_batch_capacity + hp_pool_updated = True + cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND + fused_v_transpose = cutedsl_backend and _fp4_mla_cutedsl_fused_v_transpose_enabled() + kv_cache_manager = getattr(metadata, "kv_cache_manager", None) + persistent_v_packed = None + v_packed_base = None + v_page_offset = 0 + direct_v_packed_write = False + if cutedsl_backend and not fused_v_transpose: + persistent_v_packed = _get_cutedsl_persistent_v_packed_cache( + metadata, + local_layer, + kv_cache, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + ) + # Reused/imported pages may not carry this process-local sidecar. Write + # packed V directly only when the cache pages are exclusively owned. + # The fused generation kernel handles uniform linear-MTP batches and + # updates every 16-token tile touched by the verification window. + num_gen = metadata.num_seqs - metadata.num_contexts + block_reuse = getattr(kv_cache_manager, "enable_block_reuse", True) + direct_context_v_packed_write = ( + phase == "context" + and metadata.num_contexts > 0 + and num_tokens > 0 + and block_reuse is False + ) + direct_generation_v_packed_write = ( + phase == "generation" + and num_gen > 0 + and num_tokens >= num_gen + and num_tokens % num_gen == 0 + and block_reuse is False + ) + direct_v_packed_write = direct_context_v_packed_write or direct_generation_v_packed_write + if direct_v_packed_write: + v_packed_base = _get_fp4_mla_v_packed_pool_base(metadata) + get_v_page_offset = getattr(kv_cache_manager, "get_mla_v_packed_page_offset", None) + v_page_offset = ( + int(get_v_page_offset(local_layer)) + if callable(get_v_page_offset) + else local_layer * kv_cache.shape[0] + ) + expected_row_width = metadata.page_size // 2 + required_base_rows = (v_page_offset + kv_cache.shape[0]) * v_head_dim + if ( + not isinstance(v_packed_base, torch.Tensor) + or v_packed_base.dtype != torch.uint8 + or v_packed_base.device != kv_cache.device + or v_packed_base.ndim != 2 + or v_packed_base.shape[0] < required_base_rows + or v_packed_base.shape[1] != expected_row_width + or not v_packed_base.is_contiguous() + ): + raise RuntimeError( + "FP4 MLA direct V-packed cache update requires the " + "stable full-pool base to be a contiguous uint8 tensor " + f"with at least {required_base_rows} rows and " + f"{expected_row_width} columns on {kv_cache.device}." + ) + expected_layer_ptr = v_packed_base.data_ptr() + ( + v_page_offset * v_head_dim * expected_row_width + ) + if expected_layer_ptr != persistent_v_packed.data_ptr(): + raise RuntimeError( + "FP4 MLA V-packed layer view does not match its stable " + "full-pool base and page offset." + ) + v_pack_num_valid = None + if phase == "context": + _materialize_fp4_mla_device_page_table_for_forward(metadata) + hp_pool_updated = _scatter_fp4_mla_kv_cache_2d_context( + metadata, + latent_cache, + kv_cache, + sf_cache, + v_sf, + global_scale, + rotary_cos_sin, + token_offset=token_offset, + local_layer=local_layer, + v_head_dim=v_head_dim, + head_dim=head_dim, + num_tokens=num_tokens, + num_dim_blocks=num_dim_blocks, + sf_per_token=sf_per_token, + sf_per_page=sf_per_page, + v_packed_base=v_packed_base, + v_page_offset=v_page_offset, + ) + v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices + else: + generation_state = _scatter_fp4_mla_kv_cache_2d_generation( + metadata, + latent_cache, + kv_cache, + sf_cache, + v_sf, + global_scale, + token_offset=token_offset, + local_layer=local_layer, + v_head_dim=v_head_dim, + head_dim=head_dim, + num_tokens=num_tokens, + num_dim_blocks=num_dim_blocks, + sf_per_token=sf_per_token, + sf_per_page=sf_per_page, + rotary_cos_sin=rotary_cos_sin, + q_pe=q_pe, + q_rope_out=q_rope_out, + q_quant_input=q_quant_input, + q_fp4_out=q_fp4_out, + q_sf_out=q_sf_out, + v_packed_base=v_packed_base, + v_page_offset=v_page_offset, + ) + v_pack_page_ids = _fp4_mla_generation_page_ids( + metadata, metadata.num_seqs - metadata.num_contexts + ) + if getattr(metadata, "is_cuda_graph", False): + # Frozen launch grids cannot follow the per-replay page count; + # the repack kernels stride over this device-side count instead. + v_pack_num_valid = _fp4_mla_generation_num_blocks_device(metadata) + cutedsl_repack_page_indptr = None + cutedsl_repack_kv_lens = None + cutedsl_repack_generation_lens = None + cutedsl_repack_max_touched_pages = 1 + if cutedsl_backend and not fused_v_transpose and not direct_v_packed_write: + if phase == "context": + num_contexts = metadata.num_contexts + cutedsl_v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices[ + : metadata.fp4_mla_state.num_context_blocks + ] + cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr[: num_contexts + 1] + cutedsl_repack_kv_lens = metadata.kv_lens_cuda_runtime[:num_contexts] + cutedsl_repack_generation_lens = metadata.prompt_lens_cuda_runtime[:num_contexts] + cutedsl_repack_max_touched_pages = int( + metadata.fp4_mla_state.context_repack_max_touched_pages + ) + elif generation_state is not None: + kv_lens_gen, gen_lens_gen, generation_page_ids = generation_state + num_gen = kv_lens_gen.numel() + cutedsl_v_pack_page_ids = generation_page_ids + cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode + cutedsl_repack_kv_lens = kv_lens_gen + cutedsl_repack_generation_lens = gen_lens_gen + cutedsl_repack_max_touched_pages = _ceil_div( + num_tokens // num_gen + metadata.page_size - 1, + metadata.page_size, + ) + else: + cutedsl_v_pack_page_ids = v_pack_page_ids + _repack_cutedsl_v_packed_cache( + persistent_v_packed, + kv_cache, + cutedsl_v_pack_page_ids, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + page_indptr=cutedsl_repack_page_indptr, + kv_lens=cutedsl_repack_kv_lens, + generation_lens=cutedsl_repack_generation_lens, + max_touched_pages=cutedsl_repack_max_touched_pages, + ) + _maybe_update_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + v_pack_page_ids, + num_queries=num_tokens, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + local_layer=local_layer, + v_sf=v_sf[local_layer], + num_valid_pages=v_pack_num_valid, + ) + if phase == "generation": + metadata.fp4_mla_state.generation_cache_scattered = hp_pool_updated + return hp_pool_updated diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py new file mode 100644 index 000000000000..80626bb41deb --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP4 MLA storage constants, backend selection, and kernel tuning policy.""" + +import importlib.util +import os +from typing import Any, Literal, Optional + +from tensorrt_llm._utils import get_sm_version + +HP_BLOCK_SIZE: int = 16 + + +FP4_BLOCK_SIZE: int = 16 + + +FP4_MLA_TOKENS_PER_BLOCK: int = 128 + + +FP4_MLA_SCALE_ROW_GROUP: int = 128 + + +FP4_MLA_SCALE_COL_GROUP: int = 4 + + +FP4_MLA_P_GLOBAL_SCALE: float = 448.0 * 6.0 + + +FP4_MLA_Q_STATIC_AMAX: float = 400.0 + + +FP4_MLA_KV_STATIC_AMAX: float = 30.0 + + +FP4_MLA_Q_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_Q_STATIC_AMAX + + +FP4_MLA_KV_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_KV_STATIC_AMAX + + +# Max finite e4m3 magnitude for FP4 MLA block-scale clamping. +FP4_MLA_E4M3_MAX: float = 448.0 + + +FP4_MLA_Q_RESIDUAL_DIM: int = 64 + + +FP4_MLA_K_RESIDUAL_DIM: int = FP4_MLA_Q_RESIDUAL_DIM + + +FP4_MLA_Q_PREFIX_DIM: int = 512 + + +FP4_MLA_Q_PREFIX_BLOCK_DIM: int = 256 + + +FP4_MLA_Q1_PREFIX_BLOCK_DIM: int = 512 + + +FP4_MLA_Q_LOGICAL_DIM: int = FP4_MLA_Q_PREFIX_DIM + 2 * FP4_MLA_Q_RESIDUAL_DIM + + +FP4_MLA_Q_PACKED_DIM: int = FP4_MLA_Q_LOGICAL_DIM // 2 + + +FP4_MLA_Q_SF_GROUPS: int = FP4_MLA_Q_LOGICAL_DIM // FP4_BLOCK_SIZE + + +FP4_MLA_ATTENTION_BACKEND_ENV = "TRTLLM_FP4_MLA_ATTENTION_BACKEND" + + +FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV = "TRTLLM_FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE" + + +_FP4_MLA_CUTEDSL_BACKEND = "cutedsl" + + +_FP4_MLA_K_RESIDUAL_BACKENDS = ("triton", _FP4_MLA_CUTEDSL_BACKEND) + + +_FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH = 4 + + +_FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH = 16 + + +_FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH = 32 + + +_FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD = 256 + + +_FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD = 512 + + +_FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD = 640 + + +_FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD = 768 + + +_HPUpdatePhase = Literal["context", "generation"] + + +_FP4_MLA_TRITON_PRELOAD_KEYS = "_fp4_mla_triton_preload_keys" + + +_FP4_MLA_PAGE_TABLE_TILE_SIZE = 128 + + +_FP4_MLA_MAX_GRID_Z = 65_535 + + +# Environment helpers + + +def _env_enabled_default(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None or value == "": + return default + return value.lower() in ( + "1", + "true", + "yes", + "on", + ) + + +def _fp4_mla_cutedsl_fused_v_transpose_enabled() -> bool: + return _env_enabled_default(FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, False) + + +def _env_int(name: str) -> Optional[int]: + value = os.environ.get(name) + if value is None or value == "": + return None + return int(value) + + +def _fp4_mla_attention_backend() -> str: + backend = os.getenv(FP4_MLA_ATTENTION_BACKEND_ENV) + if backend: + return backend.lower() + return _FP4_MLA_CUTEDSL_BACKEND if get_sm_version() == 107 else "triton" + + +def _cutedsl_backend_available() -> bool: + try: + return all( + importlib.util.find_spec(module) is not None + for module in ("ctm", "cutlass", "cuda.bindings.driver") + ) + except ModuleNotFoundError: + return False + + +def _fp4_mla_cutedsl_kernel_module() -> Any: + if _fp4_mla_cutedsl_fused_v_transpose_enabled(): + from . import fp4_mla_cutedsl_mufu16_fused_v_transpose + + return fp4_mla_cutedsl_mufu16_fused_v_transpose + + from . import fp4_mla_cutedsl_mufu16 + + return fp4_mla_cutedsl_mufu16 + + +def _ceil_div(lhs: int, rhs: int) -> int: + return (lhs + rhs - 1) // rhs + + +def _fp4_mla_q1_kv_blocks_per_program(num_gen: int, v_head_dim: int) -> int: + """Select the Q1 KV work per program without changing the launch boundary.""" + large_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH * FP4_BLOCK_SIZE + if num_gen >= _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD and v_head_dim % large_batch_block_dim == 0: + return _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH + medium_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH * FP4_BLOCK_SIZE + if ( + num_gen >= _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD + and v_head_dim % medium_batch_block_dim == 0 + ): + return _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH + small_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH * FP4_BLOCK_SIZE + if v_head_dim % small_batch_block_dim == 0: + return _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH + return 1 + + +def _fp4_mla_q1_prefix_blocks_per_program( + num_gen: int, + q1_kv_blocks_per_program: int, +) -> int: + """Select rolled Q-prefix work only when the batch keeps it efficient.""" + max_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM + if ( + num_gen >= _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD + and q1_kv_blocks_per_program == _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH + ): + if num_gen >= _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD: + return min(4, max_prefix_blocks) + return min(2, max_prefix_blocks) + return 1 + + +def _fp4_mla_q1_preload_variants( + max_num_sequences: int, + v_head_dim: int, +) -> tuple[tuple[int, int], ...]: + """Return every Q1 tuning variant reachable by the configured batch limit.""" + batch_sizes = [1] + for threshold in ( + _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD, + _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD, + _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD, + _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD, + ): + if threshold <= max_num_sequences: + batch_sizes.append(threshold) + + variants = [] + for batch_size in batch_sizes: + kv_blocks = _fp4_mla_q1_kv_blocks_per_program(batch_size, v_head_dim) + prefix_blocks = _fp4_mla_q1_prefix_blocks_per_program( + batch_size, + kv_blocks, + ) + variant = (kv_blocks, prefix_blocks) + if variant not in variants: + variants.append(variant) + return tuple(variants) + + +def _fp4_mla_triton_preload_key_set(metadata: Any) -> set[tuple[object, ...]]: + """Return the engine-scoped set of Triton variants loaded during warmup.""" + owner = getattr(metadata, "kv_cache_manager", None) + if owner is None: + owner = metadata + preload_keys = getattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, None) + if preload_keys is None: + preload_keys = set() + setattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, preload_keys) + return preload_keys diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cute_dsl_utils.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cute_dsl_utils.py new file mode 100644 index 000000000000..e64992ea56ac --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cute_dsl_utils.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Common CuTeDSL compilation and caller-stream handling for FP4 MLA.""" + +from collections.abc import Callable + +import cuda.bindings.driver as cuda +import cutlass.cute as cute +import torch +from cutlass.base_dsl.dsl import BaseDSL + + +def _compile_cutedsl(*args: object, **kwargs: object) -> Callable[..., object]: + """Compile the FP4 MLA kernels with their required PyIR frontend.""" + with BaseDSL.enable_pyir(): + return cute.compile(*args, **kwargs) + + +def _current_cu_stream() -> cuda.CUstream: + """Use the caller's stream for launch ordering and CUDA graph capture.""" + return cuda.CUstream(torch.cuda.current_stream().cuda_stream) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py new file mode 100644 index 000000000000..a61381cd1ed4 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py @@ -0,0 +1,1130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP4 MLA decode dispatch and Triton/CuTeDSL launch preparation.""" + +from typing import Any + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._utils import get_sm_version + +from .cache_update import _prepare_fp4_mla_q_buffers +from .config import ( + _FP4_MLA_CUTEDSL_BACKEND, + FP4_BLOCK_SIZE, + FP4_MLA_ATTENTION_BACKEND_ENV, + FP4_MLA_K_RESIDUAL_DIM, + FP4_MLA_P_GLOBAL_SCALE, + FP4_MLA_Q_LOGICAL_DIM, + FP4_MLA_Q_PACKED_DIM, + FP4_MLA_Q_RESIDUAL_DIM, + FP4_MLA_TOKENS_PER_BLOCK, + _ceil_div, + _cutedsl_backend_available, + _env_int, + _fp4_mla_attention_backend, + _fp4_mla_cutedsl_fused_v_transpose_enabled, + _fp4_mla_cutedsl_kernel_module, +) +from .layout import ( + _ensure_workspace_tensor, + _get_fp4_mla_global_scale, + _get_fp4_mla_kv_cache_tensors, + _get_fp4_mla_q_global_scale, + _get_fp4_mla_swizzled_scale_size, + _validate_fp4_mla_attention_q_shape, + _validate_fp4_mla_cache_shape, + _validate_fp4_mla_kv_storage_shape, + get_fp4_mla_v_scale_pool_view, +) +from .metadata import ( + _fp4_mla_generation_page_ids, + _fp4_mla_uniform_generation_lengths, + _get_linear_mtp_query_len_per_seq, + _infer_assume_full_pages, + _materialize_fp4_mla_device_page_table_for_forward, + _max_generation_pages, +) +from .v_cache import ( + _get_cutedsl_persistent_v_packed_cache, + _get_fp4_mla_v_packed_pool_base, + _get_fp4_mla_v_scale_pool_base, + _get_triton_v_packed_cache, + _select_triton_block_v, + _triton_can_prepack_v, + _triton_prepack_v_enabled, + _update_triton_v_packed_cache, +) + + +@triton.jit +def _cutedsl_swizzled_sf_offset(row_idx, col_idx, sf_cols: tl.constexpr): + padded_cols = ((sf_cols + 3) // 4) * 4 + return ( + col_idx % 4 + + (col_idx // 4) * (4 * 128) + + (row_idx % 32) * 16 + + ((row_idx % 128) // 32) * 4 + + (row_idx // 128) * (128 * padded_cols) + ) + + +@triton.jit +def _cutedsl_pad_q_and_sf_kernel( + q_padded_ptr, + q_ptr, + q_sf_padded_ptr, + q_sf_ptr, + num_heads, + output_heads: tl.constexpr, + packed_dim: tl.constexpr, + block_bytes: tl.constexpr, + sf_cols: tl.constexpr, + sf_cols_per_byte_block: tl.constexpr, +): + query_idx = tl.program_id(0) + byte_block = tl.program_id(1) + head_offsets = tl.arange(0, output_heads) + byte_offsets = byte_block * block_bytes + tl.arange(0, block_bytes) + head_mask = head_offsets < num_heads + byte_mask = byte_offsets < packed_dim + source_rows = query_idx * num_heads + head_offsets + destination_rows = query_idx * output_heads + head_offsets + values = tl.load( + q_ptr + source_rows[:, None] * packed_dim + byte_offsets[None, :], + mask=head_mask[:, None] & byte_mask[None, :], + other=0, + ) + tl.store( + q_padded_ptr + destination_rows[:, None] * packed_dim + byte_offsets[None, :], + values, + mask=byte_mask[None, :], + ) + sf_col_offsets = byte_block * sf_cols_per_byte_block + tl.arange(0, sf_cols_per_byte_block) + sf_col_mask = sf_col_offsets < sf_cols + source_offsets = _cutedsl_swizzled_sf_offset( + source_rows[:, None], sf_col_offsets[None, :], sf_cols + ) + destination_offsets = _cutedsl_swizzled_sf_offset( + destination_rows[:, None], sf_col_offsets[None, :], sf_cols + ) + sf_values = tl.load( + q_sf_ptr + source_offsets, + mask=head_mask[:, None] & sf_col_mask[None, :], + other=1.0, + ) + tl.store( + q_sf_padded_ptr + destination_offsets, + sf_values, + mask=sf_col_mask[None, :], + ) + + +_SM_COUNT_CACHE: dict[int, int] = {} + + +def _get_sm_count(device: torch.device) -> int: + """Return the SM (multiprocessor) count for ``device``, cached per index.""" + index = device.index if device.index is not None else torch.cuda.current_device() + count = _SM_COUNT_CACHE.get(index) + if count is None: + count = torch.cuda.get_device_properties(index).multi_processor_count + _SM_COUNT_CACHE[index] = count + return count + + +def _run_triton_attention_decode( + *, + metadata: Any, + layer_idx: int, + local_layer: int, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + src_page_ids: torch.Tensor, + kv_lens: torch.Tensor, + p_fp4: torch.Tensor, + p_sf: torch.Tensor, + max_scores: torch.Tensor, + denom: torch.Tensor, + output: torch.Tensor, + num_queries: int, + num_heads: int, + head_dim: int, + kv_lora_rank: int, + q_residual_dim: int, + query_len_per_seq: int, + max_pages: int, + sm_scale: float, + q_global_scale: torch.Tensor, +) -> None: + """Dispatch the ``triton`` FP4 MLA decode pipeline. + + Mirrors the four-stage layout used by ``fp4_mla_cutile.py`` + (page-stats with packed P -> reduce-stats -> prob-scale -> PV) but + routes through the self-contained kernels in + ``fp4_mla_triton.py``. Threads through the constexpr assume flags, + TMA descriptors, occupancy/num-warps launch meta, and pipelined PV loop. + """ + from .fp4_mla_triton import ( + _fp4_mla_attention_group_reduce_stats_kernel as _attn_group_reduce_stats_kernel, + ) + from .fp4_mla_triton import _fp4_mla_attention_page_stats_kernel as _attn_page_stats_kernel + from .fp4_mla_triton import _fp4_mla_attention_prob_scale_kernel as _attn_prob_scale_kernel + from .fp4_mla_triton import _fp4_mla_attention_pv_kernel as _attn_pv_kernel + from .fp4_mla_triton import ( + _fp4_mla_attention_pv_prepacked_v_kernel as _attn_pv_prepacked_v_kernel, + ) + from .fp4_mla_triton import _fp4_mla_attention_pv_reduce_kernel as _attn_pv_reduce_kernel + from .fp4_mla_triton import _fp4_mla_attention_reduce_stats_kernel as _attn_reduce_stats_kernel + + block_h = 128 + block_t = metadata.page_size + # Adaptive BLOCK_V: the fallback PV path uses a finer V split at small batch + # on B200 (~148 SMs). PV grid = num_queries * num_head_blocks(1) * + # (kv_lora_rank / BLOCK_V). We want >= ~2*num_SMs programs so that >1 CTA + # lands per SM and hides the L1TEX scoreboard stalls. Empirically (sweep): + # bs<=32 -> BLOCK_V=32; bs>=64 -> BLOCK_V=128. + # (BLOCK_V=16 is rejected by the V TMA descriptor min-stride requirement.) + # With prepacked V, BLOCK_V=128 avoids reloading the same P tile four times + # and matches the cutile prepacked-V tile shape. + block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) + q_storage_head_dim = head_dim + q_residual_dim + # The virtual GEMM tail evaluates QK + Q_r K + Q K_r in one reduction. + # Q and Q_r still occupy the 640-channel interleaved physical Q buffer; + # the final Q term reuses Q's main tail groups while K_r comes from the + # contiguous 64-channel tail of the primary paged KV cache. + q_head_dim = head_dim + q_residual_dim + FP4_MLA_K_RESIDUAL_DIM + # BLOCK_K = 512 aligns the K-window with the 512-channel non-residual prefix. + block_k = 512 + full_block_end = (q_head_dim // block_k) * block_k + tail_k = q_head_dim - full_block_end + tail_block_k = 1 << (tail_k - 1).bit_length() if tail_k > 0 else block_k + q_sf_per_token = q_storage_head_dim // FP4_BLOCK_SIZE + k_sf_per_token = (head_dim + FP4_MLA_K_RESIDUAL_DIM) // FP4_BLOCK_SIZE + sf_per_page = metadata.page_size // FP4_BLOCK_SIZE + num_head_blocks = triton.cdiv(num_heads, block_h) + + assume_full_heads = num_heads % block_h == 0 + assume_full_v = kv_lora_rank % block_v == 0 + # Match the cutile path: only mark pages "full" when we can prove every + # generation sequence has the same number of cached tokens AND + # query_len_per_seq == 1 (so the kv_len adjustment is a no-op). + assume_full_pages = ( + _infer_assume_full_pages(metadata, max_pages, metadata.page_size) and query_len_per_seq == 1 + ) + # Leave validity checks on. Matches cutile's default and is correctness- + # safe. The perfect-shape PV fast path (tl.ext.make_view + load_view_tko) + # remains gated off — when measured on the TileIR backend (ENABLE_TILE=1) + # it was net-slower on the bench, so the cost of enabling it isn't worth + # the win on the FP4 MLA shapes we care about. + assume_valid_pages = False + num_gen_seqs = num_queries // query_len_per_seq + if ( + not assume_valid_pages + and assume_full_pages + and src_page_ids.numel() == num_gen_seqs * max_pages + ): + assume_valid_pages = True + # cutile checks only `make_tensor_descriptor`; on the nvt backend the + # presence of TMA descriptors implies `tl.ext.make_view` is available too. + use_tma_data_load = hasattr(triton.language, "make_tensor_descriptor") + + # Install the device-side scratch allocator on every call. Triton stores + # the allocator in a ContextVar (triton.runtime._allocation), so a single + # process-wide install is not visible from worker threads / asyncio tasks + # that run with a different Context — the kernel launch would then hit the + # default NullAllocator and raise. Matches the cutile path. + if use_tma_data_load: + + def _tma_alloc(size: int, alignment: int, stream): + return torch.empty(size, device=q_fp4.device, dtype=torch.int8) + + triton.set_allocator(_tma_alloc) + + # cutile-equivalent launch meta. occupancy=2 lets two CTAs land per SM + # which improves wave-tail efficiency at the bs=32 hot point. + # NOTE: num_stages=2 (instead of the Triton 3.6 default of 3) sidesteps + # the TritonGPUAutomaticWarpSpecialization + NVWSInsertTmemAref pass that + # ICEs on the page_stats kernel under Triton 3.6.0 / sm_100. + launch_meta = {"occupancy": 2} + # The matmul kernels (page-stats QK and PV) are register-limited: at the + # Triton default of num_warps=4 the [BLOCK_H, BLOCK_T] epilogue spills the + # register file down to ~2 CTAs/SM (12.5% occupancy), so there are too few + # warps to hide the QK/PV load latency (ncu: ~0.3 eligible warps/scheduler). + # Spreading the tile epilogue over num_warps=8 halves the per-thread + # register need and roughly doubles resident warps. Matches the cutile + # ("nvt") backend, which launches page-stats at num_warps=8. Both are + # overridable for tuning. + sm_count = _get_sm_count(q_fp4.device) + # page-stats num_warps: the full-pages fast path (uniform q_len==1 decode) + # benefits from num_warps=8 (more warps hide the QK load latency); the + # masked path (q_len>1 / ragged lengths) carries extra per-thread state and + # measured markedly faster at num_warps=4 (e.g. bs256 q_len4: 131->95ms). + page_stats_num_warps = _env_int("TRTLLM_FP4_MLA_PAGE_STATS_NUM_WARPS") + if page_stats_num_warps is None: + page_stats_num_warps = 8 if assume_full_pages else 4 + page_stats_launch_meta = {"occupancy": 2, "num_warps": page_stats_num_warps} + # PV benefits from num_warps=8 across shapes measured. + pv_num_warps = _env_int("TRTLLM_FP4_MLA_PV_NUM_WARPS") or 8 + pv_launch_meta = {"occupancy": 2, "num_warps": pv_num_warps} + # PV loop pipelining. With TMA loads, num_stages>=2 lets the next page's + # loads overlap with the current MMA via mbarrier. The PV report shows + # long_scoreboard=4.5 cycles avg on V loads at PV_LOOP_STAGES=2; bumping the + # depth pays off when the grid is small enough that occupancy can absorb + # the extra in-flight tile state — i.e. medium batch / large max_pages. + # Larger pipelines hurt at small batch (more live state, fewer dim blocks). + if num_queries <= 16 or max_pages <= 4: + pv_loop_stages = 2 + else: + pv_loop_stages = 3 + + # Page-stats kernel: per (query, head_block, page) program, does QK, + # softmax stats, and packs probs into FP4 with the per-page local-max + # scaling trick. The page-max correction is applied later by + # prob_scale_kernel via p_sf in-place rescaling. + page_stats_shape = (num_queries, max_pages, num_heads) + page_max = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_page_max_buf", + page_stats_shape, + dtype=torch.float32, + device=q_fp4.device, + ) + page_sum = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_page_sum_buf", + page_stats_shape, + dtype=torch.float32, + device=q_fp4.device, + ) + + pack_prob_in_page_stats = True + _attn_page_stats_kernel[(num_queries, num_head_blocks, max_pages)]( + page_max, + page_sum, + p_fp4, + p_sf, + q_fp4, + q_sf, + kv_cache, + sf_cache, + global_scale, + q_global_scale, + src_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + q_fp4.stride(0), + q_fp4.stride(1), + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + page_max.stride(0), + page_max.stride(1), + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + q_fp4.shape[0], + sm_scale, + NUM_HEADS=num_heads, + Q_HEAD_D=q_head_dim, + Q_STORAGE_HEAD_D=q_storage_head_dim, + K_HEAD_D=head_dim, + Q_RESIDUAL_D=q_residual_dim, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + Q_SF_PER_TOKEN=q_sf_per_token, + K_SF_PER_TOKEN=k_sf_per_token, + SF_PER_PAGE=sf_per_page, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + BLOCK_H=block_h, + BLOCK_T=block_t, + BLOCK_K=block_k, + FULL_BLOCK_END=full_block_end, + TAIL_BLOCK_K=tail_block_k, + USE_TMA_DATA_LOAD=use_tma_data_load, + PACK_PROBS=pack_prob_in_page_stats, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_VALID_PAGES=assume_valid_pages, + **page_stats_launch_meta, + ) + # Two-level softmax-stats reduction. The single-level reduce launched only + # (num_queries * num_head_blocks) CTAs, each serially walking all max_pages + # twice -- at small batch that handful of CTAs left the GPU almost idle and + # the reduce cost more than the QK matmul. Level 1 parallelizes the page + # reduction across a page-group axis (online-softmax partials, pipelined); + # level 2 reuses the existing reduce kernel to fold the few groups into the + # global (max, denom). When the (query, head) grid already fills the GPU the + # group count collapses to 1 and this degenerates to the original reduce. + seqhead_ctas = num_queries * num_head_blocks + # Aim for ~3 waves of level-1 CTAs so page loads have enough memory-level + # parallelism to hide latency, while keeping the group count small enough + # that the level-2 combine loop stays short. + target_l1_ctas = 3 * sm_count + num_reduce_groups = _ceil_div(target_l1_ctas, max(seqhead_ctas, 1)) + num_reduce_groups = max(1, min(num_reduce_groups, max_pages, 64)) + # The grouped (two-level) reduce needs an auxiliary workspace, and + # _ensure_workspace_tensor can only (re)allocate it outside CUDA graph + # capture. If a warmup forward did not already size that workspace (e.g. the + # warmup batch took the single-level path), fall back to the single-level + # reduce during capture so we never allocate mid-capture. The single-level + # reduce is numerically identical (it just launches fewer CTAs). + if num_reduce_groups > 1 and torch.cuda.is_current_stream_capturing(): + gmax = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_max_buf") + gsum = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_sum_buf") + groups_ready = ( + gmax is not None + and gsum is not None + and gmax.shape[0] >= num_queries + and gmax.shape[1] >= num_reduce_groups + and gmax.shape[2] >= num_heads + and gsum.shape[0] >= num_queries + and gsum.shape[1] >= num_reduce_groups + and gsum.shape[2] >= num_heads + ) + if not groups_ready: + num_reduce_groups = 1 + if num_reduce_groups <= 1: + _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( + max_scores, + denom, + page_max, + page_sum, + max_pages, + max_scores.stride(0), + page_max.stride(0), + page_max.stride(1), + NUM_HEADS=num_heads, + MAX_PAGES=max_pages, + BLOCK_H=block_h, + **launch_meta, + ) + else: + group_pages = _ceil_div(max_pages, num_reduce_groups) + num_reduce_groups = _ceil_div(max_pages, group_pages) + group_max = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_group_max_buf", + (num_queries, num_reduce_groups, num_heads), + dtype=torch.float32, + device=q_fp4.device, + ) + group_sum = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_group_sum_buf", + (num_queries, num_reduce_groups, num_heads), + dtype=torch.float32, + device=q_fp4.device, + ) + _attn_group_reduce_stats_kernel[(num_queries, num_head_blocks, num_reduce_groups)]( + group_max, + group_sum, + page_max, + page_sum, + max_pages, + group_max.stride(0), + group_max.stride(1), + page_max.stride(0), + page_max.stride(1), + NUM_HEADS=num_heads, + GROUP_PAGES=group_pages, + BLOCK_H=block_h, + PIPELINE_STAGES=min(group_pages, 4), + **launch_meta, + ) + _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( + max_scores, + denom, + group_max, + group_sum, + num_reduce_groups, + max_scores.stride(0), + group_max.stride(0), + group_max.stride(1), + NUM_HEADS=num_heads, + MAX_PAGES=num_reduce_groups, + BLOCK_H=block_h, + **launch_meta, + ) + _attn_prob_scale_kernel[(num_queries, num_head_blocks, max_pages)]( + p_sf, + max_scores, + denom, + page_max, + metadata.fp4_mla_state.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + max_scores.stride(0), + page_max.stride(0), + page_max.stride(1), + NUM_HEADS=num_heads, + PAGE_SIZE=metadata.page_size, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + BLOCK_H=block_h, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_VALID_PAGES=assume_valid_pages, + **launch_meta, + ) + num_dim_blocks = triton.cdiv(kv_lora_rank, block_v) + v_packed = _get_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + v_head_dim=kv_lora_rank, + page_size=metadata.page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=src_page_ids, + ) + if ( + v_packed is None + and _triton_can_prepack_v(kv_lora_rank, metadata.page_size, block_v) + and not torch.cuda.is_current_stream_capturing() + ): + v_packed = _update_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + src_page_ids, + v_head_dim=kv_lora_rank, + page_size=metadata.page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + ) + use_triton_v_packed_cache = v_packed is not None + + # PV page split: partition the page range across additional programs and + # reduce in a follow-up kernel. ncu showed PV at waves/SM=0.49 for bs=32 — + # PV is L1-bandwidth bound, so raising in-flight CTAs is the lever. + # BLOCK_V is bounded below by the 16-byte TMA descriptor min-stride. + # PV page split: ncu shows that with the current shape (bs=32, max_pages=256) + # the PV kernel is L1-cache-throughput bound (long_scoreboard=4.5 cycles + # avg, L1 global LD hit-rate <40%). Increasing the program count via page + # splitting reduced waves/SM idle time but did NOT improve wall-time at + # current shapes — the per-CTA L1 thrash is the limit. Gate the split off + # by default; re-enable only for very small grids where occupancy is the + # bottleneck rather than per-CTA L1 pressure. + page_split = 1 + base_grid = num_queries * num_head_blocks * num_dim_blocks + if max_pages >= 16 and base_grid < 148: + for p in (8, 4, 2): + if max_pages % p == 0 and max_pages // p >= 16 and base_grid * p <= 148 * 4: + page_split = p + break + # The page-split PV path needs a partial-output workspace, which + # _ensure_workspace_tensor can only (re)allocate outside CUDA graph capture. + # Fall back to the unsplit PV (numerically identical) during capture unless a + # warmup forward already sized that workspace, so capture never allocates. + if page_split > 1 and torch.cuda.is_current_stream_capturing(): + pbuf = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_pv_partial_buf") + partial_ready = ( + pbuf is not None + and pbuf.shape[0] >= num_queries + and pbuf.shape[1] >= page_split + and pbuf.shape[2] >= num_heads + and pbuf.shape[3] >= kv_lora_rank + ) + if not partial_ready: + page_split = 1 + if page_split > 1: + pages_per_split = max_pages // page_split + partial_out = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_pv_partial_buf", + (num_queries, page_split, num_heads, kv_lora_rank), + dtype=torch.float32, + device=q_fp4.device, + ) + if use_triton_v_packed_cache: + _attn_pv_prepacked_v_kernel[ + (num_queries, num_head_blocks, num_dim_blocks * page_split) + ]( + output, + p_fp4, + p_sf, + v_packed, + v_sf, + global_scale, + src_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + PAGE_SPLIT=page_split, + PAGES_PER_SPLIT=pages_per_split, + PARTIAL_OUT=True, + partial_out_ptr=partial_out, + partial_s0=partial_out.stride(0), + partial_s1=partial_out.stride(1), + partial_s2=partial_out.stride(2), + partial_s3=partial_out.stride(3), + **pv_launch_meta, + ) + else: + _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks * page_split)]( + output, + p_fp4, + p_sf, + kv_cache, + kv_cache, + v_sf, + global_scale, + src_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, + USE_PREPACKED_V=False, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + PAGE_SPLIT=page_split, + PAGES_PER_SPLIT=pages_per_split, + PARTIAL_OUT=True, + partial_out_ptr=partial_out, + partial_s0=partial_out.stride(0), + partial_s1=partial_out.stride(1), + partial_s2=partial_out.stride(2), + partial_s3=partial_out.stride(3), + **pv_launch_meta, + ) + _attn_pv_reduce_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( + output, + partial_out, + global_scale, + output.stride(0), + output.stride(1), + output.stride(2), + partial_out.stride(0), + partial_out.stride(1), + partial_out.stride(2), + partial_out.stride(3), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SPLIT=page_split, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_V=assume_full_v, + **launch_meta, + ) + else: + if use_triton_v_packed_cache: + _attn_pv_prepacked_v_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( + output, + p_fp4, + p_sf, + v_packed, + v_sf, + global_scale, + src_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + **pv_launch_meta, + ) + else: + _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( + output, + p_fp4, + p_sf, + kv_cache, + kv_cache, + v_sf, + global_scale, + src_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, + USE_PREPACKED_V=False, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + **pv_launch_meta, + ) + + +def run_fp4_mla_attention_decode( + metadata: Any, + layer_idx: int, + local_layer: int, + q: torch.Tensor, + output: torch.Tensor, + *, + sm_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, + prequantized_q: torch.Tensor, + prequantized_q_sf: torch.Tensor, + q_batch_capacity: int, +) -> None: + """Run MLA decode with FP4 QK and FP4 PV tensor-core matmuls. + + Q is supplied in its assembled ``[latent, RoPE]`` layout and quantized to + FP4 directly. QK reads ``[KV-nope, K-RoPE, K-RoPE-residual]`` contiguously + from the primary cache with swizzled block scales. Softmax probabilities + are quantized to FP4 per page, and PV repacks V nibbles from the shared KV + cache while reading the auxiliary V-view scale pool. No BF16 dequantized + KV workspace is materialized on this path. Callers must supply the packed Q + and scales produced by the fused generation cache update. + """ + head_dim = kv_lora_rank + qk_rope_head_dim + if qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM: + raise ValueError( + "FP4 MLA K residual attention requires " + f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got {qk_rope_head_dim}." + ) + _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) + if metadata.page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + f"FP4 MLA attention decode requires page_size={FP4_MLA_TOKENS_PER_BLOCK}, " + f"got {metadata.page_size}." + ) + + if q.ndim != 3 or q.shape[-1] != head_dim: + raise ValueError( + "FP4 MLA attention Q must have shape " + f"[tokens, heads, {head_dim}], got {tuple(q.shape)}." + ) + if not q.is_contiguous(): + raise ValueError("FP4 MLA attention Q must be contiguous.") + + num_queries = q.shape[0] + if num_queries == 0: + raise ValueError("FP4 MLA attention decode requires at least one query token.") + num_gen_seqs = metadata.num_seqs - metadata.num_contexts + query_len_per_seq = _get_linear_mtp_query_len_per_seq( + metadata, + num_queries=num_queries, + num_gen_seqs=num_gen_seqs, + ) + + num_heads = q.shape[1] + if output.shape[:2] != (num_queries, num_heads): + raise ValueError("FP4 MLA attention output batch dimensions do not match.") + + backend = _fp4_mla_attention_backend() + if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: + raise RuntimeError( + "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." + ) + + global_scale = _get_fp4_mla_global_scale(metadata, q.device) + q_residual_dim = FP4_MLA_Q_RESIDUAL_DIM + _validate_fp4_mla_attention_q_shape(head_dim, q_residual_dim) + + if prequantized_q is None or prequantized_q_sf is None or q_batch_capacity is None: + raise RuntimeError( + "FP4 MLA decode requires Q prequantized by the fused generation cache update." + ) + + capacity = int(q_batch_capacity) + expected_q_shape = (capacity * num_heads, FP4_MLA_Q_PACKED_DIM) + expected_q_sf_shape = ( + _get_fp4_mla_swizzled_scale_size( + capacity * num_heads, + FP4_MLA_Q_LOGICAL_DIM, + ), + ) + if ( + q.dtype != torch.bfloat16 + or capacity <= 0 + or num_queries > capacity + or tuple(prequantized_q.shape) != expected_q_shape + or prequantized_q.dtype != torch.uint8 + or not prequantized_q.is_contiguous() + or tuple(prequantized_q_sf.shape) != expected_q_sf_shape + or prequantized_q_sf.dtype != torch.float8_e4m3fn + or not prequantized_q_sf.is_contiguous() + ): + raise ValueError("FP4 MLA prequantized Q does not satisfy the fused-Q contract.") + active_q_rows = num_queries * num_heads + active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( + active_q_rows, + FP4_MLA_Q_LOGICAL_DIM, + ) + q_fp4 = prequantized_q[:active_q_rows] + q_sf = prequantized_q_sf[:active_q_sf_bytes] + q_global_scale = _get_fp4_mla_q_global_scale(metadata, q.device) + + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=backend, + ) + sf_cache = sf_cache.view(torch.float8_e4m3fn) + + v_sf_pool = metadata.fp4_mla_state.v_scale_pool + v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=kv_lora_rank)[local_layer].view( + torch.float8_e4m3fn + ) + # The kv_lens runtime alias can lag at the decode anchor (seq_lens == 1) under + # CUDA graph / one-engine MTP; recover the true total per sequence so the + # per-query causal masking sees the full 1 + draft_len window (no-op when the + # alias already matches). + kv_lens, _ = _fp4_mla_uniform_generation_lengths(metadata, num_queries, num_gen_seqs) + _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens) + src_page_ids = _fp4_mla_generation_page_ids(metadata, num_gen_seqs) + max_pages = _max_generation_pages(metadata) + if max_pages == 0: + raise RuntimeError("FP4 MLA attention decode requires generation cache pages.") + if backend == _FP4_MLA_CUTEDSL_BACKEND: + if get_sm_version() != 107: + raise RuntimeError( + "FP4 MLA cutedsl attention backend requires Rubin SM107; " + f"current architecture is SM{get_sm_version()}." + ) + if not _cutedsl_backend_available(): + raise RuntimeError( + "FP4 MLA cutedsl attention backend requires the Rubin CTM and " + "CuTeDSL runtime packages." + ) + if not 0 < num_heads <= 128 or kv_lora_rank != 512: + raise ValueError( + "FP4 MLA cutedsl attention requires 1-128 local heads and " + f"kv_lora_rank=512, got num_heads={num_heads}, " + f"kv_lora_rank={kv_lora_rank}." + ) + + cutedsl_kernel = _fp4_mla_cutedsl_kernel_module() + QK_LOGICAL_DIM = cutedsl_kernel.QK_LOGICAL_DIM + QK_SF_GROUPS = cutedsl_kernel.QK_SF_GROUPS + SMEM_P4_V_N_PER_CTA = cutedsl_kernel.SMEM_P4_V_N_PER_CTA + run_trtllm_fp4_mla_decode_page_native_from_raw = ( + cutedsl_kernel.run_trtllm_fp4_mla_decode_page_native_from_raw + ) + + physical_heads = 128 + kernel_q = q_fp4 + kernel_q_sf = q_sf + if num_heads < physical_heads: + kernel_q_storage, kernel_q_sf_storage, q_batch_capacity = _prepare_fp4_mla_q_buffers( + metadata, + num_queries, + physical_heads, + q.device, + ) + active_q_rows = num_queries * physical_heads + active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( + active_q_rows, + QK_LOGICAL_DIM, + ) + kernel_q = kernel_q_storage[:active_q_rows] + kernel_q_sf = kernel_q_sf_storage[:active_q_sf_bytes] + _cutedsl_pad_q_and_sf_kernel[(num_queries, _ceil_div(QK_LOGICAL_DIM // 2, 64))]( + kernel_q, + q_fp4, + kernel_q_sf, + q_sf, + num_heads, + output_heads=physical_heads, + packed_dim=QK_LOGICAL_DIM // 2, + block_bytes=64, + sf_cols=QK_SF_GROUPS, + sf_cols_per_byte_block=8, + ) + + fused_v_transpose = _fp4_mla_cutedsl_fused_v_transpose_enabled() + if fused_v_transpose: + # The fusion kernel reads V from the canonical KV cache and uses + # the current layer V scales directly. Keep a None placeholder so + # the mufu16 and fused-V launchers share one Python call site. + core_v_packed = None + core_v_sf = v_sf + v_page_offset = 0 + else: + v_packed = _get_cutedsl_persistent_v_packed_cache( + metadata, + local_layer, + kv_cache, + v_head_dim=kv_lora_rank, + page_size=metadata.page_size, + block_v=SMEM_P4_V_N_PER_CTA, + ) + core_v_packed = _get_fp4_mla_v_packed_pool_base(metadata) + if core_v_packed is None: + raise RuntimeError("Persistent FP4 MLA V packing requires a stable full-pool base.") + get_v_page_offset = getattr( + metadata.kv_cache_manager, "get_mla_v_packed_page_offset", None + ) + v_page_offset = ( + int(get_v_page_offset(local_layer)) + if callable(get_v_page_offset) + else local_layer * kv_cache.shape[0] + ) + page_bytes = kv_lora_rank * (metadata.page_size // 2) + expected_layer_ptr = core_v_packed.data_ptr() + v_page_offset * page_bytes + if expected_layer_ptr != v_packed.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-packed layer view does not match its " + "full-pool base and page offset." + ) + + v_sf_pool_base = _get_fp4_mla_v_scale_pool_base(metadata) + if v_sf_pool_base is None: + v_sf_pool_base = v_sf_pool.flatten(0, 1).view(torch.uint8) + if v_sf_pool_base.data_ptr() != v_sf_pool.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-scale pool must flatten without a copy." + ) + if ( + not isinstance(v_sf_pool_base, torch.Tensor) + or v_sf_pool_base.dtype != torch.uint8 + or v_sf_pool_base.device != v_sf.device + or v_sf_pool_base.ndim != 2 + or v_sf_pool_base.shape[1] != v_sf_pool.shape[-1] + or not v_sf_pool_base.is_contiguous() + ): + raise RuntimeError( + "Persistent FP4 MLA V-scale pool base must be a contiguous " + "two-dimensional uint8 tensor with the configured page stride." + ) + core_v_sf = v_sf_pool_base.view(torch.float8_e4m3fn) + get_v_sf_page_offset = getattr( + metadata.kv_cache_manager, "get_mla_v_scale_page_offset", None + ) + v_sf_page_offset = ( + int(get_v_sf_page_offset(local_layer)) + if callable(get_v_sf_page_offset) + else local_layer * kv_cache.shape[0] + ) + if v_sf_page_offset != v_page_offset: + raise RuntimeError( + "Persistent FP4 MLA V-packed and V-scale pools require " + "matching encoded layer offsets." + ) + expected_v_sf_ptr = ( + core_v_sf.data_ptr() + + v_sf_page_offset * v_sf_pool.stride(1) * v_sf_pool.element_size() + ) + if expected_v_sf_ptr != v_sf.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-scale layer view does not match its " + "full-pool base and page offset." + ) + + kernel_output = output + if num_heads < physical_heads: + kernel_output = _ensure_workspace_tensor( + metadata, + "_fp4_mla_cutedsl_output_buf", + (num_queries, physical_heads, kv_lora_rank), + dtype=output.dtype, + device=output.device, + ) + + run_trtllm_fp4_mla_decode_page_native_from_raw( + kernel_q, + kernel_q_sf, + kv_cache, + sf_cache, + core_v_packed, + core_v_sf, + global_scale, + src_page_ids, + metadata.fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], + kv_lens, + kernel_output, + max_kv_len=max_pages * metadata.page_size, + sm_scale=float(sm_scale), + num_heads=physical_heads, + q_global_scale=q_global_scale, + page_size=metadata.page_size, + query_len_per_seq=query_len_per_seq, + v_page_offset=v_page_offset, + q_batch_capacity=q_batch_capacity, + partition_runtime_valid_k=bool(getattr(metadata, "is_cuda_graph", False)), + ) + if kernel_output is not output: + output.copy_(kernel_output[:, :num_heads]) + return + + total_p_rows = num_queries * max_pages * num_heads + p_fp4 = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_p_buf", + (max(total_p_rows, 1), metadata.page_size // 2), + dtype=torch.uint8, + device=q.device, + )[:total_p_rows] + p_sf = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_p_sf_buf", + (max(_get_fp4_mla_swizzled_scale_size(total_p_rows, metadata.page_size), 1),), + dtype=torch.float8_e4m3fn, + device=q.device, + ) + stats_shape = (num_queries, num_heads) + max_scores = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_max_buf", + stats_shape, + dtype=torch.float32, + device=q.device, + ) + denom = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_denom_buf", + stats_shape, + dtype=torch.float32, + device=q.device, + ) + + if backend != "triton": + raise ValueError( + f"Unsupported FP4 MLA attention backend '{backend}'. " + f"Set {FP4_MLA_ATTENTION_BACKEND_ENV} to 'triton' or " + f"'{_FP4_MLA_CUTEDSL_BACKEND}'." + ) + + # Self-contained public-Triton path: TMA-loaded QK + fused page-stats pack, + # reduce-stats, prob-scale, and PV with an optional prepacked V cache. + _run_triton_attention_decode( + metadata=metadata, + layer_idx=layer_idx, + local_layer=local_layer, + q_fp4=q_fp4, + q_sf=q_sf.contiguous().view(-1), + kv_cache=kv_cache, + sf_cache=sf_cache, + v_sf=v_sf, + global_scale=global_scale, + src_page_ids=src_page_ids, + kv_lens=kv_lens, + p_fp4=p_fp4, + p_sf=p_sf, + max_scores=max_scores, + denom=denom, + output=output, + num_queries=num_queries, + num_heads=num_heads, + head_dim=head_dim, + kv_lora_rank=kv_lora_rank, + q_residual_dim=q_residual_dim, + query_len_per_seq=query_len_per_seq, + max_pages=max_pages, + sm_scale=float(sm_scale), + q_global_scale=q_global_scale, + ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py index 4b50fa3de793..df30ec7143cc 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py @@ -21,6 +21,7 @@ import triton from tensorrt_llm._utils import get_sm_version, prefer_pinned +from tensorrt_llm.bindings import DataType from tensorrt_llm.quantization.mode import QuantMode from .fp4_mla_kernels import _fp8_mla_context_block_table_kernel @@ -92,6 +93,7 @@ class _Fp8MlaContextCacheManagerView: kv_cache_pool_pointers: torch.Tensor kv_cache_pool_mapping: torch.Tensor layer_offsets: Tuple[int, ...] + dtype: DataType = DataType.FP8 @dataclass @@ -276,23 +278,19 @@ def prepare(self, meta: "TrtllmAttentionMetadata") -> None: def _build_fp8_mla_context_attn(attn: "TrtllmAttention") -> "TrtllmAttention": """Build a direct-attribute FP8 view without per-access Python forwarding.""" - from ..fmha.fallback import FallbackFmha from ..fmha.manager import FmhaManager fp8_attn = copy.copy(attn) fp8_attn.quant_mode = int(QuantMode(0).set_fp8_kv_cache()) fp8_attn.has_fp4_kv_cache = False fp8_attn.has_fp8_kv_cache = True - # FMHA instances hold weak references to their owning attention object. - # Do not reuse the manager copied from the FP4 attention; bind this FP8 - # view explicitly to TRTLLM's regular FMHA implementation. - fp8_manager = FmhaManager(fp8_attn) - fp8_manager.fmha_libs = [FallbackFmha(fp8_attn)] - fp8_attn._fmha_manager = fp8_manager fp8_attn.local_layer_idx = 0 # This branch resolves local cache layers through layer_idx. The # disposable cache has exactly one layer, so bind the copied view to it. fp8_attn.layer_idx = 0 + # Finalize the view before capability selection. FMHA instances must hold + # weak references to this FP8 view, not to the original FP4 attention. + fp8_attn._fmha_manager = FmhaManager(fp8_attn) return fp8_attn diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index f33349506a9d..2dbb74741e8c 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -1,9 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import contextlib import math -import os import sys import threading from collections import OrderedDict @@ -19,27 +17,19 @@ import cutlass.pipeline as pipeline import torch from cutlass._mlir.dialects import llvm -from cutlass.base_dsl.dsl import BaseDSL from cutlass.cute.arch.nvvm_wrappers import inline_ptx as cute_inline_ptx from cutlass.cute.runtime import make_ptr from cutlass.experimental import cuda as cuda_tma from cutlass.experimental import primitives as prims +from .cute_dsl_utils import _compile_cutedsl, _current_cu_stream + nvvm_add_packed_f32x2 = partial(prims.add_packed_f32x2, rnd=prims.FPRoundingMode.RN) nvvm_mul_packed_f32x2 = partial(prims.mul_packed_f32x2, rnd=prims.FPRoundingMode.RN) nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 INT32_MAX = (1 << 31) - 1 -_CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" -_PYIR_STDOUT_LINES = frozenset( - { - "Enabling PyIR, it was False", - "Enabling PyIR, it is now True", - "Disabling PyIR, it was True", - "Disabling PyIR, it is now False", - } -) @ctm.dsl_user_op @@ -67,45 +57,6 @@ def _mbarrier_arrive_release_cta_shared_cluster(mbar, count=1, *, loc=None, ip=N ) -class _PyIRStdoutFilter: - def __init__(self, output): - self._output = output - self._pending = "" - - def write(self, text): - lines = (self._pending + text).split("\n") - self._pending = lines.pop() - for line in lines: - if line.rstrip("\r") not in _PYIR_STDOUT_LINES: - self._output.write(f"{line}\n") - return len(text) - - def flush(self): - self._output.flush() - - def finish(self): - if self._pending and self._pending.rstrip("\r") not in _PYIR_STDOUT_LINES: - self._output.write(self._pending) - self._pending = "" - self._output.flush() - - def __getattr__(self, name): - return getattr(self._output, name) - - -def _compile_cutedsl(*args, **kwargs): - verbose = os.getenv(_CUTEDSL_VERBOSE_COMPILE_ENV, "").strip().lower() - if verbose in {"1", "true", "yes", "on"}: - with BaseDSL.enable_pyir(): - return cute.compile(*args, **kwargs) - stdout_filter = _PyIRStdoutFilter(sys.stdout) - try: - with contextlib.redirect_stdout(stdout_filter), BaseDSL.enable_pyir(): - return cute.compile(*args, **kwargs) - finally: - stdout_filter.finish() - - def _scale_words_for_k(k_dim: int, sf_vec_size: int) -> int: return (k_dim + sf_vec_size * 4 - 1) // (sf_vec_size * 4) @@ -500,11 +451,6 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_QK_COMPLETION_MBARS = SMEM_P4_TMEM_SCORE_PIPELINE_STAGES -def _current_cu_stream() -> cuda.CUstream: - """Use the caller's stream for launch ordering and CUDA graph capture.""" - return cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - @dataclass(frozen=True) class KvCache3DLayout: num_pages: int diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py index c7a9d73120ad..6e6ded7d99be 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py @@ -2,9 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # ruff: noqa: E501, E741, F841 -import contextlib import math -import os import sys import threading from collections import OrderedDict @@ -21,27 +19,19 @@ import torch from cutlass._mlir.dialects import llvm from cutlass.base_dsl.array import EvictPriority -from cutlass.base_dsl.dsl import BaseDSL from cutlass.cute.arch.nvvm_wrappers import inline_ptx as cute_inline_ptx from cutlass.cute.runtime import make_ptr from cutlass.experimental import cuda as cuda_tma from cutlass.experimental import primitives as prims +from .cute_dsl_utils import _compile_cutedsl, _current_cu_stream + nvvm_add_packed_f32x2 = partial(prims.add_packed_f32x2, rnd=prims.FPRoundingMode.RN) nvvm_mul_packed_f32x2 = partial(prims.mul_packed_f32x2, rnd=prims.FPRoundingMode.RN) nvvm_fma_packed_f32x2 = partial(prims.fma_packed_f32x2, rnd=prims.FPRoundingMode.RN) PREPARED_BUFFER_ALIGNMENT_BYTES = 32 CUDA_GRID_Z_MAX = 65535 INT32_MAX = (1 << 31) - 1 -_CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" -_PYIR_STDOUT_LINES = frozenset( - { - "Enabling PyIR, it was False", - "Enabling PyIR, it is now True", - "Disabling PyIR, it was True", - "Disabling PyIR, it is now False", - } -) @ctm.dsl_user_op @@ -71,45 +61,6 @@ def _mbarrier_arrive_shared_cluster(mbar, count=1, *, loc=None, ip=None) -> None ) -class _PyIRStdoutFilter: - def __init__(self, output): - self._output = output - self._pending = "" - - def write(self, text): - lines = (self._pending + text).split("\n") - self._pending = lines.pop() - for line in lines: - if line.rstrip("\r") not in _PYIR_STDOUT_LINES: - self._output.write(f"{line}\n") - return len(text) - - def flush(self): - self._output.flush() - - def finish(self): - if self._pending and self._pending.rstrip("\r") not in _PYIR_STDOUT_LINES: - self._output.write(self._pending) - self._pending = "" - self._output.flush() - - def __getattr__(self, name): - return getattr(self._output, name) - - -def _compile_cutedsl(*args, **kwargs): - verbose = os.getenv(_CUTEDSL_VERBOSE_COMPILE_ENV, "").strip().lower() - if verbose in {"1", "true", "yes", "on"}: - with BaseDSL.enable_pyir(): - return cute.compile(*args, **kwargs) - stdout_filter = _PyIRStdoutFilter(sys.stdout) - try: - with contextlib.redirect_stdout(stdout_filter), BaseDSL.enable_pyir(): - return cute.compile(*args, **kwargs) - finally: - stdout_filter.finish() - - def _scale_words_for_k(k_dim: int, sf_vec_size: int) -> int: return (k_dim + sf_vec_size * 4 - 1) // (sf_vec_size * 4) @@ -502,11 +453,6 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_QK_COMPLETION_MBARS = SMEM_P4_TMEM_SCORE_PIPELINE_STAGES -def _current_cu_stream() -> cuda.CUstream: - """Use the caller's stream for launch ordering and CUDA graph capture.""" - return cuda.CUstream(torch.cuda.current_stream().cuda_stream) - - @dataclass(frozen=True) class KvCache3DLayout: num_pages: int diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py index 78418ab6ebfd..51c0fc3fdb2e 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_v_repack.py @@ -4,9 +4,6 @@ """CuTeDSL helpers for repacking the paged FP4 MLA V cache.""" -import contextlib -import os -import sys from collections.abc import Callable from dataclasses import dataclass @@ -14,73 +11,16 @@ import cutlass as ctm import cutlass.cute as cute import torch -from cutlass.base_dsl.dsl import BaseDSL from cutlass.cute.runtime import make_ptr from cutlass.experimental import cuda as cuda_exp from cutlass.experimental import primitives +from .cute_dsl_utils import _compile_cutedsl, _current_cu_stream + PREPARED_BUFFER_ALIGNMENT_BYTES = 32 TRTLLM_PAGE_SIZE = 128 SMEM_P4_V_N_PER_CTA = 128 -_CUTEDSL_VERBOSE_COMPILE_ENV = "TRTLLM_CUTEDSL_VERBOSE_COMPILE" -_PYIR_STDOUT_LINES = frozenset( - { - "Enabling PyIR, it was False", - "Enabling PyIR, it is now True", - "Disabling PyIR, it was True", - "Disabling PyIR, it is now False", - } -) - - -class _PyIRStdoutFilter: - """Drop only CuTeDSL PyIR state transitions from a text stream.""" - - def __init__(self, output): - self._output = output - self._pending = "" - - def write(self, text): - lines = (self._pending + text).split("\n") - self._pending = lines.pop() - for line in lines: - if line.rstrip("\r") not in _PYIR_STDOUT_LINES: - self._output.write(f"{line}\n") - return len(text) - - def flush(self): - self._output.flush() - - def finish(self): - if self._pending and self._pending.rstrip("\r") not in _PYIR_STDOUT_LINES: - self._output.write(self._pending) - self._pending = "" - self._output.flush() - - def __getattr__(self, name): - return getattr(self._output, name) - - -def _compile_cutedsl(*args, **kwargs): - """Compile with PyIR while suppressing only its state-transition lines.""" - verbose = os.getenv(_CUTEDSL_VERBOSE_COMPILE_ENV, "").strip().lower() - if verbose in {"1", "true", "yes", "on"}: - with BaseDSL.enable_pyir(): - return cute.compile(*args, **kwargs) - - stdout_filter = _PyIRStdoutFilter(sys.stdout) - try: - with contextlib.redirect_stdout(stdout_filter), BaseDSL.enable_pyir(): - return cute.compile(*args, **kwargs) - finally: - stdout_filter.finish() - - -def _current_cu_stream() -> cuda.CUstream: - """Use the caller's stream for launch ordering and CUDA graph capture.""" - return cuda.CUstream(torch.cuda.current_stream().cuda_stream) - @dataclass(frozen=True) class KvCache3DLayout: diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py new file mode 100644 index 000000000000..e4bdf6e2c739 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py @@ -0,0 +1,257 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP4 MLA tensor layout validation, scale views, and workspace allocation.""" + +from typing import Any + +import torch + +from .config import ( + _FP4_MLA_K_RESIDUAL_BACKENDS, + FP4_BLOCK_SIZE, + FP4_MLA_ATTENTION_BACKEND_ENV, + FP4_MLA_K_RESIDUAL_DIM, + FP4_MLA_SCALE_COL_GROUP, + FP4_MLA_SCALE_ROW_GROUP, + FP4_MLA_TOKENS_PER_BLOCK, + HP_BLOCK_SIZE, + _ceil_div, +) + +# FP4 MLA scale-layout helpers + + +def get_fp4_mla_v_scale_pool_size(v_head_dim: int, page_size: int) -> int: + """Return elements per page for the swizzled FP4 MLA V-scale pool. + + The PV matmul treats V as a RHS matrix shaped ``[v_head_dim, kv_tokens]``. + NVFP4 block scales therefore group along the token/K axis, not along the + latent dimension as the K-view cache does. The physical layout matches the + Triton block-scaled matmul scale layout: + ``[ceil(v_head_dim / 128), ceil(page_size / 16 / 4), 32, 16]``. + """ + return _get_fp4_mla_swizzled_scale_size(v_head_dim, page_size) + + +def _get_fp4_mla_swizzled_scale_size(rows: int, cols: int) -> int: + scale_cols = _ceil_div(cols, FP4_BLOCK_SIZE) + row_groups = _ceil_div(rows, FP4_MLA_SCALE_ROW_GROUP) + col_groups = _ceil_div(scale_cols, FP4_MLA_SCALE_COL_GROUP) + return row_groups * col_groups * 32 * 16 + + +def get_fp4_mla_v_scale_pool_shape( + num_layers: int, + num_pages: int, + v_head_dim: int, + page_size: int, +) -> tuple[int, int, int, int, int, int]: + """Return the logical swizzled V-scale view shape. + + The leading dimensions are ``[layer, physical_page]``. The remaining + dimensions are the preshuffled ``[N // 128, K // 16 // 4, 32, 16]`` shape + consumed by Triton block-scaled matmul for the V/PV RHS operand. + """ + token_scale_cols = _ceil_div(page_size, FP4_BLOCK_SIZE) + return ( + num_layers, + num_pages, + _ceil_div(v_head_dim, FP4_MLA_SCALE_ROW_GROUP), + _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP), + 32, + 16, + ) + + +def get_fp4_mla_v_scale_pool_view( + metadata: Any, + *, + v_head_dim: int, +) -> torch.Tensor: + """View the auxiliary MLA V-scale pool in Triton's block-scaled layout.""" + pool = getattr(metadata.fp4_mla_state, "v_scale_pool", None) + if pool is None: + raise RuntimeError("FP4 MLA V scale pool is not allocated.") + + elems_per_page = get_fp4_mla_v_scale_pool_size(v_head_dim, metadata.page_size) + if pool.shape[-1] < elems_per_page: + raise RuntimeError( + f"FP4 MLA V scale pool page stride is too small: got " + f"{pool.shape[-1]}, need {elems_per_page}." + ) + + token_scale_cols = _ceil_div(metadata.page_size, FP4_BLOCK_SIZE) + col_groups = _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP) + shape = get_fp4_mla_v_scale_pool_shape( + pool.shape[0], pool.shape[1], v_head_dim, metadata.page_size + ) + strides = ( + pool.stride(0), + pool.stride(1), + col_groups * 32 * 16, + 32 * 16, + 16, + 1, + ) + return torch.as_strided(pool, size=shape, stride=strides) + + +# Python launch helpers + + +def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: + global_scale = getattr(metadata.fp4_mla_state, "kv_global_scale", None) + if ( + not isinstance(global_scale, torch.Tensor) + or global_scale.device != device + or global_scale.dtype != torch.float32 + or global_scale.numel() != 1 + ): + raise RuntimeError("FP4 MLA requires a preallocated FP32 KV global-scale tensor.") + return global_scale + + +def _get_fp4_mla_q_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: + global_scale = getattr(metadata.fp4_mla_state, "q_global_scale", None) + if ( + not isinstance(global_scale, torch.Tensor) + or global_scale.device != device + or global_scale.dtype != torch.float32 + or global_scale.numel() != 1 + ): + raise RuntimeError("FP4 MLA requires a preallocated FP32 Q global-scale tensor.") + return global_scale + + +def _get_fp4_mla_kv_cache_tensors( + metadata: Any, layer_idx: int +) -> tuple[torch.Tensor, torch.Tensor]: + return metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) + + +def _get_fp4_mla_hp_pool_layout( + metadata: Any, + pool: torch.Tensor, +) -> tuple[int, int]: + """Return the manager-owned HP ring size and per-token head dimension.""" + manager = getattr(metadata, "kv_cache_manager", None) + if manager is None or not hasattr(manager, "fp4_mla_hp_pool_size"): + raise ValueError("FP4 MLA requires a V2 manager-owned HP ring.") + hp_pool_size = manager.fp4_mla_hp_pool_size + if ( + hp_pool_size < HP_BLOCK_SIZE + or pool.ndim != 4 + or pool.shape[2] < 1 + or pool.shape[-1] % hp_pool_size != 0 + ): + raise ValueError( + "FP4 MLA high-precision KV pool does not match its configured " + f"ring: shape={tuple(pool.shape)}, ring_size={hp_pool_size}." + ) + return hp_pool_size, pool.shape[-1] // hp_pool_size + + +def _validate_fp4_mla_hp_generation_width( + hp_pool_size: int, + generation_len: int, +) -> None: + """Ensure one target plus rewindable drafts fit without clobbering the live tail.""" + max_rewind_len = hp_pool_size - HP_BLOCK_SIZE + if generation_len <= 0 or generation_len - 1 > max_rewind_len: + raise RuntimeError( + "FP4 MLA generation exceeds the HP ring's rewind slack: " + f"generation={generation_len}, max_rewind={max_rewind_len}." + ) + + +def _validate_fp4_mla_kv_storage_shape( + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + *, + head_dim: int, + backend: str, +) -> int: + """Validate the backend-specific physical KV and scale strides.""" + residual_dim = FP4_MLA_K_RESIDUAL_DIM if backend in _FP4_MLA_K_RESIDUAL_BACKENDS else 0 + expected_storage_head_dim = head_dim + residual_dim + storage_head_dim = kv_cache.shape[-1] * 2 + if storage_head_dim != expected_storage_head_dim: + raise RuntimeError( + "FP4 MLA KV cache storage head dimension does not match the selected backend: " + f"got {storage_head_dim}, expected {expected_storage_head_dim}. Recreate the engine " + f"after setting {FP4_MLA_ATTENTION_BACKEND_ENV}." + ) + + expected_scale_columns = expected_storage_head_dim // FP4_BLOCK_SIZE + if sf_cache.shape[-1] != expected_scale_columns: + raise RuntimeError( + "FP4 MLA KV cache scale storage does not match the contiguous data layout: " + f"got {sf_cache.shape[-1]} columns, expected {expected_scale_columns}." + ) + return storage_head_dim + + +def _validate_fp4_mla_cache_shape(page_size: int, head_dim: int) -> None: + if page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + f"FP4 MLA KV cache requires tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK} " + f"for swizzled block scales, got {page_size}." + ) + + sf_per_token = head_dim // FP4_BLOCK_SIZE + if head_dim % FP4_BLOCK_SIZE != 0 or sf_per_token % 4 != 0: + raise ValueError( + f"FP4 MLA KV head_dim must produce a scale column count divisible by 4; " + f"got head_dim={head_dim}, scale_columns={sf_per_token}." + ) + + +def _validate_fp4_mla_attention_q_shape(head_dim: int, q_residual_dim: int) -> None: + if q_residual_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA Q residual_dim must be divisible by {FP4_BLOCK_SIZE}, got {q_residual_dim}." + ) + if q_residual_dim <= 0 or q_residual_dim > head_dim: + raise ValueError( + f"FP4 MLA Q residual_dim must be in (0, head_dim], got " + f"residual_dim={q_residual_dim}, head_dim={head_dim}." + ) + + q_head_dim = head_dim + q_residual_dim + q_sf_per_token = q_head_dim // FP4_BLOCK_SIZE + if q_head_dim % FP4_BLOCK_SIZE != 0 or q_sf_per_token % FP4_MLA_SCALE_COL_GROUP != 0: + raise ValueError( + f"FP4 MLA residual Q must produce a scale column count divisible " + f"by {FP4_MLA_SCALE_COL_GROUP}; got q_head_dim={q_head_dim}, " + f"scale_columns={q_sf_per_token}." + ) + + +def _ensure_workspace_tensor( + metadata: Any, + attr_name: str, + shape: tuple[int, ...], + *, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + workspaces = metadata.fp4_mla_state.workspaces + tensor = workspaces.get(attr_name) + needs_alloc = ( + tensor is None + or tensor.dtype != dtype + or tensor.device != device + or len(tensor.shape) != len(shape) + or any(tensor.shape[idx] < dim for idx, dim in enumerate(shape)) + ) + if needs_alloc: + if torch.cuda.is_current_stream_capturing(): + raise ValueError( + f"Cannot allocate {attr_name} while capturing a CUDA graph. " + "Run a warmup prepare/forward first." + ) + tensor = torch.empty(shape, dtype=dtype, device=device) + workspaces[attr_name] = tensor + + slices = tuple(slice(0, dim) for dim in shape) + return tensor[slices] diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py new file mode 100644 index 000000000000..b654dc0c2695 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py @@ -0,0 +1,897 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP4 MLA page tables, append metadata, and generation-length preparation.""" + +from typing import Any, Optional + +import torch +import triton +import triton.language as tl + +from .config import _FP4_MLA_PAGE_TABLE_TILE_SIZE, FP4_MLA_TOKENS_PER_BLOCK, _ceil_div + + +@triton.jit +def _fp4_mla_store_sequence_append_metadata( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + sequence_idx, + num_tokens, + PREFIX_BLOCK: tl.constexpr, + TOKEN_BLOCK: tl.constexpr, +): + append_len = tl.load(append_lens_ptr + sequence_idx) + prefix_offsets = tl.arange(0, PREFIX_BLOCK) + token_start = append_len - append_len + for prefix_start in tl.range(0, sequence_idx, PREFIX_BLOCK): + preceding_sequences = prefix_start + prefix_offsets + preceding_lens = tl.load( + append_lens_ptr + preceding_sequences, + mask=preceding_sequences < sequence_idx, + other=0, + ) + token_start += tl.sum(preceding_lens) + + cached_len = tl.load(kv_lens_ptr + sequence_idx) - append_len + + token_offsets = tl.arange(0, TOKEN_BLOCK) + for block_start in tl.range(0, append_len, TOKEN_BLOCK): + local_offsets = block_start + token_offsets + token_mask = (local_offsets < append_len) & (token_start + local_offsets < num_tokens) + output_offsets = token_start + local_offsets + tl.store( + batch_indices_ptr + output_offsets, + sequence_idx, + mask=token_mask, + ) + tl.store( + positions_ptr + output_offsets, + cached_len + local_offsets, + mask=token_mask, + ) + + +@triton.jit( + do_not_specialize=[ + "num_tokens", + "num_contexts", + "num_generation_sequences", + ], + do_not_specialize_on_alignment=[ + "num_tokens", + "num_contexts", + "num_generation_sequences", + ], +) +def _fp4_mla_append_metadata_kernel( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + num_tokens, + num_contexts, + num_generation_sequences, + ONE_TOKEN_GENERATION: tl.constexpr, + PREFIX_BLOCK: tl.constexpr, + TOKEN_BLOCK: tl.constexpr, + GENERATION_BLOCK: tl.constexpr, +): + program_idx = tl.program_id(0) + if ONE_TOKEN_GENERATION: + if program_idx < num_contexts: + _fp4_mla_store_sequence_append_metadata( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + program_idx, + num_tokens, + PREFIX_BLOCK, + TOKEN_BLOCK, + ) + else: + generation_offsets = (program_idx - num_contexts) * GENERATION_BLOCK + tl.arange( + 0, GENERATION_BLOCK + ) + generation_mask = generation_offsets < num_generation_sequences + sequence_indices = num_contexts + generation_offsets + output_offsets = num_tokens - num_generation_sequences + generation_offsets + generation_mask = generation_mask & (output_offsets < num_tokens) + generation_positions = ( + tl.load( + kv_lens_ptr + sequence_indices, + mask=generation_mask, + other=1, + ) + - 1 + ) + tl.store( + batch_indices_ptr + output_offsets, + sequence_indices, + mask=generation_mask, + ) + tl.store( + positions_ptr + output_offsets, + generation_positions, + mask=generation_mask, + ) + else: + _fp4_mla_store_sequence_append_metadata( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + program_idx, + num_tokens, + PREFIX_BLOCK, + TOKEN_BLOCK, + ) + + +def populate_fp4_mla_append_metadata( + append_lens: torch.Tensor, + kv_lens: torch.Tensor, + batch_indices: torch.Tensor, + positions: torch.Tensor, + *, + num_tokens: int, + num_sequences: int, + num_contexts: int, + num_context_tokens: int, +) -> None: + """Populate FP4 MLA token-to-sequence metadata in one Triton launch. + + Mixed batches vectorize their one-token generation rows. Multi-token MTP + and fallback shapes use the generic per-sequence path in the same kernel. + """ + if num_sequences <= 0 or num_tokens <= 0: + return + if not 0 <= num_contexts <= num_sequences: + raise ValueError( + f"FP4 MLA num_contexts must be in [0, {num_sequences}], got {num_contexts}." + ) + if not 0 <= num_context_tokens <= num_tokens: + raise ValueError( + f"FP4 MLA num_context_tokens must be in [0, {num_tokens}], got {num_context_tokens}." + ) + + tensors = ( + append_lens, + kv_lens, + batch_indices, + positions, + ) + if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): + raise ValueError("FP4 MLA append metadata tensors must be contiguous and one-dimensional.") + if any(tensor.dtype != torch.int32 for tensor in tensors): + raise TypeError("FP4 MLA append metadata tensors must use int32.") + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("FP4 MLA append metadata tensors must be CUDA tensors.") + if any(tensor.device != append_lens.device for tensor in tensors[1:]): + raise ValueError("FP4 MLA append metadata tensors must be on the same device.") + sequence_tensors = (append_lens, kv_lens) + if any(tensor.numel() < num_sequences for tensor in sequence_tensors): + raise ValueError( + f"FP4 MLA sequence metadata tensors need at least {num_sequences} entries." + ) + token_tensors = (batch_indices, positions) + if any(tensor.numel() < num_tokens for tensor in token_tensors): + raise ValueError(f"FP4 MLA token metadata tensors need at least {num_tokens} entries.") + + num_generation_sequences = num_sequences - num_contexts + # Each scheduled generation sequence appends at least one token. Equality + # therefore identifies the common mixed batch with one token per decode + # row without reading the device append lengths back on the host. + one_token_generation = num_tokens == num_context_tokens + num_generation_sequences + generation_block = 128 + grid = num_sequences + if one_token_generation: + grid = num_contexts + triton.cdiv(num_generation_sequences, generation_block) + + _fp4_mla_append_metadata_kernel[(grid,)]( + append_lens, + kv_lens, + batch_indices, + positions, + num_tokens, + num_contexts, + num_generation_sequences, + ONE_TOKEN_GENERATION=one_token_generation, + PREFIX_BLOCK=128, + TOKEN_BLOCK=256, + GENERATION_BLOCK=generation_block, + num_warps=4, + ) + + +@triton.jit( + do_not_specialize=["num_gen", "generation_len"], + do_not_specialize_on_alignment=["num_gen", "generation_len"], +) +def _fp4_mla_generation_lengths_kernel( + kv_lens_ptr, + prompt_lens_ptr, + corrected_kv_lens_ptr, + generation_lens_ptr, + num_gen, + generation_len, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < num_gen + kv_lens = tl.load(kv_lens_ptr + offsets, mask=mask, other=0) + prompt_lens = tl.load(prompt_lens_ptr + offsets, mask=mask, other=0) + tl.store( + corrected_kv_lens_ptr + offsets, + kv_lens - prompt_lens + generation_len, + mask=mask, + ) + tl.store(generation_lens_ptr + offsets, generation_len, mask=mask) + + +def populate_fp4_mla_generation_lengths( + kv_lens: torch.Tensor, + prompt_lens: torch.Tensor, + corrected_kv_lens: torch.Tensor, + generation_lens: torch.Tensor, + *, + num_gen_tokens: int, + num_gen: int, +) -> None: + """Populate reusable FP4 MLA generation lengths in one Triton launch.""" + if num_gen <= 0 or num_gen_tokens % num_gen != 0: + raise ValueError( + "FP4 MLA generation lengths require a positive sequence count and " + f"uniform token count, got {num_gen_tokens} tokens for {num_gen} sequences." + ) + tensors = (kv_lens, prompt_lens, corrected_kv_lens, generation_lens) + if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): + raise ValueError( + "FP4 MLA generation length tensors must be contiguous and one-dimensional." + ) + if any(tensor.dtype != torch.int32 for tensor in tensors): + raise TypeError("FP4 MLA generation length tensors must use int32.") + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("FP4 MLA generation length tensors must be CUDA tensors.") + if any(tensor.device != kv_lens.device for tensor in tensors[1:]): + raise ValueError("FP4 MLA generation length tensors must be on the same device.") + if any(tensor.numel() < num_gen for tensor in tensors): + raise ValueError(f"FP4 MLA generation length tensors need at least {num_gen} entries.") + + block = 128 + _fp4_mla_generation_lengths_kernel[(triton.cdiv(num_gen, block),)]( + kv_lens, + prompt_lens, + corrected_kv_lens, + generation_lens, + num_gen, + num_gen_tokens // num_gen, + BLOCK=block, + num_warps=4, + ) + + +def _fp4_mla_page_table_spec(kv_cache_manager: Any) -> Any: + get_spec = getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None) + if not callable(get_spec): + raise RuntimeError("FP4 MLA requires Fp4MlaKVCacheManagerV2 page metadata.") + spec = get_spec() + for field_name in ( + "cache_pool_id", + "cache_page_index_scale", + "hp_pool_id", + "hp_page_index_scale", + ): + value = getattr(spec, field_name, None) + if not isinstance(value, int) or value < 0: + raise ValueError( + f"FP4 MLA page-table spec requires non-negative {field_name}, got {value}." + ) + if spec.cache_page_index_scale <= 0 or spec.hp_page_index_scale <= 0: + raise ValueError("FP4 MLA page-index scales must be positive.") + return spec + + +# Mixed batches frequently vary by one sequence. Keep per-forward dimensions +# out of Triton's specialization key and tile page rows at one fixed width so +# those shape changes cannot trigger JIT compilation on the critical path. +@triton.jit( + do_not_specialize=["num_sequences", "num_contexts", "max_pages"], + do_not_specialize_on_alignment=["num_sequences", "num_contexts", "max_pages"], +) +def _fp4_mla_materialize_page_table_kernel( + page_ids_ptr, + paged_kv_indptr_ptr, + paged_kv_indptr_decode_ptr, + block_offsets_ptr, + kv_lens_ptr, + generation_kv_lens_ptr, + block_offsets_stride, + num_sequences, + num_contexts, + max_pages, + PAGE_SIZE: tl.constexpr, + PAGE_INDEX_SCALE: tl.constexpr, + PAGE_TILE_SIZE: tl.constexpr, +): + sequence_idx = tl.program_id(0) + page_tile_idx = tl.program_id(1) + page_offsets = page_tile_idx * PAGE_TILE_SIZE + tl.arange(0, PAGE_TILE_SIZE) + generation_idx = sequence_idx - num_contexts + is_generation = generation_idx >= 0 + context_kv_len = tl.load(kv_lens_ptr + sequence_idx) + generation_kv_len = tl.load( + generation_kv_lens_ptr + generation_idx, + mask=is_generation, + other=0, + ) + kv_len = tl.maximum(tl.where(is_generation, generation_kv_len, context_kv_len), 0) + num_active_pages = tl.minimum( + (kv_len + PAGE_SIZE - 1) // PAGE_SIZE, + max_pages, + ) + active_page_mask = page_offsets < num_active_pages + encoded_page_offsets = tl.load( + block_offsets_ptr + sequence_idx * block_offsets_stride + page_offsets, + mask=active_page_mask, + other=-1, + ) + decoded_page_ids = tl.where( + encoded_page_offsets >= 0, + encoded_page_offsets // PAGE_INDEX_SCALE, + encoded_page_offsets, + ) + page_ids = tl.where(active_page_mask, decoded_page_ids, 0) + table_offset = sequence_idx * max_pages + page_offsets + # Fixed-stride indptrs expose the whole row. Initialize inactive slots so + # masked or prefetched page-table reads cannot observe stale page IDs. + tl.store( + page_ids_ptr + table_offset, + page_ids, + mask=page_offsets < max_pages, + ) + + first_lane = page_offsets == 0 + sequence_start = sequence_idx * max_pages + tl.store( + paged_kv_indptr_ptr + sequence_idx + page_offsets, + sequence_start, + mask=first_lane, + ) + tl.store( + paged_kv_indptr_decode_ptr + generation_idx + page_offsets, + generation_idx * max_pages, + mask=first_lane & is_generation, + ) + final_sequence = sequence_idx == num_sequences - 1 + table_end = num_sequences * max_pages + num_generation_sequences = num_sequences - num_contexts + tl.store( + paged_kv_indptr_ptr + num_sequences + page_offsets, + table_end, + mask=first_lane & final_sequence, + ) + tl.store( + paged_kv_indptr_decode_ptr + num_generation_sequences + page_offsets, + num_generation_sequences * max_pages, + mask=first_lane & final_sequence, + ) + + +def configure_fp4_mla_device_page_table( + metadata: Any, + kv_lens: Optional[torch.Tensor] = None, +) -> bool: + """Configure the fixed-stride, device-materialized page table. + + Context, generation, and fresh mixed batches receive the full block-offset + table on the GPU. The materialization kernel decodes V2 page indices and + refreshes rows from the final device KV lengths before cache update. + """ + metadata.fp4_mla_state.device_page_table = False + metadata.fp4_mla_state.device_page_table_valid = False + metadata.fp4_mla_state.page_table_stride = 0 + metadata.fp4_mla_state.context_repack_max_touched_pages = 1 + + kv_cache_manager = getattr(metadata, "kv_cache_manager", None) + num_contexts = int(getattr(metadata, "num_contexts", 0)) + num_sequences = int(getattr(metadata, "num_seqs", 0)) + num_generation_sequences = num_sequences - num_contexts + num_tokens = int(getattr(metadata, "num_tokens", 0)) + num_context_tokens = int(getattr(metadata, "num_ctx_tokens", 0)) + num_generation_tokens = num_tokens - num_context_tokens + block_offsets = getattr(metadata, "kv_cache_block_offsets", None) + page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) + paged_kv_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) + paged_kv_indptr_decode = getattr(metadata.fp4_mla_state, "paged_kv_indptr_decode", None) + max_page_capacity = int(getattr(kv_cache_manager, "max_blocks_per_seq", 0) or 0) + page_spec = _fp4_mla_page_table_spec(kv_cache_manager) + page_index_scale = int(page_spec.cache_page_index_scale) + + tensors = (block_offsets, page_ids, paged_kv_indptr, paged_kv_indptr_decode) + is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) + generation_only = num_contexts == 0 + fresh_mixed = ( + not is_cuda_graph + and num_contexts > 0 + and num_generation_sequences > 0 + and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 + ) + fresh_context_only = ( + not is_cuda_graph + and num_contexts > 0 + and num_generation_sequences == 0 + and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 + ) + has_valid_generation = num_generation_sequences == 0 or ( + num_generation_tokens >= num_generation_sequences + and num_generation_tokens % num_generation_sequences == 0 + ) + # NVFP4 exposes one data pool plus its paired block-scale pool. The + # materializer reads encoded data offsets from pool 0. + supported = ( + (generation_only or fresh_mixed or fresh_context_only) + and kv_cache_manager is not None + and has_valid_generation + and int(getattr(metadata, "beam_width", 1)) == 1 + and not bool(getattr(metadata, "is_spec_dec_tree", False)) + and not bool(getattr(metadata, "locality_domain_enabled", False)) + and not bool(getattr(metadata, "enable_helix", False)) + and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK + and max_page_capacity > 0 + and page_index_scale > 0 + and all(isinstance(tensor, torch.Tensor) for tensor in tensors) + and all(tensor.dtype == torch.int32 for tensor in tensors) + and all(tensor.is_cuda for tensor in tensors) + ) + if not supported: + return False + + hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) + max_pool_id = max(page_spec.cache_pool_id, page_spec.hp_pool_id) + if ( + not isinstance(hp_page_ids, torch.Tensor) + or hp_page_ids.dtype != torch.int32 + or not hp_page_ids.is_cuda + or block_offsets.shape[0] <= max_pool_id + ): + return False + metadata.fp4_mla_state.cache_pool_id = int(page_spec.cache_pool_id) + metadata.fp4_mla_state.cache_page_index_scale = int(page_spec.cache_page_index_scale) + metadata.fp4_mla_state.hp_pool_id = int(page_spec.hp_pool_id) + metadata.fp4_mla_state.hp_page_index_scale = int(page_spec.hp_page_index_scale) + + max_pages = max_page_capacity + host_kv_lens_available = ( + isinstance(kv_lens, torch.Tensor) + and kv_lens.device.type == "cpu" + and kv_lens.ndim == 1 + and kv_lens.numel() >= num_sequences + ) + if fresh_mixed and not host_kv_lens_available: + return False + if not is_cuda_graph and host_kv_lens_available: + generation_tokens_per_sequence = ( + num_generation_tokens // num_generation_sequences if num_generation_sequences > 0 else 0 + ) + # Eager execution can narrow the fixed row stride to the current + # batch. CUDA Graph metadata retains full configured capacity so a + # replay never changes tensor addresses or launch dimensions. + max_kv_len = int(kv_lens[:num_sequences].max().item()) + max( + 0, + generation_tokens_per_sequence - 1, + ) + max_pages = min( + max_page_capacity, + max(1, _ceil_div(max_kv_len, FP4_MLA_TOKENS_PER_BLOCK)), + ) + if num_contexts > 0: + max_context_len = int(kv_lens[:num_contexts].max().item()) + max_context_pages = _ceil_div( + max_context_len, + FP4_MLA_TOKENS_PER_BLOCK, + ) + metadata.fp4_mla_state.context_repack_max_touched_pages = min( + max_pages, + triton.next_power_of_2(max(1, max_context_pages)), + ) + + assert isinstance(block_offsets, torch.Tensor) + assert isinstance(page_ids, torch.Tensor) + assert isinstance(paged_kv_indptr, torch.Tensor) + assert isinstance(paged_kv_indptr_decode, torch.Tensor) + required_page_ids = num_sequences * max_pages + buffers_cover_table = ( + block_offsets.ndim == 4 + and block_offsets.shape[0] >= 1 + and block_offsets.shape[1] >= num_sequences + and block_offsets.shape[2] >= 1 + and block_offsets.shape[3] >= max_pages + and page_ids.ndim == 1 + and page_ids.numel() >= required_page_ids + and paged_kv_indptr.ndim == 1 + and paged_kv_indptr.numel() >= num_sequences + 1 + and paged_kv_indptr_decode.ndim == 1 + and paged_kv_indptr_decode.numel() >= num_generation_sequences + 1 + ) + if not buffers_cover_table: + return False + if metadata.fp4_mla_state.hp_page_indices.numel() < required_page_ids: + return False + + metadata.fp4_mla_state.device_page_table = True + metadata.fp4_mla_state.num_sequences = num_sequences + metadata.fp4_mla_state.page_table_stride = max_pages + metadata.fp4_mla_state.num_blocks = None + metadata.fp4_mla_state.num_context_blocks = num_contexts * max_pages + metadata.fp4_mla_state.num_generation_blocks = num_generation_sequences * max_pages + return True + + +def materialize_fp4_mla_device_page_table( + metadata: Any, + kv_lens: torch.Tensor, + generation_kv_lens: Optional[torch.Tensor] = None, +) -> None: + """Refresh the fixed-stride context and generation page table once per forward.""" + if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): + raise RuntimeError("FP4 MLA requires fixed-stride device page metadata.") + if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): + return + + num_contexts = int(metadata.num_contexts) + num_sequences = int(metadata.num_seqs) + num_generation_sequences = num_sequences - num_contexts + max_pages = int(metadata.fp4_mla_state.page_table_stride) + if num_sequences <= 0 or max_pages <= 0: + raise RuntimeError( + "FP4 MLA device page metadata requires positive sequence and page capacities." + ) + if ( + kv_lens.dtype != torch.int32 + or not kv_lens.is_cuda + or kv_lens.ndim != 1 + or kv_lens.stride(0) != 1 + or kv_lens.numel() < num_sequences + ): + raise ValueError( + "FP4 MLA device page metadata requires a contiguous CUDA int32 " + f"KV-length tensor with at least {num_sequences} entries." + ) + + if generation_kv_lens is None: + generation_kv_lens = kv_lens[num_contexts:num_sequences] + if ( + generation_kv_lens.dtype != torch.int32 + or not generation_kv_lens.is_cuda + or generation_kv_lens.ndim != 1 + or generation_kv_lens.stride(0) != 1 + or generation_kv_lens.numel() < num_generation_sequences + ): + raise ValueError( + "FP4 MLA device page metadata requires a contiguous CUDA int32 " + "generation KV-length tensor with at least " + f"{num_generation_sequences} entries." + ) + + cache_pool_id = int(getattr(metadata.fp4_mla_state, "cache_pool_id", 0)) + block_offsets = metadata.kv_cache_block_offsets[ + cache_pool_id, + :num_sequences, + 0, + :max_pages, + ] + page_ids = metadata.fp4_mla_state._paged_kv_indices[: num_sequences * max_pages] + page_index_scale = int(metadata.fp4_mla_state.cache_page_index_scale) + if page_index_scale <= 0: + raise RuntimeError("FP4 MLA device page metadata requires a positive page-index scale.") + grid = ( + num_sequences, + triton.cdiv(max_pages, _FP4_MLA_PAGE_TABLE_TILE_SIZE), + ) + _fp4_mla_materialize_page_table_kernel[grid]( + page_ids, + metadata.fp4_mla_state._paged_kv_indptr, + metadata.fp4_mla_state.paged_kv_indptr_decode, + block_offsets, + kv_lens, + generation_kv_lens, + block_offsets.stride(0), + num_sequences, + num_contexts, + max_pages, + PAGE_SIZE=metadata.page_size, + PAGE_INDEX_SCALE=page_index_scale, + PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, + num_warps=4, + ) + hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) + if not isinstance(hp_page_ids, torch.Tensor): + raise RuntimeError("FP4 MLA requires an HP page-table output tensor.") + hp_pool_id = int(metadata.fp4_mla_state.hp_pool_id) + hp_page_index_scale = int(metadata.fp4_mla_state.hp_page_index_scale) + hp_block_offsets = metadata.kv_cache_block_offsets[ + hp_pool_id, + :num_sequences, + 0, + :max_pages, + ] + _fp4_mla_materialize_page_table_kernel[grid]( + hp_page_ids, + metadata.fp4_mla_state._paged_kv_indptr, + metadata.fp4_mla_state.paged_kv_indptr_decode, + hp_block_offsets, + kv_lens, + generation_kv_lens, + hp_block_offsets.stride(0), + num_sequences, + num_contexts, + max_pages, + PAGE_SIZE=metadata.page_size, + PAGE_INDEX_SCALE=hp_page_index_scale, + PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, + num_warps=4, + ) + metadata.fp4_mla_state.device_page_table_valid = True + + +def _host_int_list_during_forward(value: Any, start: int, end: int) -> Optional[list[int]]: + if torch.cuda.is_current_stream_capturing(): + return None + return _host_int_list(value, start, end) + + +def _fp4_mla_generation_num_blocks_device(metadata: Any) -> torch.Tensor: + """Device-side scalar view holding the generation page-table capacity. + + ``paged_kv_indptr_decode[num_gen]`` is the fixed-stride generation-table + endpoint. Device kernels combine it with the live KV lengths, so inactive + slots are never consumed. + """ + num_gen = metadata.num_seqs - metadata.num_contexts + return metadata.fp4_mla_state.paged_kv_indptr_decode[num_gen : num_gen + 1] + + +def _fp4_mla_uniform_generation_lengths( + metadata: Any, num_gen_tokens: int, num_gen: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Return preallocated CUDA generation KV and append lengths.""" + if num_gen <= 0 or num_gen_tokens % num_gen != 0: + raise RuntimeError("FP4 MLA generation requires a non-empty uniform request batch.") + + num_contexts = metadata.num_contexts + num_seqs = metadata.num_seqs + kv_lens_gen = metadata.kv_lens_cuda_runtime[num_contexts:num_seqs] + prompt_lens_gen = metadata.prompt_lens_cuda_runtime[num_contexts:num_seqs] + corrected_kv_lens = getattr(metadata.fp4_mla_state, "generation_kv_lens", None) + generation_lens = getattr(metadata.fp4_mla_state, "generation_append_lens", None) + tensors = ( + kv_lens_gen, + prompt_lens_gen, + corrected_kv_lens, + generation_lens, + ) + if ( + not all(isinstance(tensor, torch.Tensor) for tensor in tensors) + or corrected_kv_lens.numel() < num_gen + or generation_lens.numel() < num_gen + or not all(tensor.is_cuda for tensor in tensors) + ): + raise RuntimeError("FP4 MLA generation lengths require preallocated CUDA buffers.") + + record_for_capture = bool( + getattr(metadata, "is_cuda_graph", False) + and torch.cuda.is_current_stream_capturing() + and not getattr(metadata.fp4_mla_state, "generation_lengths_capture_recorded", False) + ) + precomputed = ( + not record_for_capture + and metadata.fp4_mla_state.generation_lengths_num_tokens == num_gen_tokens + and metadata.fp4_mla_state.generation_lengths_num_seqs == num_gen + and metadata.fp4_mla_state.generation_lengths_num_contexts == num_contexts + ) + if not precomputed: + populate_fp4_mla_generation_lengths( + kv_lens_gen, + prompt_lens_gen, + corrected_kv_lens[:num_gen], + generation_lens[:num_gen], + num_gen_tokens=num_gen_tokens, + num_gen=num_gen, + ) + metadata.fp4_mla_state.generation_lengths_num_tokens = num_gen_tokens + metadata.fp4_mla_state.generation_lengths_num_seqs = num_gen + metadata.fp4_mla_state.generation_lengths_num_contexts = num_contexts + if record_for_capture: + metadata.fp4_mla_state.generation_lengths_capture_recorded = True + return corrected_kv_lens[:num_gen], generation_lens[:num_gen] + + +def _materialize_fp4_mla_device_page_table_for_forward( + metadata: Any, + generation_kv_lens: Optional[torch.Tensor] = None, +) -> None: + """Materialize all fixed-stride rows from final per-forward device lengths.""" + if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): + raise RuntimeError("FP4 MLA cache update requires fixed-stride device page metadata.") + if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): + return + + num_contexts = int(metadata.num_contexts) + num_sequences = int(metadata.num_seqs) + num_generation_sequences = num_sequences - num_contexts + if generation_kv_lens is None: + if num_generation_sequences > 0: + num_generation_tokens = int(metadata.num_tokens) - int(metadata.num_ctx_tokens) + generation_kv_lens, _ = _fp4_mla_uniform_generation_lengths( + metadata, + num_generation_tokens, + num_generation_sequences, + ) + else: + generation_kv_lens = metadata.kv_lens_cuda_runtime[num_contexts:num_sequences] + materialize_fp4_mla_device_page_table( + metadata, + metadata.kv_lens_cuda_runtime[:num_sequences], + generation_kv_lens, + ) + + +def _max_generation_pages(metadata: Any) -> int: + num_gen = metadata.num_seqs - metadata.num_contexts + if num_gen <= 0: + return 0 + if not getattr(metadata.fp4_mla_state, "device_page_table", False): + raise RuntimeError("FP4 MLA generation requires fixed-stride device page metadata.") + max_pages = int(metadata.fp4_mla_state.page_table_stride) + if max_pages <= 0: + raise RuntimeError("FP4 MLA device page-table stride must be positive.") + return max_pages + + +def _fp4_mla_generation_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: + """Return the fixed-stride generation page-table view.""" + expected_num_gen = metadata.num_seqs - metadata.num_contexts + if num_gen_seqs != expected_num_gen: + raise RuntimeError( + "FP4 MLA generation sequence count does not match metadata: " + f"{num_gen_seqs} != {expected_num_gen}." + ) + max_pages = _max_generation_pages(metadata) + page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) + start = metadata.num_contexts * max_pages + end = start + num_gen_seqs * max_pages + if ( + not isinstance(page_ids, torch.Tensor) + or page_ids.ndim != 1 + or page_ids.dtype != torch.int32 + or not page_ids.is_contiguous() + or page_ids.numel() < end + ): + raise RuntimeError("FP4 MLA fixed-stride generation page-table backing is invalid.") + return page_ids[start:end] + + +def _fp4_mla_generation_hp_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: + """Return generation rows from the fixed-stride V2 HP page table.""" + expected_num_gen = metadata.num_seqs - metadata.num_contexts + if num_gen_seqs != expected_num_gen: + raise RuntimeError( + "FP4 MLA generation sequence count does not match metadata: " + f"{num_gen_seqs} != {expected_num_gen}." + ) + max_pages = _max_generation_pages(metadata) + page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) + start = metadata.num_contexts * max_pages + end = start + num_gen_seqs * max_pages + if ( + not isinstance(page_ids, torch.Tensor) + or page_ids.ndim != 1 + or page_ids.dtype != torch.int32 + or not page_ids.is_contiguous() + or page_ids.numel() < end + ): + raise RuntimeError("FP4 MLA fixed-stride generation HP page-table backing is invalid.") + return page_ids[start:end] + + +def _host_int_list(value: Any, start: int, end: int) -> Optional[list[int]]: + if value is None: + return None + if isinstance(value, torch.Tensor): + if value.is_cuda: + return None + return [int(item) for item in value[start:end].tolist()] + try: + return [int(item) for item in value[start:end]] + except (TypeError, ValueError): + return None + + +def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> bool: + if getattr(metadata, "is_cuda_graph", False): + return False + + start = metadata.num_contexts + end = metadata.num_seqs + block_counts = _host_int_list(getattr(metadata.fp4_mla_state, "num_blocks", None), start, end) + if block_counts is not None and ( + not block_counts or min(block_counts) != max_pages or max(block_counts) != max_pages + ): + return False + + kv_lens_cuda = getattr(metadata, "kv_lens_cuda_runtime", None) + if isinstance(kv_lens_cuda, torch.Tensor): + cache_key = ( + start, + end, + max_pages, + page_size, + tuple(block_counts) if block_counts is not None else None, + kv_lens_cuda.data_ptr(), + ) + cache = getattr(metadata.fp4_mla_state, "full_pages_cache", None) + if cache is not None and cache[0] == cache_key: + return bool(cache[1]) + kv_lens = [int(item) for item in kv_lens_cuda[start:end].detach().cpu().tolist()] + result = bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size + setattr(metadata.fp4_mla_state, "full_pages_cache", (cache_key, result)) + return result + + kv_cache_params = getattr(metadata, "kv_cache_params", None) + cached_token_lens = _host_int_list( + getattr(kv_cache_params, "num_cached_tokens_per_seq", None), + start, + end, + ) + seq_lens_kv = _host_int_list(getattr(metadata, "seq_lens_kv", None), start, end) + if cached_token_lens is not None and seq_lens_kv is not None: + if len(cached_token_lens) != len(seq_lens_kv): + return False + kv_lens = [ + cached_len + seq_len for cached_len, seq_len in zip(cached_token_lens, seq_lens_kv) + ] + elif kv_cache_params is None: + kv_lens = _host_int_list(getattr(metadata, "prompt_lens_cpu_runtime", None), start, end) + else: + return False + + return bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size + + +def _get_linear_mtp_query_len_per_seq( + metadata: Any, + *, + num_queries: int, + num_gen_seqs: int, +) -> int: + """Return the uniform generation query length required by linear MTP. + + Derives the length from the real query-token count (``num_queries``, taken + from the q shape) and the generation sequence count, which are reliable in + every representation. The host ``prompt_lens``/``seq_lens`` mirror can lag at + the decode anchor (== 1) under CUDA graph / one-engine MTP, so it is only + consulted to produce a precise diagnostic when the counts do not divide + evenly (a genuinely non-uniform batch, which the no-dequant path does not + support). + """ + if num_gen_seqs <= 0: + return 1 + + if num_queries % num_gen_seqs == 0: + return num_queries // num_gen_seqs + + start = metadata.num_contexts + end = metadata.num_seqs + query_lens = _host_int_list_during_forward( + getattr(metadata, "prompt_lens_cpu_runtime", None), start, end + ) + if query_lens is None: + query_lens = _host_int_list_during_forward(getattr(metadata, "seq_lens", None), start, end) + raise NotImplementedError( + "FP4 MLA no-dequant attention requires a uniform linear MTP generation " + f"query length; got {num_queries} query tokens for {num_gen_seqs} " + f"sequences (per-sequence lengths {query_lens})." + ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py index ea5465a1b2a5..db4aefb823b9 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py @@ -9,13 +9,8 @@ import torch -from . import ( - FP4_MLA_KV_GLOBAL_SCALE, - FP4_MLA_Q_GLOBAL_SCALE, - HP_BLOCK_SIZE, - configure_fp4_mla_device_page_table, - populate_fp4_mla_append_metadata, -) +from .config import FP4_MLA_KV_GLOBAL_SCALE, FP4_MLA_Q_GLOBAL_SCALE, HP_BLOCK_SIZE +from .metadata import configure_fp4_mla_device_page_table, populate_fp4_mla_append_metadata if TYPE_CHECKING: from ....memory_buffer_utils import Buffers @@ -143,6 +138,8 @@ def invalidate_generation_lengths(self) -> None: self.generation_lengths_capture_recorded = False def prepare(self, metadata: TrtllmAttentionMetadata, kv_lens: torch.Tensor) -> None: + # All local layers reuse this FP8 view only within the current batch. + self.fp8_context_state = None self.invalidate_generation_lengths() if metadata.kv_cache_manager is None or metadata.request_ids is None: raise RuntimeError( diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py new file mode 100644 index 000000000000..d7408491b862 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py @@ -0,0 +1,409 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP4 MLA packed-V views, repacking, and cache validity tracking.""" + +import os +from typing import Any, Optional + +import torch +import triton +import triton.language as tl + +from .config import ( + FP4_MLA_TOKENS_PER_BLOCK, + _ceil_div, + _env_enabled_default, + _env_int, + _fp4_mla_attention_backend, +) +from .layout import _ensure_workspace_tensor + + +def _shared_v_pack_storage_enabled() -> bool: + return os.getenv("TRTLLM_FP4_MLA_SHARE_V_PACK_STORAGE", "1").lower() not in ( + "0", + "false", + "no", + "off", + ) + + +def _select_triton_block_v(num_queries: int, *, prefer_prepacked_v: bool = False) -> int: + env_block_v = _env_int("TRTLLM_FP4_MLA_BLOCK_V") + if env_block_v is not None: + return env_block_v + if prefer_prepacked_v: + return 128 + return 32 if num_queries <= 32 else 128 + + +def _v_packed_shape( + kv_cache: torch.Tensor, + v_head_dim: int, + page_size: int, + block_v: int, +) -> tuple[int, int]: + return (kv_cache.shape[0] * _ceil_div(v_head_dim, block_v) * block_v, page_size // 2) + + +def _get_fp4_mla_v_packed_pool(metadata: Any, local_layer: int) -> Optional[torch.Tensor]: + return metadata.kv_cache_manager.get_mla_v_packed_pool(local_layer) + + +def _get_fp4_mla_v_packed_pool_base(metadata: Any) -> Optional[torch.Tensor]: + return metadata.kv_cache_manager.get_mla_v_packed_pool_base() + + +def _get_fp4_mla_v_scale_pool_base(metadata: Any) -> Optional[torch.Tensor]: + return metadata.kv_cache_manager.get_mla_v_scale_pool_base() + + +def _get_cutedsl_persistent_v_packed_cache( + metadata: Any, + local_layer: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + block_v: int, +) -> torch.Tensor: + v_packed = _get_fp4_mla_v_packed_pool(metadata, local_layer) + if v_packed is None: + raise RuntimeError( + "CuTeDSL FP4 MLA requires the manager-owned persistent V-packed " + "pool; the scratch full-repack fallback has been removed." + ) + expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) + if ( + not isinstance(v_packed, torch.Tensor) + or v_packed.dtype != torch.uint8 + or v_packed.device != kv_cache.device + or tuple(v_packed.shape) != expected_shape + or not v_packed.is_contiguous() + ): + raise RuntimeError( + "FP4 MLA persistent V-packed pool must be a contiguous uint8 tensor " + f"with shape {expected_shape} on {kv_cache.device}; got " + f"{type(v_packed).__name__}, " + f"shape={getattr(v_packed, 'shape', None)}, " + f"dtype={getattr(v_packed, 'dtype', None)}, " + f"device={getattr(v_packed, 'device', None)}." + ) + return v_packed + + +def _repack_cutedsl_v_packed_cache( + v_packed: torch.Tensor, + kv_cache: torch.Tensor, + page_ids: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + block_v: int, + page_indptr: Optional[torch.Tensor] = None, + kv_lens: Optional[torch.Tensor] = None, + generation_lens: Optional[torch.Tensor] = None, + max_touched_pages: int = 1, +) -> None: + if page_ids.numel() == 0: + return + from .fp4_mla_cutedsl_v_repack import fp4_mla_repack_v_cache + + fp4_mla_repack_v_cache( + v_packed, + kv_cache, + page_ids, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + page_indptr=page_indptr, + kv_lens=kv_lens, + generation_lens=generation_lens, + max_touched_pages=max_touched_pages, + ) + + +def _v_packed_cache_tag( + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> tuple[Any, ...]: + v_sf_tag = ( + None + if v_sf is None + else ( + int(v_sf.data_ptr()), + str(v_sf.device), + str(v_sf.dtype), + tuple(int(dim) for dim in v_sf.shape), + tuple(int(stride) for stride in v_sf.stride()), + ) + ) + page_ids_tag = ( + None + if page_ids is None + else ( + int(page_ids.data_ptr()), + str(page_ids.device), + str(page_ids.dtype), + tuple(int(dim) for dim in page_ids.shape), + tuple(int(stride) for stride in page_ids.stride()), + ) + ) + return ( + int(layer_idx), + None if local_layer is None else int(local_layer), + int(kv_cache.data_ptr()), + str(kv_cache.device), + str(kv_cache.dtype), + tuple(int(dim) for dim in kv_cache.shape), + tuple(int(stride) for stride in kv_cache.stride()), + int(v_head_dim), + int(page_size), + int(block_v), + v_sf_tag, + page_ids_tag, + ) + + +def _triton_prepack_v_enabled() -> bool: + if _fp4_mla_attention_backend() != "triton": + return False + default = _env_enabled_default("TRTLLM_FP4_MLA_PREPACK_V", True) + return _env_enabled_default("TRTLLM_FP4_MLA_TRITON_PREPACK_V", default) + + +def _triton_can_prepack_v(v_head_dim: int, page_size: int, block_v: int) -> bool: + return ( + _triton_prepack_v_enabled() + and hasattr(tl, "make_tensor_descriptor") + and block_v in (32, 128) + and v_head_dim % block_v == 0 + and page_size == FP4_MLA_TOKENS_PER_BLOCK + ) + + +def _triton_v_packed_attr(layer_idx: int) -> str: + if _shared_v_pack_storage_enabled(): + return "_fp4_mla_triton_attention_v_packed_buf" + return f"_fp4_mla_triton_attention_v_packed_buf_l{layer_idx}" + + +def _triton_v_packed_valid_attr(layer_idx: int) -> str: + return f"_fp4_mla_triton_attention_v_packed_valid_l{layer_idx}" + + +def _triton_shared_v_packed_valid_attr() -> str: + return "_fp4_mla_triton_attention_v_packed_valid_tag" + + +def _triton_v_packed_cache_tag( + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> tuple[Any, ...]: + return ( + "triton", + _v_packed_cache_tag( + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ), + ) + + +def _set_triton_v_packed_cache_valid( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> None: + valid_attr = ( + _triton_shared_v_packed_valid_attr() + if _shared_v_pack_storage_enabled() + else _triton_v_packed_valid_attr(layer_idx) + ) + metadata.fp4_mla_state.v_packed_cache_tags[valid_attr] = _triton_v_packed_cache_tag( + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ) + + +def _is_triton_v_packed_cache_valid( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> bool: + valid_attr = ( + _triton_shared_v_packed_valid_attr() + if _shared_v_pack_storage_enabled() + else _triton_v_packed_valid_attr(layer_idx) + ) + return metadata.fp4_mla_state.v_packed_cache_tags.get(valid_attr) == _triton_v_packed_cache_tag( + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ) + + +def _get_triton_v_packed_cache( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> Optional[torch.Tensor]: + if not _triton_can_prepack_v(v_head_dim, page_size, block_v): + return None + if not _is_triton_v_packed_cache_valid( + metadata, + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ): + return None + v_packed = metadata.fp4_mla_state.workspaces.get(_triton_v_packed_attr(layer_idx)) + expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) + if ( + v_packed is None + or v_packed.dtype != torch.uint8 + or v_packed.device != kv_cache.device + or len(v_packed.shape) != 2 + or v_packed.shape[0] < expected_shape[0] + or v_packed.shape[1] < expected_shape[1] + ): + return None + return v_packed[: expected_shape[0], : expected_shape[1]] + + +def _update_triton_v_packed_cache( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + page_ids: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + block_v: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + num_valid_pages: Optional[torch.Tensor] = None, +) -> Optional[torch.Tensor]: + if not _triton_can_prepack_v(v_head_dim, page_size, block_v): + return None + if page_ids.numel() == 0: + return None + from .fp4_mla_triton import fp4_mla_repack_v_cache_triton + + def _tma_alloc(size: int, alignment: int, stream): + return torch.empty(size, device=kv_cache.device, dtype=torch.int8) + + triton.set_allocator(_tma_alloc) + attr_name = _triton_v_packed_attr(layer_idx) + v_packed = _ensure_workspace_tensor( + metadata, + attr_name, + _v_packed_shape(kv_cache, v_head_dim, page_size, block_v), + dtype=torch.uint8, + device=kv_cache.device, + ) + fp4_mla_repack_v_cache_triton( + v_packed, + kv_cache, + page_ids, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + num_valid_pages=num_valid_pages, + ) + _set_triton_v_packed_cache_valid( + metadata, + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ) + return v_packed + + +def _maybe_update_triton_v_packed_cache( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + page_ids: torch.Tensor, + *, + num_queries: int, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + num_valid_pages: Optional[torch.Tensor] = None, +) -> None: + block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) + _update_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + page_ids, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + num_valid_pages=num_valid_pages, + ) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 7a6a83dc0a7c..59ca359e3516 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -872,10 +872,6 @@ def restore_after_draft_forward(self, saved_state: dict | None) -> None: return None def prepare(self) -> None: - # The FP8 scratch metadata view is shared by every local FP4 MLA layer - # in one eager context forward and must be rebuilt for the next batch. - if self.fp4_mla_state is not None: - self.fp4_mla_state.fp8_context_state = None super().prepare() # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers. self._invalidate_mla_scheduler_buffers() @@ -2740,7 +2736,6 @@ def _fp4_mla_rope_generation( "FP4 MLA generation requires fused BF16 RoPE/cache update " "with an FP32 rotary table.") - metadata.fp4_mla_state.generation_cache_scattered = False hp_pool_updated = scatter_fp4_mla_kv_cache( metadata, latent_cache, @@ -2757,7 +2752,6 @@ def _fp4_mla_rope_generation( if not hp_pool_updated: raise RuntimeError( "Fused FP4 MLA RoPE/cache scatter did not update the HP pool.") - metadata.fp4_mla_state.generation_cache_scattered = True def can_fuse_fp4_mla_q_quant( self, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index c63cd14d93ba..fe97e51ae823 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1580,12 +1580,6 @@ def get_cache_size_per_token(model_config: ModelConfigPython, # get head dim mla = hasattr(config, "kv_lora_rank") and config.kv_lora_rank is not None - quant_config = model_config.quant_config - if (mla and quant_config is not None - and quant_config.quant_mode.has_fp4_kv_cache()): - raise ValueError( - "FP4 MLA cache sizing requires Fp4MlaKVCacheManagerV2; " - "KVCacheManager V1 is not supported.") if mla: head_dim = config.kv_lora_rank + config.qk_rope_head_dim kv_factor = 1 @@ -1602,6 +1596,7 @@ def get_cache_size_per_token(model_config: ModelConfigPython, # K and V mem_per_token = kv_factor * num_attention_layers * head_dim # The data type bytes. + quant_config = model_config.quant_config if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( ): mem_per_token *= 1 diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index c6df26fd383e..fe6dde0fdf67 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -10,6 +10,8 @@ import tensorrt_llm import tensorrt_llm._torch.attention.backends.fp4_mla as fp4_mla_backend +import tensorrt_llm._torch.attention.backends.fp4_mla.cache_update as fp4_mla_cache_update +import tensorrt_llm._torch.attention.backends.fp4_mla.metadata as fp4_mla_metadata from tensorrt_llm._torch.attention.backends.fp4_mla import ( FP4_BLOCK_SIZE, FP4_MLA_ATTENTION_BACKEND_ENV, @@ -114,7 +116,7 @@ def populate_generation_lengths(*args, **kwargs) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) monkeypatch.setattr( - fp4_mla_backend, + fp4_mla_metadata, "populate_fp4_mla_generation_lengths", populate_generation_lengths, ) @@ -992,7 +994,7 @@ def record_repack(*args, **kwargs) -> None: } ) - monkeypatch.setattr(fp4_mla_backend, "_repack_cutedsl_v_packed_cache", record_repack) + monkeypatch.setattr(fp4_mla_cache_update, "_repack_cutedsl_v_packed_cache", record_repack) _assert_fp4_mla_attention_decode_accuracy( monkeypatch, From 711a199191739d097246c171308e1b59cdeea65c Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:30:48 -0700 Subject: [PATCH 17/21] test: preserve cache hits in FMHA sanity check coverage Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tests/unittest/_torch/attention/test_fmha_manager.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 7cfb59eee093..16f5dc64a81e 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -739,6 +739,7 @@ def test_fmha_cache_sanity_check_logs_mismatched_inputs() -> None: relative_attention_max_distance=7, ) + # Keep K/V presence unchanged so both inputs hit the same cache key. with ( patch.object(fmha_manager, "_is_fmha_cache_enabled", return_value=True), patch.object(fmha_manager.logger, "error") as log_error, @@ -748,7 +749,7 @@ def test_fmha_cache_sanity_check_logs_mismatched_inputs() -> None: attn, torch.empty((1, 4)), torch.empty((1, 2)), - None, + torch.empty((1, 2)), cached_metadata, cached_args, ) @@ -769,7 +770,7 @@ def test_fmha_cache_sanity_check_logs_mismatched_inputs() -> None: assert "cached=FakeFmha, uncached=FakeFmha" in message assert "q: cached=shape=(1, 4), uncached=shape=(1, 8)" in message assert "k: cached=shape=(1, 2), uncached=shape=(1, 3)" in message - assert "v: cached=None, uncached=shape=(1, 3)" in message + assert "v: cached=shape=(1, 2), uncached=shape=(1, 3)" in message assert "metadata.beam_width: cached=1, uncached=2" in message assert "forward_args.relative_attention_max_distance: cached=0, uncached=7" in message From 7241ddff6c9514ce91299549a027d26225f44e6c Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 18 Sep 2026 09:27:44 -0700 Subject: [PATCH 18/21] refactor: simplify FP4 MLA cache policy and RoPE sizing Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 27 ++--------------- .../_torch/attention/backends/fmha/fp4_mla.py | 4 +-- .../_torch/attention/backends/fmha/manager.py | 30 ++----------------- .../_torch/attention/backends/trtllm.py | 8 ++--- .../_torch/attention/test_fmha_manager.py | 5 ++-- 5 files changed, 14 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index d283b3511eb1..a89c39e5a73f 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -428,30 +428,9 @@ The FMHA package is split by role: MLA uses `query_input` with `is_fused_qkv=False`. - `fmha/combined.py` composes different context and generation implementations for non-MLA mixed batches. -- `fmha/fp4_mla.py` implements FP4 MLA context and no-dequant decode. - The shared FMHA availability check admits FP4 MLA only to libraries that - declare `supports_fp4_mla`; it does not restrict the existing FP4 GQA path. - Selection caches validation in `_is_supported()`. Its cache key distinguishes - K/V input presence, sparse/sinks inputs, and prepared FP4 state, so changing - those conditions triggers validation again instead of reusing a valid entry. - The core implementation requires dense TRTLLM MLA, BF16 absorption weights, - fused RoPE with duplicated rotary tables, and KV Cache Manager V2. It uses - FP8 context attention with an FP4 cache update and FP4 generation attention. - Chunked prefill and context parallelism are rejected before KV allocation. - Cached-context attention is not implemented, so executor block reuse remains - disabled even though the manager supports full-block reuse of its pools. - Disaggregated serving is reserved for the follow-up integration. On SM107, - this dense TRTLLM path keeps NVFP4 KV quantization; unsupported profiles retain - the existing FP8 fallback. - `fp4_mla/state.py` owns batch-shared page tables, HP/V-scale views, append - metadata, and scratch caches. `TrtllmAttentionMetadata` keeps one optional - state reference and forwards prepare/MTP lifecycle updates to it. These - buffers continue to use the metadata allocator for CUDA-graph address - stability; the layer-local FP8 attention view is allocated lazily by FMHA. - The `fp4_mla` package exports entry points from focused `config`, `layout`, - `metadata`, `cache_update`, `v_cache`, and `decode` modules. The FP8 context - view has its own FMHA manager and uses normal capability-based selection; - its manager-owned scratch is allocated once and reused across layers. +- `fmha/fp4_mla.py` implements FP4 MLA using FP8 context attention with FP4 + cache updates and FP4 no-dequant decode. It uses KV Cache Manager V2; + batch state, cache storage, and kernels live in `fp4_mla/`. - `fmha/triton_custom_mask.py` implements the Triton custom-mask context phase. Custom-mask data applies to context requests; for mixed batches, `TrtllmAttention` can pair it with a later causal-generation provider through diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index 3ce18d0792d7..6fab3a59a10e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -70,8 +70,8 @@ def _is_supported( *, phase: Optional[FmhaPhase] = None, ) -> bool: - # Input presence/readiness and mask/output formats are in the selection - # cache key; cache geometry and sparse compression are model invariants. + # Mask/output formats and the attention phase are represented in the cache key. + # K/V presence is fixed per phase; optional features and pools are instance invariants. del q, phase if forward_args.output_sf is not None: raise NotImplementedError("FP4 MLA does not support quantized attention output.") diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index 0f2a014c53de..c9afff66e1d8 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -135,12 +135,6 @@ class _FmhaCacheKey(NamedTuple): attention_mask_type: AttentionMaskType use_spec_decoding: bool has_block_sparse_inputs: bool - # These support conditions can change without changing the query grid. - has_k_input: bool - has_v_input: bool - has_attention_sinks: bool - has_sparse_indices: bool - fp4_mla_state_ready: bool # LoRA can change the effective output from packed NVFP4 to unpacked BF16 # without changing the request shape. Keep those selection regimes apart. output_dtype: torch.dtype | None @@ -323,9 +317,6 @@ def _make_cache_key( q: torch.Tensor, metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, - *, - k: torch.Tensor | None = None, - v: torch.Tensor | None = None, ) -> _FmhaCacheKey: """Build the dynamic FMHA cache key for one attention instance. @@ -376,29 +367,14 @@ def _make_cache_key( generation_seq_len_q, _FMHA_CACHE_SEQ_LEN_Q_GRID ) - sparse = forward_args.sparse_runtime_params - has_sparse_indices = ( - metadata.num_sparse_topk > 0 - or (sparse.sparse_kv_indices is not None and sparse.sparse_kv_indices.numel() > 0) - or (sparse.sparse_attn_indices is not None and sparse.sparse_attn_indices.numel() > 0) - ) - fp4_state = getattr(metadata, "fp4_mla_state", None) + block_sparse_inputs = forward_args.sparse_runtime_params.block_sparse_inputs return _FmhaCacheKey( context_batch_size=context_batch_size, generation_batch_size=generation_batch_size, generation_seq_len_q=generation_seq_len_q, attention_mask_type=attention_mask_type, use_spec_decoding=metadata.use_spec_decoding, - has_block_sparse_inputs=sparse.block_sparse_inputs is not None, - has_k_input=k is not None, - has_v_input=v is not None, - has_attention_sinks=forward_args.attention_sinks is not None, - has_sparse_indices=has_sparse_indices, - fp4_mla_state_ready=( - fp4_state is not None - and fp4_state.hp_pool is not None - and fp4_state.v_scale_pool is not None - ), + has_block_sparse_inputs=block_sparse_inputs is not None, output_dtype=output_dtype, output_sf_dtype=output_sf_dtype, ) @@ -416,7 +392,7 @@ def select( if not _is_fmha_cache_enabled(): return self._select_uncached(attn, q, k, v, metadata, forward_args) - cache_key = self._make_cache_key(q, metadata, forward_args, k=k, v=v) + cache_key = self._make_cache_key(q, metadata, forward_args) fmha = self._cache.get(cache_key) if fmha is not None: if self._cache_sanity_check_enabled: diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 59ca359e3516..8b413cf23af0 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1914,11 +1914,11 @@ def _ensure_rope_table_size(self, required_max_positions: int) -> None: 2 if self.rope_params.duplicate_data else 1) table_max_positions = (self.rotary_cos_sin.numel() // floats_per_position - if self.rotary_cos_sin is not None - and floats_per_position > 0 else 0) + if self.rotary_cos_sin is not None else 0) + self.rope_params.max_positions = max(self.rope_params.max_positions, + table_max_positions, + required_max_positions) if required_max_positions > table_max_positions: - self.rope_params.max_positions = max(required_max_positions, - self.rope_params.max_positions) self.rotary_inv_freq, self.rotary_cos_sin = ( self.rope_params.create_rope_const_params()) diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 16f5dc64a81e..7cfb59eee093 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -739,7 +739,6 @@ def test_fmha_cache_sanity_check_logs_mismatched_inputs() -> None: relative_attention_max_distance=7, ) - # Keep K/V presence unchanged so both inputs hit the same cache key. with ( patch.object(fmha_manager, "_is_fmha_cache_enabled", return_value=True), patch.object(fmha_manager.logger, "error") as log_error, @@ -749,7 +748,7 @@ def test_fmha_cache_sanity_check_logs_mismatched_inputs() -> None: attn, torch.empty((1, 4)), torch.empty((1, 2)), - torch.empty((1, 2)), + None, cached_metadata, cached_args, ) @@ -770,7 +769,7 @@ def test_fmha_cache_sanity_check_logs_mismatched_inputs() -> None: assert "cached=FakeFmha, uncached=FakeFmha" in message assert "q: cached=shape=(1, 4), uncached=shape=(1, 8)" in message assert "k: cached=shape=(1, 2), uncached=shape=(1, 3)" in message - assert "v: cached=shape=(1, 2), uncached=shape=(1, 3)" in message + assert "v: cached=None, uncached=shape=(1, 3)" in message assert "metadata.beam_width: cached=1, uncached=2" in message assert "forward_args.relative_attention_max_distance: cached=0, uncached=7" in message From 70f6f2e694d901ab94af62259e35766070001bc3 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:02:50 -0700 Subject: [PATCH 19/21] refactor: integrate upstream FP4 MLA cache policy and helpers Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../_torch/attention/backends/fmha/fp4_mla.py | 161 ++++- .../attention/backends/fp4_mla/__init__.py | 4 + .../backends/fp4_mla/cache_gather.py | 154 +++++ .../backends/fp4_mla/cache_manager.py | 507 ++++++++++++++-- .../backends/fp4_mla/cache_update.py | 69 ++- .../backends/fp4_mla/fp4_mla_context.py | 48 +- .../backends/fp4_mla/fp4_mla_kernels.py | 284 ++++++++- .../backends/fp4_mla/fp4_mla_triton.py | 553 +++++++++++++++++- .../attention/backends/fp4_mla/metadata.py | 21 +- .../attention/backends/fp4_mla/state.py | 7 +- .../attention/backends/fp4_mla/v_cache.py | 256 +++++++- .../_torch/attention/backends/trtllm.py | 6 +- .../kv_cache/kv_cache_manager_v2.py | 29 +- .../_torch/pyexecutor/model_loader.py | 3 +- .../unittest/_torch/attention/test_fp4_mla.py | 518 ++++++++++++++++ 15 files changed, 2484 insertions(+), 136 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index 6fab3a59a10e..9e65fa7c522a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -25,6 +25,7 @@ from tensorrt_llm._torch.attention.backends.fp4_mla.fp4_mla_context import ( _FP8_CONTEXT_SCRATCH_ATTR, _build_fp8_mla_context_attn, + _build_fp8_mla_context_metadata, _execute_fp8_context_with_cache_update, _Fp8MlaContextScratch, _get_fp8_mla_context_metadata, @@ -125,14 +126,47 @@ def _is_supported( raise NotImplementedError("FP4 MLA requires a context-only or generation-only call.") return True + def _get_fp8_context_resources( + self, + metadata: "TrtllmAttentionMetadata", + q: torch.Tensor, + ) -> tuple["TrtllmAttention", _Fp8MlaContextScratch]: + """Reuse one manager-owned scratch and one layer-local FP8 attention view.""" + attn = self.attn + kv_cache_manager = metadata.kv_cache_manager + if kv_cache_manager is None: + raise RuntimeError("FP8 MLA context scratch requires a KV cache manager.") + scratch = getattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, None) + scratch_head_dim = (attn.kv_lora_rank or 0) + (attn.qk_rope_head_dim or 0) + if scratch is None: + scratch = _Fp8MlaContextScratch.create( + metadata, + device=q.device, + head_dim=scratch_head_dim, + ) + setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch) + else: + assert isinstance(scratch, _Fp8MlaContextScratch) + assert scratch.matches(metadata, device=q.device, head_dim=scratch_head_dim), ( + "FP8 MLA context scratch geometry must be shared by all layers of its KV manager." + ) + + fp8_attention = self._fp8_attention + if fp8_attention is None: + fp8_attention = _build_fp8_mla_context_attn(attn) + self._fp8_attention = fp8_attention + fp8_attention.rotary_inv_freq = attn.rotary_inv_freq + fp8_attention.rotary_cos_sin = attn.rotary_cos_sin + return fp8_attention, scratch + def run_mla_context(self, params: FmhaParams) -> None: attn = params.attn metadata = params.meta forward_args = params.fwd - q = params.qkv_input + q = params.query_input k = params.key_input v = params.value_input - output = params.context_buf + output = params.output if q is None or k is None or v is None: raise RuntimeError("FP4 MLA context requires expanded Q, K, and V tensors.") if output is None: @@ -154,9 +188,8 @@ def run_mla_context(self, params: FmhaParams) -> None: num_tokens = q.shape[0] output = output.view(num_tokens, -1) - local_layer = attn.get_local_layer_idx(metadata) + local_layer = attn.get_fp4_mla_local_layer_idx(metadata) kv_lora_rank = attn.kv_lora_rank or 0 - qk_rope_head_dim = attn.qk_rope_head_dim or 0 latent_cache = forward_args.latent_cache[:num_tokens] @@ -174,30 +207,7 @@ def update_fp4_cache() -> None: if not hp_pool_updated: raise RuntimeError("Fused FP4 MLA context scatter did not update the HP pool.") - kv_cache_manager = metadata.kv_cache_manager - if kv_cache_manager is None: - raise RuntimeError("FP8 MLA context scratch requires a KV cache manager.") - scratch = getattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, None) - scratch_head_dim = kv_lora_rank + qk_rope_head_dim - if scratch is None: - scratch = _Fp8MlaContextScratch.create( - metadata, - device=q.device, - head_dim=scratch_head_dim, - ) - setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch) - else: - assert isinstance(scratch, _Fp8MlaContextScratch) - assert scratch.matches(metadata, device=q.device, head_dim=scratch_head_dim), ( - "FP8 MLA context scratch geometry must be shared by all layers of its KV manager." - ) - - fp8_attention = self._fp8_attention - if fp8_attention is None: - fp8_attention = _build_fp8_mla_context_attn(attn) - self._fp8_attention = fp8_attention - fp8_attention.rotary_inv_freq = attn.rotary_inv_freq - fp8_attention.rotary_cos_sin = attn.rotary_cos_sin + fp8_attention, scratch = self._get_fp8_context_resources(metadata, q) fp8_metadata = _get_fp8_mla_context_metadata(metadata, scratch) fp8_forward_args = replace( forward_args, @@ -218,12 +228,101 @@ def run_fp8_context() -> None: scratch.cache_done_event, ) + def forward_context_partition( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + kv_lens_cuda: torch.Tensor, + kv_lens_cpu: torch.Tensor, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run one explicit-KV partition through the FP8 MLA context kernel. + + Returns whatever ``TrtllmAttention.forward`` returns for the FP8 + context copy: the attention output plus its optional output scale. + """ + if forward_args.output is None: + raise RuntimeError("FP4 MLA context partition requires an output buffer.") + if forward_args.latent_cache is not None: + raise RuntimeError("FP4 MLA context partition expects explicit K/V tensors.") + if forward_args.output_sf is not None: + raise NotImplementedError( + "FP4 MLA context partition does not support quantized attention output." + ) + if forward_args.attention_mask not in ( + PredefinedAttentionMask.CAUSAL, + PredefinedAttentionMask.FULL, + ): + raise NotImplementedError( + "FP4 MLA context partition requires a causal or full attention mask." + ) + if forward_args.attention_mask_data is not None: + raise NotImplementedError( + "FP4 MLA context partition does not support custom attention masks." + ) + if forward_args.attention_sinks is not None: + raise NotImplementedError("FP4 MLA context partition does not support attention sinks.") + if metadata.is_cuda_graph: + raise NotImplementedError("FP4 MLA chunked prefill does not support CUDA graphs.") + if metadata.num_contexts <= 0 or q.shape[0] != metadata.num_ctx_tokens: + raise RuntimeError( + "FP4 MLA context partition query token count must match context metadata." + ) + if k.shape[0] != v.shape[0]: + raise RuntimeError("FP4 MLA context partition K/V token counts do not match.") + + sparse_runtime_params = forward_args.sparse_runtime_params + if ( + ( + sparse_runtime_params.sparse_kv_indices is not None + and sparse_runtime_params.sparse_kv_indices.numel() > 0 + ) + or ( + sparse_runtime_params.sparse_attn_indices is not None + and sparse_runtime_params.sparse_attn_indices.numel() > 0 + ) + or metadata.num_sparse_topk > 0 + ): + raise NotImplementedError("FP4 MLA chunked prefill does not support sparse attention.") + + require_fp4_mla_fp8_context_support() + fp8_attention, scratch = self._get_fp8_context_resources(metadata, q) + + expected_kv_tokens = int(kv_lens_cpu[: metadata.num_contexts].sum().item()) + if k.shape[0] != expected_kv_tokens: + raise RuntimeError( + "FP4 MLA context partition K/V token count does not match " + f"the KV lengths: got {k.shape[0]}, expected {expected_kv_tokens}." + ) + scratch.prepare( + metadata, + context_lengths_cuda=kv_lens_cuda, + context_lengths_cpu=kv_lens_cpu, + ) + fp8_metadata = _build_fp8_mla_context_metadata( + metadata, + scratch, + kv_lens_cuda=kv_lens_cuda[: metadata.num_contexts], + kv_lens_cpu=kv_lens_cpu[: metadata.num_contexts], + ) + fp8_forward_args = replace( + forward_args, + output_sf=None, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + latent_cache=None, + ) + return fp8_attention.forward(q, k, v, fp8_metadata, fp8_forward_args) + def run_mla_generation(self, params: FmhaParams) -> None: attn = params.attn metadata = params.meta forward_args = params.fwd - q = params.qkv_input - output = params.context_buf + q = params.query_input + output = params.output if q is None: raise RuntimeError("FP4 MLA generation requires a fused query input.") if output is None: @@ -233,7 +332,7 @@ def run_mla_generation(self, params: FmhaParams) -> None: if metadata.num_generations <= 0: raise RuntimeError("FP4 MLA generation requires generation requests.") - local_layer = attn.get_local_layer_idx(metadata) + local_layer = attn.get_fp4_mla_local_layer_idx(metadata) kv_lora_rank = attn.kv_lora_rank or 0 qk_rope_head_dim = attn.qk_rope_head_dim or 0 fused_head_dim = kv_lora_rank + qk_rope_head_dim diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py index ccfde0c4fb15..596b208194a7 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """FP4 MLA entry points; implementation lives in the focused sibling modules.""" +from .cache_gather import load_fp4_mla_chunked_kv_cache from .cache_update import ( _get_fp4_mla_context_start_positions as _get_fp4_mla_context_start_positions, ) @@ -140,8 +141,11 @@ from .v_cache import _update_triton_v_packed_cache as _update_triton_v_packed_cache from .v_cache import _v_packed_cache_tag as _v_packed_cache_tag from .v_cache import _v_packed_shape as _v_packed_shape +from .v_cache import rebuild_fp4_mla_disagg_imported_cache __all__ = [ + "load_fp4_mla_chunked_kv_cache", + "rebuild_fp4_mla_disagg_imported_cache", "FP4_BLOCK_SIZE", "FP4_MLA_ATTENTION_BACKEND_ENV", "FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV", diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py new file mode 100644 index 000000000000..fa8879296835 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Gather cached FP4 MLA prefix partitions for explicit-KV context attention.""" + +from typing import Any + +import torch +import triton + +from .config import FP4_BLOCK_SIZE, FP4_MLA_K_RESIDUAL_DIM, _fp4_mla_attention_backend +from .fp4_mla_kernels import _fp4_mla_chunked_cache_gather_kernel +from .layout import ( + _get_fp4_mla_global_scale, + _get_fp4_mla_kv_cache_tensors, + _validate_fp4_mla_kv_storage_shape, +) +from .metadata import _materialize_fp4_mla_device_page_table_for_forward + + +def load_fp4_mla_chunked_kv_cache( + metadata: Any, + layer_idx: int, + *, + num_ctx_cached_tokens: int, + cu_chunked_seq_len: torch.Tensor, + chunked_global_offset: torch.Tensor, + chunked_max_seq_len: int, + out_dtype: torch.dtype, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather one cached-prefix partition from dense FP4 MLA V2 storage.""" + if out_dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError(f"FP4 MLA chunk gather does not support output dtype {out_dtype}.") + if num_ctx_cached_tokens < 0: + raise ValueError("FP4 MLA chunk gather token count must be non-negative.") + if chunked_max_seq_len < 0: + raise ValueError("FP4 MLA chunk gather max sequence length must be non-negative.") + if num_ctx_cached_tokens > 0 and chunked_max_seq_len == 0: + raise ValueError("A non-empty FP4 MLA chunk gather requires a positive max length.") + if ( + kv_lora_rank <= 0 + or kv_lora_rank % FP4_BLOCK_SIZE != 0 + or qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM + ): + raise ValueError( + "FP4 MLA chunk gather requires a positive kv_lora_rank and " + f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got " + f"{kv_lora_rank} and {qk_rope_head_dim}." + ) + + num_contexts = int(metadata.num_contexts) + if num_ctx_cached_tokens > 0 and num_contexts <= 0: + raise ValueError("A non-empty FP4 MLA chunk gather requires context requests.") + tensors = (cu_chunked_seq_len, chunked_global_offset) + if any(not isinstance(tensor, torch.Tensor) or not tensor.is_cuda for tensor in tensors): + raise ValueError("FP4 MLA chunk gather metadata must use CUDA tensors.") + if ( + cu_chunked_seq_len.dtype != torch.int64 + or cu_chunked_seq_len.ndim != 1 + or cu_chunked_seq_len.numel() < num_contexts + 1 + or not cu_chunked_seq_len.is_contiguous() + ): + raise ValueError( + "FP4 MLA chunk gather requires a contiguous int64 cumulative-length tensor." + ) + if ( + chunked_global_offset.dtype != torch.int64 + or chunked_global_offset.ndim != 1 + or chunked_global_offset.numel() < num_contexts + or not chunked_global_offset.is_contiguous() + ): + raise ValueError("FP4 MLA chunk gather requires contiguous int64 global offsets.") + if cu_chunked_seq_len.device != chunked_global_offset.device: + raise ValueError("FP4 MLA chunk gather metadata tensors must share one device.") + + compressed_kv = torch.empty( + (num_ctx_cached_tokens, kv_lora_rank), + dtype=out_dtype, + device=cu_chunked_seq_len.device, + ) + k_pe = torch.empty( + (num_ctx_cached_tokens, qk_rope_head_dim), + dtype=out_dtype, + device=cu_chunked_seq_len.device, + ) + if num_ctx_cached_tokens == 0: + return compressed_kv, k_pe + + if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): + raise RuntimeError("FP4 MLA chunk gather requires fixed-stride device page metadata.") + _materialize_fp4_mla_device_page_table_for_forward(metadata) + page_table_stride = int(metadata.fp4_mla_state.page_table_stride) + page_ids = metadata.fp4_mla_state._paged_kv_indices + if ( + page_table_stride <= 0 + or not isinstance(page_ids, torch.Tensor) + or page_ids.dtype != torch.int32 + or not page_ids.is_cuda + or page_ids.device != cu_chunked_seq_len.device + or page_ids.numel() < num_contexts * page_table_stride + ): + raise RuntimeError("FP4 MLA chunk gather received invalid device page metadata.") + + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + head_dim = kv_lora_rank + qk_rope_head_dim + storage_head_dim = _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=_fp4_mla_attention_backend(), + ) + if storage_head_dim != head_dim + FP4_MLA_K_RESIDUAL_DIM: + raise RuntimeError( + "FP4 MLA chunk gather requires K residual storage for BF16 reconstruction." + ) + sf_cache = sf_cache.view(torch.float8_e4m3fn) + global_scale = _get_fp4_mla_global_scale(metadata, kv_cache.device) + token_block = 16 + grid = ( + triton.cdiv(chunked_max_seq_len, token_block), + num_contexts, + triton.cdiv(head_dim, FP4_BLOCK_SIZE), + ) + _fp4_mla_chunked_cache_gather_kernel[grid]( + compressed_kv, + k_pe, + kv_cache, + sf_cache, + page_ids, + cu_chunked_seq_len, + chunked_global_offset, + global_scale, + chunked_max_seq_len, + page_table_stride, + kv_cache.shape[0], + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + compressed_kv.stride(0), + k_pe.stride(0), + KV_LORA_RANK=kv_lora_rank, + QK_ROPE_HEAD_DIM=qk_rope_head_dim, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=storage_head_dim // FP4_BLOCK_SIZE, + TOKEN_BLOCK=token_block, + num_warps=4, + ) + return compressed_kv, k_pe + + +__all__ = ["load_fp4_mla_chunked_kv_cache"] diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index e83da387f9a1..5aa8429afe24 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -4,10 +4,13 @@ import math from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, List, Optional +from types import MethodType +from typing import TYPE_CHECKING, Iterable, List, Optional import torch +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind +from tensorrt_llm._torch.kimi_k3_cache_policy import get_kimi_k3_bf16_kv_layer_ids from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned @@ -50,22 +53,57 @@ class Fp4MlaPageTableSpec: hp_is_paged: bool = True -class Fp4MlaKVCacheManagerV2(KVCacheManagerV2): - """V2 manager for canonical FP4 MLA pages and compact HP sequence state. +class Fp4MlaV2CacheLayoutPolicy: + """FP4 MLA storage policy for a KV cache manager V2 lifecycle. K, K block scales, V scales, and optional packed V share one full-history lifecycle. Each model layer also has a virtual sliding layer whose page is a compact ``16 + max_rewind`` BF16 ring. The HP role therefore follows V2 - allocation, rewind, and prefix lifecycles without storing BF16 values for - the full sequence. + allocation, rewind, prefix, and disaggregated-transfer lifecycles without + storing BF16 values for the full sequence. + + ``Fp4MlaKVCacheManagerV2`` applies this policy directly for pure MLA + models. Hybrid linear-attention managers compose the same policy with + their SSM lifecycle and forward only attention-layout operations here. """ def __init__(self, *args, **kwargs) -> None: kv_cache_config = args[0] if args else kwargs.get("kv_cache_config") dtype = kwargs.get("dtype", DataType.HALF) kv_cache_type = args[1] if len(args) > 1 else kwargs.get("kv_cache_type") - tokens_per_block = kwargs.get("tokens_per_block") - head_dim = kwargs.get("head_dim") + kv_cache_config = self.configure_manager( + self, + kv_cache_config, + kv_cache_type, + dtype=dtype, + tokens_per_block=kwargs.get("tokens_per_block"), + head_dim=kwargs.get("head_dim"), + pretrained_config=kwargs.get("pretrained_config"), + spec_config=kwargs.get("spec_config"), + is_disagg=kwargs.get("is_disagg", False), + ) + if args: + args = (kv_cache_config, *args[1:]) + else: + kwargs["kv_cache_config"] = kv_cache_config + + super().__init__(*args, **kwargs) + self.finalize_manager(self) + + @staticmethod + def configure_manager( + manager, + kv_cache_config, + kv_cache_type, + *, + dtype: DataType, + tokens_per_block: Optional[int], + head_dim: Optional[int], + pretrained_config=None, + spec_config=None, + is_disagg: bool = False, + ): + """Validate and install FP4 MLA layout state before V2 construction.""" if dtype != DataType.NVFP4 or kv_cache_type != CacheTypeCpp.SELFKONLY: raise ValueError("Fp4MlaKVCacheManagerV2 requires NVFP4 SELFKONLY cache storage.") if kv_cache_config is None or kv_cache_config.dtype not in ("auto", "nvfp4"): @@ -83,10 +121,6 @@ def __init__(self, *args, **kwargs) -> None: config_updates["enable_partial_reuse"] = False if config_updates: kv_cache_config = kv_cache_config.model_copy(update=config_updates) - if args: - args = (kv_cache_config, *args[1:]) - else: - kwargs["kv_cache_config"] = kv_cache_config if "enable_partial_reuse" in config_updates: logger.info( "FP4 MLA KV cache manager V2 disables partial block reuse " @@ -100,53 +134,205 @@ def __init__(self, *args, **kwargs) -> None: if not isinstance(head_dim, int) or head_dim <= FP4_MLA_K_RESIDUAL_DIM: raise ValueError(f"FP4 MLA V2 requires a positive scalar MLA head_dim, got {head_dim}.") - pretrained_config = kwargs.get("pretrained_config") - self.mla_v_scale_head_dim = int( + manager.mla_v_scale_head_dim = int( getattr(pretrained_config, "kv_lora_rank", head_dim - FP4_MLA_K_RESIDUAL_DIM) ) - if not 0 < self.mla_v_scale_head_dim < head_dim: + if not 0 < manager.mla_v_scale_head_dim < head_dim: raise ValueError( "FP4 MLA V2 requires kv_lora_rank in (0, head_dim), got " - f"{self.mla_v_scale_head_dim} and {head_dim}." + f"{manager.mla_v_scale_head_dim} and {head_dim}." ) - self._fp4_mla_storage_backend = _fp4_mla_attention_backend() - if self._fp4_mla_storage_backend not in _FP4_MLA_K_RESIDUAL_BACKENDS: + manager._bf16_mla_global_layer_ids = get_kimi_k3_bf16_kv_layer_ids(pretrained_config) + if manager._bf16_mla_global_layer_ids: + logger.info( + "Kimi K3 MLA layers using BF16 KV-cache fallback: " + f"{sorted(manager._bf16_mla_global_layer_ids)}." + ) + manager._fp4_mla_storage_backend = _fp4_mla_attention_backend() + if manager._fp4_mla_storage_backend not in _FP4_MLA_K_RESIDUAL_BACKENDS: raise ValueError( "Fp4MlaKVCacheManagerV2 supports only the triton and cutedsl " - f"backends, got {self._fp4_mla_storage_backend!r}." + f"backends, got {manager._fp4_mla_storage_backend!r}." ) - self.fp4_mla_k_residual_dim = ( + manager.fp4_mla_k_residual_dim = ( FP4_MLA_K_RESIDUAL_DIM - if self._fp4_mla_storage_backend in _FP4_MLA_K_RESIDUAL_BACKENDS + if manager._fp4_mla_storage_backend in _FP4_MLA_K_RESIDUAL_BACKENDS else 0 ) - self.mla_v_head_dim = ( - self.mla_v_scale_head_dim - if self._fp4_mla_storage_backend == _FP4_MLA_CUTEDSL_BACKEND + manager.mla_v_head_dim = ( + manager.mla_v_scale_head_dim + if manager._fp4_mla_storage_backend == _FP4_MLA_CUTEDSL_BACKEND and not _fp4_mla_cutedsl_fused_v_transpose_enabled() else None ) - spec_config = kwargs.get("spec_config") max_rewind_len = int(spec_config.tokens_per_gen_step - 1) if spec_config else 0 - self._fp4_mla_hp_pool_size = HP_BLOCK_SIZE + max_rewind_len - self._fp4_mla_view_cache: dict[tuple, torch.Tensor] = {} + manager._fp4_mla_hp_pool_size = HP_BLOCK_SIZE + max_rewind_len + manager._fp4_mla_view_cache = {} + manager._attention_cache_layout_policy = Fp4MlaV2CacheLayoutPolicy + Fp4MlaV2CacheLayoutPolicy.install_manager_capabilities(manager) + return kv_cache_config - super().__init__(*args, **kwargs) + @staticmethod + def install_manager_capabilities(manager) -> None: + """Attach attention-layout hooks to a non-policy lifecycle manager.""" + if isinstance(manager, Fp4MlaV2CacheLayoutPolicy): + return + manager.fp4_mla_hp_pool_size = manager._fp4_mla_hp_pool_size + method_names = ( + "_bf16_mla_local_layer_indices", + "_fp4_mla_local_layer_indices", + "_fp4_mla_compact_layer_idx", + "_bf16_mla_bytes_per_token", + "_get_buffer_roles_for_layer", + "_storage_head_dim", + "get_buffers", + "get_kv_cache_dtype", + "get_kv_cache_num_blocks", + "_v_scale_bytes_per_page", + "_v_packed_bytes_per_page", + "_hp_bytes_per_page", + "get_layer_bytes_per_token", + "_extra_buffers_per_layer", + "_prepare_page_table_tensor", + "_validate_fp4_mla_layer_groups", + "get_fp4_mla_page_table_spec", + "_role_encoded_page_capacity", + "_role_view", + "get_fp4_mla_cache_buffers", + "_iter_fp4_mla_physical_pool_views", + "_all_layer_role_view", + "_role_pool_base_view", + "get_mla_v_scale_pool", + "get_mla_v_scale_pool_base", + "get_mla_v_scale_page_offset", + "get_mla_v_packed_pool", + "get_mla_v_packed_pool_base", + "get_mla_v_packed_page_offset", + "get_fp4_mla_hp_pool", + "get_disagg_role_mapper_kinds", + "get_disagg_transfer_roles", + "get_disagg_global_layer_ids", + "_get_runtime_cache_size_layer_components", + "_get_generation_request_capacity", + ) + for method_name in method_names: + method = getattr(Fp4MlaV2CacheLayoutPolicy, method_name) + setattr(manager, method_name, MethodType(method, manager)) - self._validate_fp4_mla_layer_groups() + @staticmethod + def finalize_manager(manager) -> None: + """Validate allocated pools and initialize derived FP4 state.""" + Fp4MlaV2CacheLayoutPolicy._validate_fp4_mla_layer_groups(manager) # Partial pages leave unused V-scale tiles untouched while the # fixed-width PV path can load the complete scale page. Match V1's # deterministic initialization before any warmup or graph capture. - with torch.cuda.stream(self._stream): - self.get_mla_v_scale_pool_base().zero_() - self._stream.synchronize() + with torch.cuda.stream(manager._stream): + Fp4MlaV2CacheLayoutPolicy.get_mla_v_scale_pool_base(manager).zero_() + manager._stream.synchronize() @property def fp4_mla_hp_pool_size(self) -> int: """Number of BF16 tokens in each layer's rewind-capable HP ring.""" return self._fp4_mla_hp_pool_size + def _bf16_mla_local_layer_indices(self) -> list[int]: + fallback_layers = getattr(self, "_bf16_mla_global_layer_ids", frozenset()) + is_linear_layer = getattr(self, "_is_local_mamba_layer", None) + return [ + local_layer + for local_layer in range(self.num_local_layers) + if self.pp_layers[local_layer] in fallback_layers + and (not callable(is_linear_layer) or not is_linear_layer(local_layer)) + ] + + def _fp4_mla_local_layer_indices(self) -> list[int]: + fallback_layers = set(self._bf16_mla_local_layer_indices()) + is_linear_layer = getattr(self, "_is_local_mamba_layer", None) + return [ + local_layer + for local_layer in range(self.num_local_layers) + if local_layer not in fallback_layers + and (not callable(is_linear_layer) or not is_linear_layer(local_layer)) + ] + + def _fp4_mla_compact_layer_idx(self, local_layer: int) -> int: + try: + return self._fp4_mla_local_to_compact[local_layer] + except KeyError as error: + raise ValueError( + f"Local layer {local_layer} is not an FP4 MLA attention layer." + ) from error + + def _bf16_mla_bytes_per_token(self, local_layer_idx: int) -> int: + return ( + self.num_kv_heads_per_layer[local_layer_idx] + * self.head_dim_per_layer[local_layer_idx] + * torch.empty((), dtype=torch.bfloat16).element_size() + ) + + def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: + if local_layer_idx in self._bf16_mla_local_layer_indices(): + return [Role.KEY] + return KVCacheManagerV2._get_buffer_roles_for_layer(self, local_layer_idx) + + def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + """Return the layer's primary paged-cache view. + + BF16 fallback layers use a buffer role and physical page geometry that + differ from the manager-wide NVFP4 layout. Construct their view from + the role descriptor directly so FMHA dispatch observes the same BF16 + dtype and logical shape as an all-BF16 cache manager. + """ + local_layer = self.layer_offsets[layer_idx] + is_linear_layer = getattr(self, "_is_local_mamba_layer", None) + if callable(is_linear_layer) and is_linear_layer(local_layer): + return None + if local_layer not in self._bf16_mla_local_layer_indices(): + return KVCacheManagerV2.get_buffers(self, layer_idx, kv_layout) + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported kv_layout: {kv_layout}") + + page_shape = ( + ( + self.kv_factor, + self.tokens_per_block, + self.num_kv_heads_per_layer[local_layer], + self.head_dim_per_layer[local_layer], + ) + if kv_layout == "NHD" + else ( + self.kv_factor, + self.num_kv_heads_per_layer[local_layer], + self.tokens_per_block, + self.head_dim_per_layer[local_layer], + ) + ) + return self._role_view( + LayerId(local_layer), + Role.KEY, + torch.bfloat16, + page_shape, + ) + + def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: + if layer_idx is None: + return self.dtype + local_layer = self.layer_offsets[layer_idx] + if local_layer in self._bf16_mla_local_layer_indices(): + return DataType.BF16 + return self.dtype + + def get_kv_cache_num_blocks(self, layer_idx: int) -> int: + local_layer = self.layer_offsets[layer_idx] + if local_layer in self._bf16_mla_local_layer_indices(): + manager_layer = LayerId(local_layer) + else: + manager_layer = self._cache_manager_layer_ids[ + self._fp4_mla_compact_layer_idx(local_layer) + ] + return self._role_encoded_page_capacity(manager_layer, Role.KEY) + @property def blocks_in_primary_pool(self) -> int: return self._role_encoded_page_capacity(self._cache_manager_layer_ids[0], Role.KEY) @@ -170,6 +356,12 @@ def _hp_bytes_per_page(self, local_layer_idx: int) -> int: ) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) -> int: + if local_layer_idx in self._bf16_mla_local_layer_indices(): + if data_role in (Role.KEY, Role.ALL): + return self._bf16_mla_bytes_per_token(local_layer_idx) + raise ValueError(f"Invalid BF16 MLA V2 data role: {data_role}") + if local_layer_idx not in self._fp4_mla_local_layer_indices(): + return KVCacheManagerV2.get_layer_bytes_per_token(self, local_layer_idx, data_role) storage_head_dim = self._storage_head_dim(local_layer_idx) role_sizes = { Role.KEY: math.ceil(storage_head_dim / 2), @@ -193,7 +385,7 @@ def _extra_buffers_per_layer( self, *, tokens_per_block: int ) -> Optional[dict[int, List[BufferConfig]]]: result = {} - for local_layer in range(self.num_local_layers): + for local_layer in self._fp4_mla_local_layer_indices(): buffers = [ BufferConfig( role=Role.MLA_V_SCALE, @@ -212,9 +404,39 @@ def _extra_buffers_per_layer( def _build_cache_config(self, config: KVCacheManagerConfig) -> KVCacheManagerConfig: cache_layers = list(config.layers) - self._cache_manager_layer_ids = [LayerId(i) for i in range(self.num_local_layers)] + bf16_local_layers = self._bf16_mla_local_layer_indices() + fp4_local_layers = self._fp4_mla_local_layer_indices() + if not fp4_local_layers: + raise ValueError("FP4 MLA V2 layout requires at least one local attention layer.") + self._bf16_mla_manager_layer_ids = [ + LayerId(local_layer) for local_layer in bf16_local_layers + ] + if bf16_local_layers: + bf16_lifecycle_window = ( + self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1 + ) + for local_layer in bf16_local_layers: + cache_layers[local_layer] = AttentionLayerConfig( + layer_id=LayerId(local_layer), + buffers=[ + BufferConfig( + role=Role.KEY, + size=self._bf16_mla_bytes_per_token(local_layer) + * self.tokens_per_block, + ) + ], + # A distinct, non-evicting lifecycle gives BF16 pages their + # own attention-op pool without changing the model window. + sliding_window_size=bf16_lifecycle_window, + num_sink_tokens=None, + ) + self._fp4_mla_local_to_compact = { + local_layer: compact_layer for compact_layer, local_layer in enumerate(fp4_local_layers) + } + self._fp4_mla_compact_to_local = fp4_local_layers + self._cache_manager_layer_ids = [LayerId(i) for i in fp4_local_layers] self._hp_manager_layer_ids = [] - for local_layer in range(self.num_local_layers): + for local_layer in fp4_local_layers: layer_id = LayerId(len(cache_layers)) self._hp_manager_layer_ids.append(layer_id) cache_layers.append( @@ -239,6 +461,12 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: hp_pool_id = int(self.impl.get_layer_group_id(hp_layer)) if cache_pool_id == hp_pool_id: raise RuntimeError("FP4 MLA full-history and HP state must use distinct lifecycles.") + bf16_pool_id = None + if self._bf16_mla_manager_layer_ids: + bf16_layer = self._bf16_mla_manager_layer_ids[0] + bf16_pool_id = int(self.impl.get_layer_group_id(bf16_layer)) + if bf16_pool_id in (cache_pool_id, hp_pool_id): + raise RuntimeError("FP4 MLA, BF16 MLA, and HP state must use distinct lifecycles.") num_pools = len(self.impl.layer_grouping) pool_pointers = [[[0, 0], [0, 0]] for _ in range(num_pools)] @@ -253,6 +481,18 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: [int(self.impl.get_mem_pool_base_address(hp_layer, Role.MLA_HP_TAIL)), 0], [0, 0], ] + if bf16_pool_id is not None: + pool_pointers[bf16_pool_id] = [ + [ + int( + self.impl.get_mem_pool_base_address( + self._bf16_mla_manager_layer_ids[0], Role.KEY + ) + ), + 0, + ], + [0, 0], + ] self.kv_cache_pool_pointers = torch.tensor( pool_pointers, dtype=torch.int64, @@ -260,10 +500,18 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: pin_memory=prefer_pinned(), ) - mapping = [] - for layer_id in self._cache_manager_layer_ids: + mapping = [[0, 0] for _ in range(self.num_local_layers)] + for compact_layer, layer_id in enumerate(self._cache_manager_layer_ids): converter = self.impl.get_page_index_converter(layer_id, Role.KEY) - mapping.append([cache_pool_id, int(converter.layer_offset)]) + local_layer = self._fp4_mla_compact_to_local[compact_layer] + mapping[local_layer] = [cache_pool_id, int(converter.layer_offset)] + if bf16_pool_id is not None: + for layer_id in self._bf16_mla_manager_layer_ids: + converter = self.impl.get_page_index_converter(layer_id, Role.KEY) + mapping[int(layer_id)] = [ + bf16_pool_id, + int(converter.layer_offset), + ] self.kv_cache_pool_mapping = torch.tensor( mapping, dtype=torch.int32, @@ -280,6 +528,12 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: self.index_scales[hp_pool_id] = int( self.impl.get_page_index_converter(hp_layer, Role.MLA_HP_TAIL).scale ) + if bf16_pool_id is not None: + self.index_scales[bf16_pool_id] = int( + self.impl.get_page_index_converter( + self._bf16_mla_manager_layer_ids[0], Role.KEY + ).scale + ) self.kv_offset = torch.zeros_like(self.index_scales) self._index_scale_ints = self.index_scales.tolist() self.num_attention_op_pools = num_pools @@ -306,10 +560,22 @@ def _validate_fp4_mla_layer_groups(self) -> None: "FP4 MLA V2 requires one full-history and one HP sliding layer group; " f"got cache={sorted(cache_groups)}, hp={sorted(hp_groups)}." ) + bf16_groups = { + int(self.impl.get_layer_group_id(layer_id)) + for layer_id in self._bf16_mla_manager_layer_ids + } + if self._bf16_mla_manager_layer_ids and ( + len(bf16_groups) != 1 or bf16_groups == cache_groups or bf16_groups == hp_groups + ): + raise RuntimeError( + "BF16 MLA fallback layers require one lifecycle distinct from " + f"FP4 and HP; got bf16={sorted(bf16_groups)}." + ) cache_roles = [Role.KEY, Role.KEY_BLOCK_SCALE, Role.MLA_V_SCALE] if self.mla_v_head_dim is not None: cache_roles.append(Role.MLA_V_PACKED) - for local_layer, layer_id in enumerate(self._cache_manager_layer_ids): + for compact_layer, layer_id in enumerate(self._cache_manager_layer_ids): + local_layer = self._fp4_mla_compact_to_local[compact_layer] converters = [ self.impl.get_page_index_converter(layer_id, role) for role in cache_roles ] @@ -327,12 +593,14 @@ def _validate_fp4_mla_layer_groups(self) -> None: f"one encoded page geometry for local layer {local_layer}; " f"got {sorted(geometries)}." ) - for local_layer, layer_id in enumerate(self._hp_manager_layer_ids): + num_fp4_layers = len(self._hp_manager_layer_ids) + for compact_layer, layer_id in enumerate(self._hp_manager_layer_ids): + local_layer = self._fp4_mla_compact_to_local[compact_layer] converter = self.impl.get_page_index_converter(layer_id, Role.MLA_HP_TAIL) if ( - int(converter.scale) != self.num_local_layers + int(converter.scale) != num_fp4_layers or int(converter.expansion) != 1 - or int(converter.layer_offset) != local_layer + or int(converter.layer_offset) != compact_layer ): raise RuntimeError( "FP4 MLA V2 HP roles require one coalesced page per local " @@ -340,9 +608,14 @@ def _validate_fp4_mla_layer_groups(self) -> None: ) def get_fp4_mla_page_table_spec(self, layer_idx: Optional[int] = None) -> Fp4MlaPageTableSpec: - local_layer = 0 if layer_idx is None else self.layer_offsets[layer_idx] - cache_layer = self._cache_manager_layer_ids[local_layer] - hp_layer = self._hp_manager_layer_ids[local_layer] + local_layer = ( + self._fp4_mla_compact_to_local[0] + if layer_idx is None + else self.layer_offsets[layer_idx] + ) + compact_layer = self._fp4_mla_compact_layer_idx(local_layer) + cache_layer = self._cache_manager_layer_ids[compact_layer] + hp_layer = self._hp_manager_layer_ids[compact_layer] return Fp4MlaPageTableSpec( cache_pool_id=int(self.impl.get_layer_group_id(cache_layer)), # copy_batch_block_offsets already applies the V2 converter scale. @@ -409,7 +682,8 @@ def get_fp4_mla_cache_buffers( if kv_layout != "NHD": raise ValueError("FP4 MLA V2 cache buffers support only NHD layout.") local_layer = self.layer_offsets[layer_idx] - manager_layer = self._cache_manager_layer_ids[local_layer] + compact_layer = self._fp4_mla_compact_layer_idx(local_layer) + manager_layer = self._cache_manager_layer_ids[compact_layer] storage_head_dim = self._storage_head_dim(local_layer) kv_cache = self._role_view( manager_layer, @@ -425,6 +699,55 @@ def get_fp4_mla_cache_buffers( ) return kv_cache, sf_cache + def _iter_fp4_mla_physical_pool_views(self) -> Iterable[torch.Tensor]: + """Yield page-major views spanning each physical MLA pool once.""" + local_layer = self._fp4_mla_compact_to_local[0] + cache_layer = self._cache_manager_layer_ids[0] + storage_head_dim = self._storage_head_dim(local_layer) + yield self._role_pool_base_view( + cache_layer, + Role.KEY, + torch.uint8, + (self.kv_factor, self.tokens_per_block, 1, storage_head_dim // 2), + ) + yield self._role_pool_base_view( + cache_layer, + Role.KEY_BLOCK_SCALE, + torch.uint8, + (self.tokens_per_block, storage_head_dim // FP4_BLOCK_SIZE), + ).view(torch.float8_e4m3fn) + if self._bf16_mla_manager_layer_ids: + bf16_layer = self._bf16_mla_manager_layer_ids[0] + bf16_local_layer = int(bf16_layer) + yield self._role_pool_base_view( + bf16_layer, + Role.KEY, + torch.bfloat16, + ( + self.kv_factor, + self.tokens_per_block, + self.num_kv_heads_per_layer[bf16_local_layer], + self.head_dim_per_layer[bf16_local_layer], + ), + ) + yield self.get_mla_v_scale_pool_base().view(torch.float8_e4m3fn) + if self.mla_v_head_dim is not None: + yield self._role_pool_base_view( + cache_layer, + Role.MLA_V_PACKED, + torch.uint8, + (self.mla_v_head_dim, self.tokens_per_block // 2), + ) + yield self._role_pool_base_view( + self._hp_manager_layer_ids[0], + Role.MLA_HP_TAIL, + torch.bfloat16, + ( + 1, + self._fp4_mla_hp_pool_size * self.head_dim_per_layer[local_layer], + ), + ) + def _all_layer_role_view( self, layer_ids: list[LayerId], @@ -514,14 +837,16 @@ def get_mla_v_scale_pool_base(self) -> torch.Tensor: ) def get_mla_v_scale_page_offset(self, local_layer: int) -> int: + if not 0 <= local_layer < len(self._cache_manager_layer_ids): + raise IndexError(f"Invalid compact FP4 MLA layer {local_layer}.") layer_id = self._cache_manager_layer_ids[local_layer] return int(self.impl.get_page_index_converter(layer_id, Role.MLA_V_SCALE).layer_offset) def get_mla_v_packed_pool(self, local_layer: int) -> Optional[torch.Tensor]: if self.mla_v_head_dim is None: return None - if not 0 <= local_layer < self.num_local_layers: - raise IndexError(f"Invalid FP4 MLA local layer {local_layer}.") + if not 0 <= local_layer < len(self._cache_manager_layer_ids): + raise IndexError(f"Invalid compact FP4 MLA layer {local_layer}.") pool = self._role_view( self._cache_manager_layer_ids[local_layer], Role.MLA_V_PACKED, @@ -542,6 +867,8 @@ def get_mla_v_packed_pool_base(self) -> Optional[torch.Tensor]: return pool.reshape(pool.shape[0] * pool.shape[1], pool.shape[2]) def get_mla_v_packed_page_offset(self, local_layer: int) -> int: + if not 0 <= local_layer < len(self._cache_manager_layer_ids): + raise IndexError(f"Invalid compact FP4 MLA layer {local_layer}.") layer_id = self._cache_manager_layer_ids[local_layer] return int(self.impl.get_page_index_converter(layer_id, Role.MLA_V_PACKED).layer_offset) @@ -550,20 +877,64 @@ def get_fp4_mla_hp_pool(self) -> torch.Tensor: self._hp_manager_layer_ids, Role.MLA_HP_TAIL, torch.bfloat16, - (1, self._fp4_mla_hp_pool_size * self.head_dim_per_layer[0]), + ( + 1, + self._fp4_mla_hp_pool_size + * self.head_dim_per_layer[self._fp4_mla_compact_to_local[0]], + ), ).permute(1, 0, 2, 3) + def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: + return { + Role.ALL: MapperKind.NHD, + Role.MLA_HP_TAIL: MapperKind.REPLICATED, + } + + def get_disagg_transfer_roles(self) -> Optional[frozenset[DataRole]]: + return frozenset((Role.KEY, Role.KEY_BLOCK_SCALE, Role.MLA_HP_TAIL)) + + def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: + local_layer_ids = list(self.impl.layer_grouping[layer_group_id]) + cache_local = { + int(layer_id): self._fp4_mla_compact_to_local[compact_layer] + for compact_layer, layer_id in enumerate(self._cache_manager_layer_ids) + } + cache_local.update( + {int(layer_id): int(layer_id) for layer_id in self._bf16_mla_manager_layer_ids} + ) + hp_local = { + int(layer_id): self._fp4_mla_compact_to_local[compact_layer] + for compact_layer, layer_id in enumerate(self._hp_manager_layer_ids) + } + result = [] + for layer_id in local_layer_ids: + internal_layer = int(layer_id) + if internal_layer in cache_local: + result.append(2 * int(self.pp_layers[cache_local[internal_layer]])) + elif internal_layer in hp_local: + result.append(2 * int(self.pp_layers[hp_local[internal_layer]]) + 1) + else: + raise ValueError(f"Unknown FP4 MLA V2 internal layer {internal_layer}.") + return result + def _get_runtime_cache_size_layer_components(self) -> tuple[list[int], list[Optional[int]]]: + fp4_local_layers = self._fp4_mla_local_layer_indices() + bf16_local_layers = self._bf16_mla_local_layer_indices() sizes = [ self.get_layer_bytes_per_token(local_layer, Role.ALL) - for local_layer in range(self.num_local_layers) + for local_layer in fp4_local_layers ] - windows: list[Optional[int]] = [None] * self.num_local_layers + windows: list[Optional[int]] = [None] * len(fp4_local_layers) + sizes.extend( + self.get_layer_bytes_per_token(local_layer, Role.ALL) + for local_layer in bf16_local_layers + ) + windows.extend([None] * len(bf16_local_layers)) sizes.extend( self.get_layer_bytes_per_token(local_layer, Role.MLA_HP_TAIL) - for local_layer in range(self.num_local_layers) + for local_layer in fp4_local_layers ) - windows.extend([self._fp4_mla_hp_pool_size] * self.num_local_layers) + windows.extend([self._fp4_mla_hp_pool_size] * len(fp4_local_layers)) return sizes, windows def _get_generation_request_capacity(self) -> int: @@ -607,9 +978,35 @@ def get_cache_size_per_token( local_layers = KVCacheManager._resolve_num_attention_layers( model_config, mapping, num_layers ) + bf16_layer_ids = get_kimi_k3_bf16_kv_layer_ids(config) + if bf16_layer_ids: + local_global_layers = set(mapping.pp_layers(int(config.num_hidden_layers))) + num_bf16_layers = len(bf16_layer_ids & local_global_layers) + if num_bf16_layers > local_layers: + raise ValueError( + "Kimi K3 BF16 KV-cache fallback layer count exceeds the local MLA layer count." + ) + else: + num_bf16_layers = 0 + num_fp4_layers = local_layers - num_bf16_layers rewind = int(spec_config.tokens_per_gen_step - 1) if spec_config else 0 - hp_bytes = (HP_BLOCK_SIZE + rewind) * logical_head_dim * 2 * local_layers - return per_layer * local_layers, hp_bytes * (max_batch_size or 0) * mapping.pp_size + hp_bytes = (HP_BLOCK_SIZE + rewind) * logical_head_dim * 2 * num_fp4_layers + max_batch_size = int(max_batch_size or 0) + cache_bytes = ( + per_layer * num_fp4_layers + + logical_head_dim + * torch.empty((), dtype=torch.bfloat16).element_size() + * num_bf16_layers + ) + return cache_bytes, hp_bytes * max_batch_size * mapping.pp_size + + +class Fp4MlaKVCacheManagerV2(Fp4MlaV2CacheLayoutPolicy, KVCacheManagerV2): + """Compatibility manager applying the FP4 MLA V2 layout policy.""" -__all__ = ["Fp4MlaKVCacheManagerV2", "Fp4MlaPageTableSpec"] +__all__ = [ + "Fp4MlaKVCacheManagerV2", + "Fp4MlaPageTableSpec", + "Fp4MlaV2CacheLayoutPolicy", +] diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py index 0ff9a6442e7d..0456f702fafd 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py @@ -247,6 +247,8 @@ def _scatter_fp4_mla_kv_cache_2d_context( v_sf: torch.Tensor, global_scale: torch.Tensor, rotary_cos_sin: Optional[torch.Tensor], + q_context: Optional[torch.Tensor], + q_nope_head_dim: Optional[int], *, token_offset: int, local_layer: int, @@ -278,6 +280,34 @@ def _scatter_fp4_mla_kv_cache_2d_context( ) rotary_cos_sin_ptr = rotary_cos_sin if rotary_cos_sin is not None else latent_cache + apply_q_rope = q_context is not None + block_q_heads = 16 + if apply_q_rope: + if rotary_cos_sin is None or q_nope_head_dim is None: + raise ValueError("FP4 MLA fused context Q-RoPE requires a rotary table and Q layout.") + q_head_dim = q_nope_head_dim + rope_dim + if ( + q_context.dtype != torch.bfloat16 + or q_context.device != latent_cache.device + or q_context.ndim != 2 + or q_context.shape[0] != num_tokens + or q_context.shape[1] <= 0 + or q_nope_head_dim <= 0 + or q_context.shape[1] % q_head_dim != 0 + or not q_context.is_contiguous() + ): + raise ValueError( + "FP4 MLA fused context Q-RoPE requires a contiguous same-device BF16 " + f"tensor shaped [tokens, heads * ({q_nope_head_dim} + {rope_dim})]." + ) + num_q_heads = q_context.shape[1] // q_head_dim + q_context_view = q_context.view(num_tokens, num_q_heads, q_head_dim) + q_head_blocks = triton.cdiv(num_q_heads, block_q_heads) + else: + num_q_heads = 0 + q_context_view = latent_cache + q_head_blocks = 0 + hp_pool = getattr(metadata.fp4_mla_state, "hp_pool", None) if not isinstance(hp_pool, torch.Tensor): raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") @@ -313,7 +343,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( _fp4_mla_context_cache_update_kernel[ ( num_tokens, - num_dim_blocks, + num_dim_blocks + q_head_blocks, ) ]( kv_cache, @@ -321,6 +351,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( v_sf, v_packed_output, latent_cache, + q_context_view, global_scale, rotary_cos_sin_ptr, hp_pool, @@ -347,6 +378,9 @@ def _scatter_fp4_mla_kv_cache_2d_context( sf_cache.stride(0), latent_cache.stride(0), latent_cache.stride(1), + q_context_view.stride(0), + q_context_view.stride(1) if apply_q_rope else 0, + q_context_view.stride(2) if apply_q_rope else 0, v_sf.stride(0), v_sf.stride(1), v_packed_s0, @@ -364,6 +398,11 @@ def _scatter_fp4_mla_kv_cache_2d_context( STORE_K_RESIDUAL=(_fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS), ROPE_DIM=rope_dim, APPLY_K_ROPE=apply_k_rope, + APPLY_Q_ROPE=apply_q_rope, + NUM_DIM_BLOCKS=num_dim_blocks, + NUM_Q_HEADS=num_q_heads, + Q_NOPE_DIM=q_nope_head_dim if q_nope_head_dim is not None else 0, + BLOCK_Q_HEADS=block_q_heads, POOL_HEAD_D=pool_head_dim, STORE_HP_TAIL=store_hp_tail, WRITE_V_PACKED=write_v_packed, @@ -677,6 +716,7 @@ def launch_generation_update( v_sf, v_packed_output, latent_cache, + latent_cache, global_scale, rotary_table, pool, @@ -703,6 +743,9 @@ def launch_generation_update( sf_cache.stride(0), latent_cache.stride(0), latent_cache.stride(1), + latent_cache.stride(0), + 0, + 0, v_sf.stride(0), v_sf.stride(1), v_packed_s0, @@ -720,6 +763,11 @@ def launch_generation_update( STORE_K_RESIDUAL=store_k_residual, ROPE_DIM=rope_dim, APPLY_K_ROPE=True, + APPLY_Q_ROPE=False, + NUM_DIM_BLOCKS=num_dim_blocks, + NUM_Q_HEADS=0, + Q_NOPE_DIM=0, + BLOCK_Q_HEADS=16, POOL_HEAD_D=hp_head_dim, STORE_HP_TAIL=True, WRITE_V_PACKED=write_v_packed, @@ -780,6 +828,8 @@ def scatter_fp4_mla_kv_cache( q_pe: Optional[torch.Tensor] = None, q_rope_out: Optional[torch.Tensor] = None, q_quant_input: Optional[torch.Tensor] = None, + q_context: Optional[torch.Tensor] = None, + q_nope_head_dim: Optional[int] = None, ) -> bool: """Quantize MLA latent tokens and scatter them into the paged FP4 cache. @@ -799,8 +849,10 @@ def scatter_fp4_mla_kv_cache( layouts. Tail K-only dimensions use K's per-token 1D scales. For exclusively owned CuTeDSL pages, context scatter also writes the persistent packed-V sidecar. - Context scatter can rotate the K tail directly from the unassembled latent - tensor. Generation scatter rewrites each touched 16-token tile by reading + Context scatter can rotate Q in place and rotate the K tail directly from + the unassembled latent tensor. When context Q is supplied for chunked + prefill, it also writes rotated K back for current-chunk attention. + Generation scatter rewrites each touched 16-token tile by reading old tokens from the HP pool and new tokens from ``latent_cache``. The static-scale generation specialization can also rotate Q and new K tails while updating the HP pool. @@ -816,6 +868,11 @@ def scatter_fp4_mla_kv_cache( metadata.fp4_mla_state.q_batch_capacity = None if latent_cache.numel() == 0: raise ValueError("FP4 MLA cache scatter requires at least one latent token.") + if q_context is not None and not latent_cache.is_contiguous(): + raise ValueError( + "FP4 MLA fused context Q/K RoPE requires contiguous latent_cache " + "storage for in-place current-K update." + ) latent_cache = latent_cache.reshape(latent_cache.shape[0], -1).contiguous() num_tokens = latent_cache.shape[0] @@ -883,8 +940,12 @@ def scatter_fp4_mla_kv_cache( if phase == "context": if any(arg is not None for arg in (q_pe, q_rope_out, q_quant_input)): raise ValueError("FP4 MLA context cache update does not accept generation Q tensors.") + if (q_context is None) != (q_nope_head_dim is None): + raise ValueError("FP4 MLA context Q and q_nope_head_dim must be provided together.") hp_pool_updated = False else: + if q_context is not None or q_nope_head_dim is not None: + raise ValueError("FP4 MLA generation cache update does not accept context Q tensors.") if not all(arg is not None for arg in generation_inputs): raise ValueError( "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " @@ -985,6 +1046,8 @@ def scatter_fp4_mla_kv_cache( v_sf, global_scale, rotary_cos_sin, + q_context, + q_nope_head_dim, token_offset=token_offset, local_layer=local_layer, v_head_dim=v_head_dim, diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py index df30ec7143cc..482f4c9cd7ce 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py @@ -107,6 +107,7 @@ class _Fp8MlaContextScratch: max_num_sequences: int max_blocks_per_seq: int capacity_blocks: int + max_num_tokens: int page_size: int head_dim: int cache_stream: torch.cuda.Stream @@ -131,6 +132,12 @@ def create( max_num_sequences = int(meta.max_num_sequences or meta.max_num_requests) max_blocks_per_seq = int(kv_cache_manager.max_blocks_per_seq) max_num_tokens = int(meta.max_num_tokens) + runtime_features = meta.runtime_features + if runtime_features.chunked_prefill: + max_num_tokens *= max( + 1, + int(runtime_features.chunked_prefill_buffer_batch_size), + ) max_nonempty_sequences = min(max_num_sequences, max_num_tokens) capacity_blocks = max( 1, @@ -191,6 +198,7 @@ def create( max_num_sequences=max_num_sequences, max_blocks_per_seq=max_blocks_per_seq, capacity_blocks=capacity_blocks, + max_num_tokens=max_num_tokens, page_size=page_size, head_dim=head_dim, cache_stream=torch.cuda.Stream(device=device), @@ -207,19 +215,43 @@ def matches( head_dim: int, ) -> bool: kv_cache_manager = meta.kv_cache_manager + required_max_num_tokens = int(meta.max_num_tokens) + if meta.runtime_features.chunked_prefill: + required_max_num_tokens *= max( + 1, + int(meta.runtime_features.chunked_prefill_buffer_batch_size), + ) return ( kv_cache_manager is not None and self.pool.device == device and self.head_dim == head_dim + and self.max_num_tokens >= required_max_num_tokens and self.page_size == meta.tokens_per_block and self.cache_manager_view.max_seq_len == int(kv_cache_manager.max_seq_len) and self.max_num_sequences >= int(meta.max_num_sequences or meta.max_num_requests) and self.max_blocks_per_seq >= int(kv_cache_manager.max_blocks_per_seq) ) - def prepare(self, meta: "TrtllmAttentionMetadata") -> None: + def prepare( + self, + meta: "TrtllmAttentionMetadata", + *, + context_lengths_cuda: Optional[torch.Tensor] = None, + context_lengths_cpu: Optional[torch.Tensor] = None, + ) -> None: + if context_lengths_cuda is None: + context_lengths_cuda = meta.prompt_lens_cuda_runtime + if context_lengths_cpu is None: + context_lengths_cpu = meta.prompt_lens_cpu_runtime + if ( + context_lengths_cpu.device.type != "cpu" + or context_lengths_cpu.dtype != torch.int32 + or context_lengths_cpu.ndim != 1 + or context_lengths_cpu.numel() < meta.num_contexts + ): + raise ValueError("FP8 MLA context metadata requires a CPU int32 context-length tensor.") context_lengths = tuple( - int(length) for length in meta.prompt_lens_cpu_runtime[: meta.num_contexts].tolist() + int(length) for length in context_lengths_cpu[: meta.num_contexts].tolist() ) self.host_total_kv_lens[0] = sum(context_lengths) self.host_total_kv_lens[1] = 0 @@ -248,7 +280,6 @@ def prepare(self, meta: "TrtllmAttentionMetadata") -> None: f"{self.capacity_blocks} were allocated." ) - context_lengths_cuda = meta.prompt_lens_cuda_runtime if ( context_lengths_cuda.dtype != torch.int32 or not context_lengths_cuda.is_cuda @@ -297,8 +328,15 @@ def _build_fp8_mla_context_attn(attn: "TrtllmAttention") -> "TrtllmAttention": def _build_fp8_mla_context_metadata( meta: "TrtllmAttentionMetadata", scratch: _Fp8MlaContextScratch, + *, + kv_lens_cuda: Optional[torch.Tensor] = None, + kv_lens_cpu: Optional[torch.Tensor] = None, ) -> "TrtllmAttentionMetadata": """Route the mandatory FP8 cache write through a direct metadata view.""" + if kv_lens_cuda is None: + kv_lens_cuda = meta.prompt_lens_cuda_runtime[: meta.num_contexts] + if kv_lens_cpu is None: + kv_lens_cpu = meta.prompt_lens_cpu_runtime[: meta.num_contexts] fp8_meta = copy.copy(meta) fp8_meta.fp4_mla_state = None fp8_meta.kv_cache_manager = scratch.cache_manager_view @@ -310,8 +348,8 @@ def _build_fp8_mla_context_metadata( # The disposable cache represents only this context invocation. Expose # exact context-only lengths so cached prefixes, trailing generation # metadata, and stale totals cannot extend FP8 K/V quantization. - fp8_meta.kv_lens_cuda_runtime = meta.prompt_lens_cuda_runtime[: meta.num_contexts] - fp8_meta.kv_lens_runtime = meta.prompt_lens_cpu_runtime[: meta.num_contexts] + fp8_meta.kv_lens_cuda_runtime = kv_lens_cuda + fp8_meta.kv_lens_runtime = kv_lens_cpu fp8_meta.host_total_kv_lens = scratch.host_total_kv_lens # Scratch lengths intentionally start from zero. Preserve the actual # absolute positions for Q/K RoPE through the native kernel's explicit diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py index 5295c6f788d0..f844195e0df2 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_kernels.py @@ -286,6 +286,96 @@ def _fp4_mla_swizzled_sf_offset( # FP4 conversion and cache kernels +@triton.jit( + do_not_specialize=[ + "num_page_ids", + "num_pages", + "num_layers", + "local_layer", + ], + do_not_specialize_on_alignment=[ + "num_page_ids", + "num_pages", + "num_layers", + "local_layer", + ], +) +def _fp4_mla_rebuild_v_scale_from_k_scale_kernel( + k_sf_ptr, + v_sf_ptr, + page_ids_ptr, + page_valid_tokens_ptr, + num_page_ids, + num_pages, + num_layers, + local_layer, + page_size, + k_sf_s0, + v_sf_s0, + v_sf_s1, + V_HEAD_D: tl.constexpr, + HP_BLOCK: tl.constexpr, + SF_PER_TOKEN: tl.constexpr, + SF_PER_PAGE: tl.constexpr, + BLOCK_TOKEN_GROUPS: tl.constexpr, +): + """Rebuild dim-major MLA V scales from transferred token-major K scales. + + Context quantization uses one shared scale for every 16x16 compressed-V + tile. That byte is repeated across the tile's 16 K rows and 16 V rows, + so the process-local V layout can be reconstructed exactly from the first + valid K row of every token tile. + """ + page_work_idx = tl.program_id(0) + dim_block = tl.program_id(1) + if page_work_idx >= num_page_ids: + return + if (local_layer < 0) | (local_layer >= num_layers): + return + + physical_page = tl.load(page_ids_ptr + page_work_idx).to(tl.int64) + if (physical_page < 0) | (physical_page >= num_pages): + return + + valid_tokens = tl.load(page_valid_tokens_ptr + page_work_idx).to(tl.int32) + valid_tokens = tl.maximum(0, tl.minimum(valid_tokens, page_size)) + token_groups = tl.arange(0, BLOCK_TOKEN_GROUPS) + token_group_valid = (token_groups * HP_BLOCK < valid_tokens) & ( + token_groups * HP_BLOCK < page_size + ) + + # Scale tensors are passed as uint8 views so this is a bit-exact copy of + # the FP8 E4M3 encoding, including signed zero if it is ever produced. + k_rows = token_groups * HP_BLOCK + k_sf_offsets = _fp4_mla_swizzled_sf_offset( + k_rows, + dim_block, + SF_PER_TOKEN, + ) + scale_bits = tl.load( + k_sf_ptr + physical_page * k_sf_s0 + k_sf_offsets, + mask=token_group_valid, + other=0, + ) + + dims = dim_block * HP_BLOCK + tl.arange(0, HP_BLOCK) + dim_valid = dims < V_HEAD_D + safe_dims = tl.where(dim_valid, dims, 0) + v_sf_offsets = _fp4_mla_swizzled_sf_offset( + safe_dims[:, None], + token_groups[None, :], + SF_PER_PAGE, + ) + v_sf_base = tl.cast(local_layer, tl.int64) * tl.cast( + v_sf_s0, tl.int64 + ) + physical_page * tl.cast(v_sf_s1, tl.int64) + tl.store( + v_sf_ptr + v_sf_base + v_sf_offsets.to(tl.int64), + scale_bits[None, :], + mask=dim_valid[:, None] & (token_groups[None, :] * HP_BLOCK < page_size), + ) + + @triton.jit def _fp4_e2m1_to_f32(nibble): magnitude = nibble & 0x7 @@ -314,6 +404,130 @@ def _fp4_e2m1_to_f32(nibble): return tl.where(sign, -value, value) +@triton.jit +def _fp4_mla_chunked_cache_gather_kernel( + compressed_kv_ptr, + k_pe_ptr, + kv_cache_ptr, + sf_cache_ptr, + page_ids_ptr, + cu_chunked_seq_len_ptr, + chunked_global_offset_ptr, + global_scale_ptr, + max_chunk_len, + page_table_stride, + num_pages, + page_size, + kv_s0, + kv_s2, + kv_s4, + sf_s0, + compressed_kv_s0, + k_pe_s0, + KV_LORA_RANK: tl.constexpr, + QK_ROPE_HEAD_DIM: tl.constexpr, + FP4_BLOCK: tl.constexpr, + SF_PER_TOKEN: tl.constexpr, + TOKEN_BLOCK: tl.constexpr, +): + """Gather one logical MLA prefix chunk from the canonical FP4 K view.""" + token_block = tl.program_id(0) + batch_idx = tl.program_id(1) + dim_block = tl.program_id(2) + + local_tokens = token_block * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK) + chunk_start = tl.load(cu_chunked_seq_len_ptr + batch_idx).to(tl.int64) + chunk_end = tl.load(cu_chunked_seq_len_ptr + batch_idx + 1).to(tl.int64) + chunk_len = chunk_end - chunk_start + valid_tokens = (local_tokens < chunk_len) & (local_tokens < max_chunk_len) + + logical_positions = tl.load(chunked_global_offset_ptr + batch_idx).to( + tl.int64 + ) + local_tokens.to(tl.int64) + logical_page = logical_positions // page_size + page_position = logical_positions - logical_page * page_size + page_table_offsets = batch_idx * page_table_stride + logical_page + valid_tokens = valid_tokens & (logical_page >= 0) & (logical_page < page_table_stride) + physical_pages = tl.load(page_ids_ptr + page_table_offsets, mask=valid_tokens, other=0).to( + tl.int64 + ) + valid_tokens = valid_tokens & (physical_pages >= 0) & (physical_pages < num_pages) + + dims = dim_block * FP4_BLOCK + tl.arange(0, FP4_BLOCK) + head_dim = KV_LORA_RANK + QK_ROPE_HEAD_DIM + valid_dims = dims < head_dim + packed_cols = dims // 2 + packed = tl.load( + kv_cache_ptr + + physical_pages[:, None] * kv_s0 + + page_position[:, None] * kv_s2 + + packed_cols[None, :] * kv_s4, + mask=valid_tokens[:, None] & valid_dims[None, :], + other=0, + ).to(tl.uint8) + nibbles = tl.where((dims[None, :] & 1) == 0, packed & 0x0F, packed >> 4) + + sf_offsets = _fp4_mla_swizzled_sf_offset( + page_position, + dim_block, + SF_PER_TOKEN, + ) + scale = tl.load( + sf_cache_ptr + physical_pages * sf_s0 + sf_offsets, + mask=valid_tokens, + other=0.0, + ).to(tl.float32) + global_scale = tl.load(global_scale_ptr).to(tl.float32) + values = _fp4_e2m1_to_f32(nibbles) * scale[:, None] / global_scale + + if dim_block * FP4_BLOCK >= KV_LORA_RANK: + residual_group = dim_block - KV_LORA_RANK // FP4_BLOCK + residual_packed_cols = ( + head_dim // 2 + residual_group * (FP4_BLOCK // 2) + dims % FP4_BLOCK // 2 + ) + residual_packed = tl.load( + kv_cache_ptr + + physical_pages[:, None] * kv_s0 + + page_position[:, None] * kv_s2 + + residual_packed_cols[None, :] * kv_s4, + mask=valid_tokens[:, None] & valid_dims[None, :], + other=0, + ).to(tl.uint8) + residual_nibbles = tl.where( + (dims[None, :] & 1) == 0, + residual_packed & 0x0F, + residual_packed >> 4, + ) + residual_sf_col = head_dim // FP4_BLOCK + residual_group + residual_sf_offsets = _fp4_mla_swizzled_sf_offset( + page_position, + residual_sf_col, + SF_PER_TOKEN, + ) + residual_scale = tl.load( + sf_cache_ptr + physical_pages * sf_s0 + residual_sf_offsets, + mask=valid_tokens, + other=0.0, + ).to(tl.float32) + values += _fp4_e2m1_to_f32(residual_nibbles) * residual_scale[:, None] / global_scale + + output_tokens = chunk_start + local_tokens + if dim_block * FP4_BLOCK < KV_LORA_RANK: + output_dims = dims + tl.store( + compressed_kv_ptr + output_tokens[:, None] * compressed_kv_s0 + output_dims[None, :], + values, + mask=valid_tokens[:, None] & (output_dims[None, :] < KV_LORA_RANK), + ) + else: + output_dims = dims - KV_LORA_RANK + tl.store( + k_pe_ptr + output_tokens[:, None] * k_pe_s0 + output_dims[None, :], + values, + mask=valid_tokens[:, None] & (output_dims[None, :] < QK_ROPE_HEAD_DIM), + ) + + @triton.jit def _fp4_mla_floor_half_up(abs_value, multiplier, bias): return tl.inline_asm_elementwise( @@ -404,6 +618,7 @@ def _fp4_mla_context_cache_update_kernel( v_sf_ptr, v_packed_ptr, latent_cache_ptr, + q_context_ptr, global_scale_ptr, rotary_cos_sin_ptr, hp_pool_ptr, @@ -430,6 +645,9 @@ def _fp4_mla_context_cache_update_kernel( sf_s0, lc_s0, lc_s1, + q_s0, + q_s1, + q_s2, vsf_s0, vsf_s1, v_packed_s0, @@ -447,12 +665,17 @@ def _fp4_mla_context_cache_update_kernel( STORE_K_RESIDUAL: tl.constexpr, ROPE_DIM: tl.constexpr, APPLY_K_ROPE: tl.constexpr, + APPLY_Q_ROPE: tl.constexpr, + NUM_DIM_BLOCKS: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + Q_NOPE_DIM: tl.constexpr, + BLOCK_Q_HEADS: tl.constexpr, POOL_HEAD_D: tl.constexpr, STORE_HP_TAIL: tl.constexpr, WRITE_V_PACKED: tl.constexpr, ): token_idx = tl.program_id(0) - dim_block = tl.program_id(1) + work_idx = tl.program_id(1) if (local_layer < 0) | (local_layer >= num_layers): return if token_idx >= num_tokens: @@ -466,6 +689,47 @@ def _fp4_mla_context_cache_update_kernel( position = tl.load(positions_ptr + metadata_token_idx).to(tl.int64) if (batch_idx < 0) | (batch_idx + 1 >= indptr_len) | (position < 0): return + + if APPLY_Q_ROPE and work_idx >= NUM_DIM_BLOCKS: + q_head_block = work_idx - NUM_DIM_BLOCKS + q_heads = q_head_block * BLOCK_Q_HEADS + tl.arange(0, BLOCK_Q_HEADS) + q_head_mask = q_heads < NUM_Q_HEADS + pair_offsets = tl.arange(0, ROPE_DIM // 2) + rotary_offsets = position * (ROPE_DIM * 2) + pair_offsets * 2 + cos = tl.load(rotary_cos_sin_ptr + rotary_offsets).to(tl.float32) + sin = tl.load(rotary_cos_sin_ptr + rotary_offsets + 1).to(tl.float32) + q_base = token_idx * q_s0 + q_heads[:, None].to(tl.int64) * q_s1 + Q_NOPE_DIM * q_s2 + q_even_offsets = q_base + (pair_offsets[None, :] * 2).to(tl.int64) * q_s2 + q_odd_offsets = q_even_offsets + q_s2 + q_even = tl.load( + q_context_ptr + q_even_offsets, + mask=q_head_mask[:, None], + other=0.0, + ).to(tl.float32) + q_odd = tl.load( + q_context_ptr + q_odd_offsets, + mask=q_head_mask[:, None], + other=0.0, + ).to(tl.float32) + q_roped_even, q_roped_odd = _fp4_mla_rope_fp32( + q_even, + q_odd, + cos[None, :], + sin[None, :], + ) + tl.store( + q_context_ptr + q_even_offsets, + q_roped_even.to(tl.bfloat16), + mask=q_head_mask[:, None], + ) + tl.store( + q_context_ptr + q_odd_offsets, + q_roped_odd.to(tl.bfloat16), + mask=q_head_mask[:, None], + ) + return + + dim_block = work_idx if position % HP_BLOCK != 0: return @@ -551,6 +815,24 @@ def _fp4_mla_context_cache_update_kernel( roped_even, roped_odd = _fp4_mla_rope_fp32(even_values, odd_values, cos, sin) roped_even = roped_even.to(tl.bfloat16).to(tl.float32) roped_odd = roped_odd.to(tl.bfloat16).to(tl.float32) + if APPLY_Q_ROPE: + # Only chunked prefill consumes the rotated current K from + # latent_cache. Normal context overlaps this cache update with + # FP8 attention on another stream, so its input must stay read-only. + tl.store( + latent_cache_ptr + + safe_token_candidates[:, None] * lc_s0 + + safe_even_d[None, :] * lc_s1, + roped_even, + mask=rotary_mask, + ) + tl.store( + latent_cache_ptr + + safe_token_candidates[:, None] * lc_s0 + + safe_odd_d[None, :] * lc_s1, + roped_odd, + mask=rotary_mask, + ) even_values = tl.where(rotary_mask, roped_even, even_values) odd_values = tl.where(rotary_mask, roped_odd, odd_values) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py index a9f2b83f78c4..2a3a44eb417b 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py @@ -56,6 +56,30 @@ def _fp4_mla_swizzled_sf_offset_row_block( return col_part + row_part + row_group * (128 * padded_cols) +@triton.jit +def _fp4_e2m1_quantize(x): + abs_x = tl.abs(x) + magnitude = tl.where( + abs_x < 0.25, + 0, + tl.where( + abs_x < 0.75, + 1, + tl.where( + abs_x < 1.25, + 2, + tl.where( + abs_x < 1.75, + 3, + tl.where(abs_x < 2.5, 4, tl.where(abs_x < 3.5, 5, tl.where(abs_x < 5.0, 6, 7))), + ), + ), + ), + ) + sign = tl.where(x < 0.0, 8, 0) + return (magnitude | sign).to(tl.uint8) + + @triton.jit def _fp4_pack_low_nibbles(even_packed, odd_packed): """PTX helper: pack the low nibbles of two bytes into one byte (low + high<<4).""" @@ -288,7 +312,7 @@ def _fp4_mla_attention_v_repack_pages_dyn_kernel( out_desc.store([row_base.to(tl.int32), 0], v_vals) -def fp4_mla_repack_v_cache_triton( +def fp4_mla_repack_v_cache( v_packed: Any, kv_cache: Any, page_ids: Optional[Any] = None, @@ -990,6 +1014,533 @@ def _fp4_mla_attention_page_stats_kernel( tl.store(page_sum_ptr + out_offsets, page_sum, mask=mask_h) +@triton.jit +def _fp4_mla_attention_page_stats_grouped_kernel( + page_max_ptr, + page_sum_ptr, + p_fp4_ptr, + p_sf_ptr, + q_fp4_ptr, + q_sf_ptr, + kv_cache_ptr, + sf_cache_ptr, + global_scale_ptr, + q_global_scale_ptr, + src_page_ids_ptr, + paged_kv_indptr_decode_ptr, + kv_lens_ptr, + page_ids_len, + num_pages, + q_fp4_s0, + q_fp4_s1, + kv_s0, + kv_s2, + kv_s4, + sf_s0, + page_stats_s0, + page_stats_s1, + p_s0, + p_s1, + p_num_rows, + q_num_rows, + sm_scale, + NUM_HEADS: tl.constexpr, + Q_HEAD_D: tl.constexpr, + K_HEAD_D: tl.constexpr, + Q_RESIDUAL_D: tl.constexpr, + PAGE_SIZE: tl.constexpr, + FP4_BLOCK: tl.constexpr, + Q_SF_PER_TOKEN: tl.constexpr, + K_SF_PER_TOKEN: tl.constexpr, + SF_PER_PAGE: tl.constexpr, + P_GLOBAL_SCALE: tl.constexpr, + QUERY_LEN_PER_SEQ: tl.constexpr, + MAX_PAGES: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_K: tl.constexpr, + FULL_BLOCK_END: tl.constexpr, + TAIL_BLOCK_K: tl.constexpr, + GROUP_PAGES: tl.constexpr, + ASSUME_FULL_PAGES: tl.constexpr, + ASSUME_VALID_PAGES: tl.constexpr, + PAGE_LOOP_STAGES: tl.constexpr, + occupancy: tl.constexpr = 1, +): + """Grouped page-stats QK + softmax-stats + FP4 P pack. + + Functionally identical to ``_fp4_mla_attention_page_stats_kernel`` for the + perfect decode shape (NUM_HEADS == BLOCK_H, TMA + PACK_PROBS, the standard + 640/576/64 residual-Q layout) but each program owns one + ``(query, head_block, page_group)`` and walks ``GROUP_PAGES`` pages in a + pipelined loop. Q (and its scales) and the TMA descriptors are loaded once + and reused across the group, eliminating the per-page Q reload and the tiny + per-CTA prologue that made the one-page-per-CTA kernel work-bound at long + context. Per-page outputs (page_max/page_sum and packed P) are written + exactly as the one-page kernel writes them, so every downstream stage is + unchanged. + """ + query_idx = tl.program_id(0) + head_block = tl.program_id(1) + page_group = tl.program_id(2) + seq_idx = query_idx // QUERY_LEN_PER_SEQ + query_offset = query_idx - seq_idx * QUERY_LEN_PER_SEQ + + head_start = head_block * BLOCK_H + offs_h = head_start + tl.arange(0, BLOCK_H) + offs_t = tl.arange(0, BLOCK_T) + q_row_base = query_idx * NUM_HEADS + + if ASSUME_FULL_PAGES: + kv_len = 0 + else: + kv_len = tl.load(kv_lens_ptr + seq_idx) - (QUERY_LEN_PER_SEQ - 1 - query_offset) + kv_len = tl.maximum(kv_len, 0) + page_table_start = tl.load(paged_kv_indptr_decode_ptr + seq_idx).to(tl.int64) + + q_gscale = tl.load(q_global_scale_ptr) + + residual_groups = Q_RESIDUAL_D // FP4_BLOCK + non_residual_groups = K_HEAD_D // FP4_BLOCK - residual_groups + + # ---- Hoisted, page-independent index tensors and Q tiles ---- + # Main window (q_start == 0): the first FULL_BLOCK_END elements sit entirely + # in the non-residual region, so Q and K map 1:1 and load contiguously. + # Q rows are global rows (query_idx * NUM_HEADS + head); the swizzle's + # row_group term selects this query's scale/tail block, so q_row_base must + # be folded in (the main q_vals descriptor load already does this via its + # row coordinate). Keep int32 to match the one-page kernel -- int64 swizzle + # math is emulated and was measured ~2x slower; the max global row index + # (num_queries * NUM_HEADS * stride) stays well within int32. + q_rows = q_row_base + offs_h + scale_offsets = tl.arange(0, BLOCK_K // FP4_BLOCK) + q_sf_cols = scale_offsets + q_sf_offsets = _fp4_mla_swizzled_sf_offset(q_rows[:, None], q_sf_cols[None, :], Q_SF_PER_TOKEN) + k_sf_offsets_main = _fp4_mla_swizzled_sf_offset( + offs_t[:, None], q_sf_cols[None, :], K_SF_PER_TOKEN + ) + q_scales = tl.load(q_sf_ptr + q_sf_offsets) + + # Tail window (q_start == FULL_BLOCK_END): the residual-Q groups, each of + # which maps onto a duplicated K residual group. + tail_packed_offsets = tl.arange(0, TAIL_BLOCK_K // 2) + tail_scale_offsets = tl.arange(0, TAIL_BLOCK_K // FP4_BLOCK) + qt_elem = FULL_BLOCK_END + tail_packed_offsets * 2 + qt_group = qt_elem // FP4_BLOCK + kt_group = tl.where( + qt_group < non_residual_groups, + qt_group, + non_residual_groups + (qt_group - non_residual_groups) // 2, + ) + byte_t = (qt_elem % FP4_BLOCK) // 2 + packed_qt_cols = FULL_BLOCK_END // 2 + tail_packed_offsets + packed_kt_cols = kt_group * (FP4_BLOCK // 2) + byte_t + qt_sf_cols = FULL_BLOCK_END // FP4_BLOCK + tail_scale_offsets + kt_sf_cols = tl.where( + qt_sf_cols < non_residual_groups, + qt_sf_cols, + non_residual_groups + (qt_sf_cols - non_residual_groups) // 2, + ) + qt_sf_offsets = _fp4_mla_swizzled_sf_offset( + q_rows[:, None], qt_sf_cols[None, :], Q_SF_PER_TOKEN + ) + kt_sf_offsets = _fp4_mla_swizzled_sf_offset( + offs_t[:, None], kt_sf_cols[None, :], K_SF_PER_TOKEN + ) + q_tail_scales = tl.load(q_sf_ptr + qt_sf_offsets) + q_tail_vals = tl.load( + q_fp4_ptr + q_rows[:, None] * q_fp4_s0 + packed_qt_cols[None, :] * q_fp4_s1 + ) + + tl.assume(q_fp4_s0 % 8 == 0) + tl.assume(q_fp4_s1 == 1) + tl.assume(kv_s0 % 8 == 0) + tl.assume(kv_s2 % 8 == 0) + tl.assume(kv_s4 == 1) + tl.assume(p_s0 % 8 == 0) + tl.assume(p_s1 == 1) + q_desc = tl.make_tensor_descriptor( + q_fp4_ptr, + shape=[q_num_rows, Q_HEAD_D // 2], + strides=[q_fp4_s0, q_fp4_s1], + block_shape=[BLOCK_H, BLOCK_K // 2], + ) + k_desc = tl.make_tensor_descriptor( + kv_cache_ptr, + shape=[num_pages, BLOCK_T, K_HEAD_D // 2], + strides=[kv_s0, kv_s2, kv_s4], + block_shape=[1, BLOCK_T, BLOCK_K // 2], + ) + p_desc = tl.make_tensor_descriptor( + p_fp4_ptr, + shape=[p_num_rows, PAGE_SIZE // 2], + strides=[p_s0, p_s1], + block_shape=[BLOCK_H, PAGE_SIZE // 2], + ) + q_vals = q_desc.load([(q_row_base + head_start).to(tl.int32), 0]) + + scale_cols = tl.arange(0, SF_PER_PAGE) + byte_offsets = tl.arange(0, FP4_BLOCK // 2) + byte_cols = scale_cols[:, None] * (FP4_BLOCK // 2) + byte_offsets[None, :] + + page_lo = page_group * GROUP_PAGES + page_hi = page_lo + GROUP_PAGES + for page_rel in tl.range(page_lo, page_hi, num_stages=PAGE_LOOP_STAGES): + if page_rel < MAX_PAGES: + page_start = page_rel * PAGE_SIZE + page_max = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) + page_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) + if ASSUME_FULL_PAGES or page_start < kv_len: + compact_page = page_table_start + page_rel + if ASSUME_VALID_PAGES: + physical_page = tl.load(src_page_ids_ptr + compact_page).to(tl.int64) + safe_physical_page = physical_page + else: + valid_compact_page = (compact_page >= 0) & (compact_page < page_ids_len) + safe_compact_page = tl.where(valid_compact_page, compact_page, 0) + physical_page = tl.load( + src_page_ids_ptr + safe_compact_page, mask=valid_compact_page, other=-1 + ).to(tl.int64) + valid_physical_page = ( + valid_compact_page & (physical_page >= 0) & (physical_page < num_pages) + ) + safe_physical_page = tl.where(valid_physical_page, physical_page, 0) + + global_scale = tl.load(global_scale_ptr) + qk_scale = sm_scale / (q_gscale * global_scale) + + k_vals = k_desc.load([safe_physical_page.to(tl.int32), 0, 0]) + k_vals = tl.reshape(k_vals, (BLOCK_T, BLOCK_K // 2)) + if not ASSUME_VALID_PAGES: + k_vals = tl.where(valid_physical_page, k_vals, 0) + k_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + k_sf_offsets_main) + scores = tl.dot_scaled( + q_vals, + q_scales, + "e2m1", + k_vals.T, + k_scales, + "e2m1", + fast_math=True, + rhs_k_pack=True, + ) + + kt_ptrs = ( + kv_cache_ptr + + safe_physical_page * kv_s0 + + offs_t[:, None].to(tl.int64) * kv_s2 + + packed_kt_cols[None, :] * kv_s4 + ) + if ASSUME_VALID_PAGES: + kt_vals = tl.load(kt_ptrs) + else: + kt_vals = tl.load(kt_ptrs, mask=valid_physical_page, other=0) + kt_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + kt_sf_offsets) + scores = tl.dot_scaled( + q_tail_vals, + q_tail_scales, + "e2m1", + kt_vals.T, + kt_scales, + "e2m1", + acc=scores, + fast_math=True, + rhs_k_pack=True, + ) + + if ASSUME_FULL_PAGES: + scores = scores * qk_scale + page_max = tl.max(scores, axis=1) + exp_scores = tl.math.exp2((scores - page_max[:, None]) * _LOG2_E) + page_sum = tl.sum(exp_scores, axis=1) + else: + valid_t = page_start + offs_t < kv_len + scores = tl.where(valid_t[None, :], scores * qk_scale, -float("inf")) + page_max = tl.max(scores, axis=1) + exp_scores = tl.math.exp2((scores - page_max[:, None]) * _LOG2_E) + exp_scores = tl.where(valid_t[None, :], exp_scores, 0.0) + page_sum = tl.sum(exp_scores, axis=1) + + grouped_probs = tl.reshape(exp_scores, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK)) + amax = tl.max(grouped_probs, axis=2) + inv_local_scale = tl.where(amax > 0.0, 6.0 / amax, 1.0) + stored_scale = tl.where( + amax > 0.0, + tl.minimum(amax * (P_GLOBAL_SCALE / 6.0), 448.0), + 1.0, + ) + scaled_probs = grouped_probs * tl.reshape( + inv_local_scale, (BLOCK_H, SF_PER_PAGE, 1) + ) + pairs = tl.reshape(scaled_probs, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK // 2, 2)) + even_probs, odd_probs = tl.split(pairs) + packed = _fp4_e2m1_quantize_packed(even_probs, odd_probs) + + p_page = query_idx * MAX_PAGES + page_rel + if ASSUME_VALID_PAGES and NUM_HEADS == 128 and BLOCK_H == 128: + sf_offsets = _fp4_mla_swizzled_sf_offset_row_block( + p_page, offs_h[:, None], scale_cols[None, :], SF_PER_PAGE + ) + else: + p_rows = (p_page * NUM_HEADS + offs_h).to(tl.int64) + sf_offsets = _fp4_mla_swizzled_sf_offset( + p_rows[:, None], scale_cols[None, :], SF_PER_PAGE + ) + if ASSUME_VALID_PAGES: + tl.store(p_sf_ptr + sf_offsets, stored_scale) + p_desc.store( + [(p_page * NUM_HEADS + head_start).to(tl.int32), 0], + tl.reshape(packed, (BLOCK_H, PAGE_SIZE // 2)), + ) + else: + tl.store(p_sf_ptr + sf_offsets, stored_scale, mask=valid_compact_page) + p_rows = (p_page * NUM_HEADS + offs_h).to(tl.int64) + tl.store( + p_fp4_ptr + p_rows[:, None, None] * p_s0 + byte_cols[None, :, :] * p_s1, + packed, + mask=valid_compact_page, + ) + + out_offsets = query_idx * page_stats_s0 + page_rel * page_stats_s1 + offs_h + tl.store(page_max_ptr + out_offsets, page_max) + tl.store(page_sum_ptr + out_offsets, page_sum) + + +@triton.jit +def _fp4_mla_attention_page_stats_mtp_kernel( + page_max_ptr, + page_sum_ptr, + p_fp4_ptr, + p_sf_ptr, + q_fp4_ptr, + q_sf_ptr, + kv_cache_ptr, + sf_cache_ptr, + global_scale_ptr, + q_global_scale_ptr, + src_page_ids_ptr, + paged_kv_indptr_decode_ptr, + kv_lens_ptr, + page_ids_len, + num_pages, + q_fp4_s0, + q_fp4_s1, + kv_s0, + kv_s2, + kv_s4, + sf_s0, + page_stats_s0, + page_stats_s1, + p_s0, + p_s1, + p_num_rows, + q_num_rows, + sm_scale, + NUM_HEADS: tl.constexpr, + Q_HEAD_D: tl.constexpr, + K_HEAD_D: tl.constexpr, + Q_RESIDUAL_D: tl.constexpr, + PAGE_SIZE: tl.constexpr, + FP4_BLOCK: tl.constexpr, + Q_SF_PER_TOKEN: tl.constexpr, + K_SF_PER_TOKEN: tl.constexpr, + SF_PER_PAGE: tl.constexpr, + P_GLOBAL_SCALE: tl.constexpr, + QUERY_LEN_PER_SEQ: tl.constexpr, + MAX_PAGES: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_K: tl.constexpr, + FULL_BLOCK_END: tl.constexpr, + TAIL_BLOCK_K: tl.constexpr, + occupancy: tl.constexpr = 1, +): + """MTP-fused page-stats: one CTA owns (seq, head_block, page) and processes + all QUERY_LEN_PER_SEQ linear-MTP query rows of the sequence, loading the + page's K (and K scales) once and reusing it across the q_len QK matmuls. + + The per-query-row kernel reloads K once per query row (q_len times per + page); at decode the QK is load-latency bound (one K load feeds one MMA), + so amortizing the K load over q_len rows lifts the load:MMA ratio. Per-row + outputs are written identically to the one-page kernel's masked path + (ASSUME_FULL_PAGES/VALID_PAGES are always False for q_len>1), so all + downstream stages are unchanged. Restricted to the perfect decode shape. + """ + seq_idx = tl.program_id(0) + head_block = tl.program_id(1) + page_rel = tl.program_id(2) + head_start = head_block * BLOCK_H + offs_h = head_start + tl.arange(0, BLOCK_H) + offs_t = tl.arange(0, BLOCK_T) + page_start = page_rel * PAGE_SIZE + + kv_len_base = tl.load(kv_lens_ptr + seq_idx) + page_table_start = tl.load(paged_kv_indptr_decode_ptr + seq_idx).to(tl.int64) + q_gscale = tl.load(q_global_scale_ptr) + + residual_groups = Q_RESIDUAL_D // FP4_BLOCK + non_residual_groups = K_HEAD_D // FP4_BLOCK - residual_groups + + # ---- r-independent index tensors (main + residual-tail column maps) ---- + scale_offsets = tl.arange(0, BLOCK_K // FP4_BLOCK) + q_sf_cols = scale_offsets + k_sf_offsets_main = _fp4_mla_swizzled_sf_offset( + offs_t[:, None], q_sf_cols[None, :], K_SF_PER_TOKEN + ) + tail_packed_offsets = tl.arange(0, TAIL_BLOCK_K // 2) + tail_scale_offsets = tl.arange(0, TAIL_BLOCK_K // FP4_BLOCK) + qt_elem = FULL_BLOCK_END + tail_packed_offsets * 2 + qt_group = qt_elem // FP4_BLOCK + kt_group = tl.where( + qt_group < non_residual_groups, + qt_group, + non_residual_groups + (qt_group - non_residual_groups) // 2, + ) + byte_t = (qt_elem % FP4_BLOCK) // 2 + packed_qt_cols = FULL_BLOCK_END // 2 + tail_packed_offsets + packed_kt_cols = kt_group * (FP4_BLOCK // 2) + byte_t + qt_sf_cols = FULL_BLOCK_END // FP4_BLOCK + tail_scale_offsets + kt_sf_cols = tl.where( + qt_sf_cols < non_residual_groups, + qt_sf_cols, + non_residual_groups + (qt_sf_cols - non_residual_groups) // 2, + ) + kt_sf_offsets = _fp4_mla_swizzled_sf_offset( + offs_t[:, None], kt_sf_cols[None, :], K_SF_PER_TOKEN + ) + scale_cols = tl.arange(0, SF_PER_PAGE) + byte_offsets = tl.arange(0, FP4_BLOCK // 2) + byte_cols = scale_cols[:, None] * (FP4_BLOCK // 2) + byte_offsets[None, :] + + tl.assume(q_fp4_s0 % 8 == 0) + tl.assume(q_fp4_s1 == 1) + tl.assume(kv_s0 % 8 == 0) + tl.assume(kv_s2 % 8 == 0) + tl.assume(kv_s4 == 1) + tl.assume(p_s0 % 8 == 0) + tl.assume(p_s1 == 1) + q_desc = tl.make_tensor_descriptor( + q_fp4_ptr, + shape=[q_num_rows, Q_HEAD_D // 2], + strides=[q_fp4_s0, q_fp4_s1], + block_shape=[BLOCK_H, BLOCK_K // 2], + ) + k_desc = tl.make_tensor_descriptor( + kv_cache_ptr, + shape=[num_pages, BLOCK_T, K_HEAD_D // 2], + strides=[kv_s0, kv_s2, kv_s4], + block_shape=[1, BLOCK_T, BLOCK_K // 2], + ) + # p_desc = tl.make_tensor_descriptor( + # p_fp4_ptr, + # shape=[p_num_rows, PAGE_SIZE // 2], + # strides=[p_s0, p_s1], + # block_shape=[BLOCK_H, PAGE_SIZE // 2], + # ) + + # ---- Load this page's K once (shared across all query rows). ---- + compact_page = page_table_start + page_rel + valid_compact_page = (compact_page >= 0) & (compact_page < page_ids_len) + safe_compact_page = tl.where(valid_compact_page, compact_page, 0) + physical_page = tl.load( + src_page_ids_ptr + safe_compact_page, mask=valid_compact_page, other=-1 + ).to(tl.int64) + valid_physical_page = valid_compact_page & (physical_page >= 0) & (physical_page < num_pages) + safe_physical_page = tl.where(valid_physical_page, physical_page, 0) + global_scale = tl.load(global_scale_ptr) + qk_scale = sm_scale / (q_gscale * global_scale) + + k_vals = k_desc.load([safe_physical_page.to(tl.int32), 0, 0]) + k_vals = tl.reshape(k_vals, (BLOCK_T, BLOCK_K // 2)) + k_vals = tl.where(valid_physical_page, k_vals, 0) + k_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + k_sf_offsets_main) + kt_vals = tl.load( + kv_cache_ptr + + safe_physical_page * kv_s0 + + offs_t[:, None].to(tl.int64) * kv_s2 + + packed_kt_cols[None, :] * kv_s4, + mask=valid_physical_page, + other=0, + ) + kt_scales = tl.load(sf_cache_ptr + safe_physical_page * sf_s0 + kt_sf_offsets) + + for r in tl.static_range(QUERY_LEN_PER_SEQ): + query_idx_r = seq_idx * QUERY_LEN_PER_SEQ + r + kv_len_r = tl.maximum(kv_len_base - (QUERY_LEN_PER_SEQ - 1 - r), 0) + q_row_base_r = query_idx_r * NUM_HEADS + q_rows_r = q_row_base_r + offs_h + page_max = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) + page_sum = tl.zeros((BLOCK_H,), dtype=tl.float32) + if page_start < kv_len_r: + q_vals = q_desc.load([(q_row_base_r + head_start).to(tl.int32), 0]) + q_sf_offsets = _fp4_mla_swizzled_sf_offset( + q_rows_r[:, None], q_sf_cols[None, :], Q_SF_PER_TOKEN + ) + q_scales = tl.load(q_sf_ptr + q_sf_offsets) + scores = tl.dot_scaled( + q_vals, + q_scales, + "e2m1", + k_vals.T, + k_scales, + "e2m1", + fast_math=True, + rhs_k_pack=True, + ) + qt_sf_offsets = _fp4_mla_swizzled_sf_offset( + q_rows_r[:, None], qt_sf_cols[None, :], Q_SF_PER_TOKEN + ) + q_tail_scales = tl.load(q_sf_ptr + qt_sf_offsets) + q_tail_vals = tl.load( + q_fp4_ptr + q_rows_r[:, None] * q_fp4_s0 + packed_qt_cols[None, :] * q_fp4_s1 + ) + scores = tl.dot_scaled( + q_tail_vals, + q_tail_scales, + "e2m1", + kt_vals.T, + kt_scales, + "e2m1", + acc=scores, + fast_math=True, + rhs_k_pack=True, + ) + + valid_t = page_start + offs_t < kv_len_r + scores = tl.where(valid_t[None, :], scores * qk_scale, -float("inf")) + page_max = tl.max(scores, axis=1) + exp_scores = tl.math.exp2((scores - page_max[:, None]) * _LOG2_E) + exp_scores = tl.where(valid_t[None, :], exp_scores, 0.0) + page_sum = tl.sum(exp_scores, axis=1) + + grouped_probs = tl.reshape(exp_scores, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK)) + amax = tl.max(grouped_probs, axis=2) + inv_local_scale = tl.where(amax > 0.0, 6.0 / amax, 1.0) + stored_scale = tl.where( + amax > 0.0, tl.minimum(amax * (P_GLOBAL_SCALE / 6.0), 448.0), 1.0 + ) + scaled_probs = grouped_probs * tl.reshape(inv_local_scale, (BLOCK_H, SF_PER_PAGE, 1)) + pairs = tl.reshape(scaled_probs, (BLOCK_H, SF_PER_PAGE, FP4_BLOCK // 2, 2)) + even_probs, odd_probs = tl.split(pairs) + packed = _fp4_e2m1_quantize_packed(even_probs, odd_probs) + + p_page = query_idx_r * MAX_PAGES + page_rel + p_rows = (p_page * NUM_HEADS + offs_h).to(tl.int64) + sf_offsets = _fp4_mla_swizzled_sf_offset( + p_rows[:, None], scale_cols[None, :], SF_PER_PAGE + ) + tl.store(p_sf_ptr + sf_offsets, stored_scale, mask=valid_compact_page) + tl.store( + p_fp4_ptr + p_rows[:, None, None] * p_s0 + byte_cols[None, :, :] * p_s1, + packed, + mask=valid_compact_page, + ) + + out_offsets = query_idx_r * page_stats_s0 + page_rel * page_stats_s1 + offs_h + tl.store(page_max_ptr + out_offsets, page_max) + tl.store(page_sum_ptr + out_offsets, page_sum) + + @triton.jit def _fp4_mla_attention_reduce_stats_kernel( max_ptr, diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py index b654dc0c2695..3527338f8bc6 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py @@ -276,7 +276,7 @@ def populate_fp4_mla_generation_lengths( def _fp4_mla_page_table_spec(kv_cache_manager: Any) -> Any: get_spec = getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None) if not callable(get_spec): - raise RuntimeError("FP4 MLA requires Fp4MlaKVCacheManagerV2 page metadata.") + raise RuntimeError("FP4 MLA requires V2 cache-layout page metadata.") spec = get_spec() for field_name in ( "cache_pool_id", @@ -386,7 +386,7 @@ def configure_fp4_mla_device_page_table( ) -> bool: """Configure the fixed-stride, device-materialized page table. - Context, generation, and fresh mixed batches receive the full block-offset + Eager context and generation batches receive the full block-offset table on the GPU. The materialization kernel decodes V2 page indices and refreshes rows from the final device KV lengths before cache update. """ @@ -413,18 +413,7 @@ def configure_fp4_mla_device_page_table( tensors = (block_offsets, page_ids, paged_kv_indptr, paged_kv_indptr_decode) is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) generation_only = num_contexts == 0 - fresh_mixed = ( - not is_cuda_graph - and num_contexts > 0 - and num_generation_sequences > 0 - and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 - ) - fresh_context_only = ( - not is_cuda_graph - and num_contexts > 0 - and num_generation_sequences == 0 - and int(getattr(metadata, "num_ctx_cached_tokens", 0) or 0) == 0 - ) + eager_context = not is_cuda_graph and num_contexts > 0 has_valid_generation = num_generation_sequences == 0 or ( num_generation_tokens >= num_generation_sequences and num_generation_tokens % num_generation_sequences == 0 @@ -432,7 +421,7 @@ def configure_fp4_mla_device_page_table( # NVFP4 exposes one data pool plus its paired block-scale pool. The # materializer reads encoded data offsets from pool 0. supported = ( - (generation_only or fresh_mixed or fresh_context_only) + (generation_only or eager_context) and kv_cache_manager is not None and has_valid_generation and int(getattr(metadata, "beam_width", 1)) == 1 @@ -470,7 +459,7 @@ def configure_fp4_mla_device_page_table( and kv_lens.ndim == 1 and kv_lens.numel() >= num_sequences ) - if fresh_mixed and not host_kv_lens_available: + if eager_context and num_generation_sequences > 0 and not host_kv_lens_available: return False if not is_cuda_graph and host_kv_lens_available: generation_tokens_per_sequence = ( diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py index db4aefb823b9..49fc36d9cc6e 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py @@ -86,18 +86,19 @@ def create(cls, metadata: TrtllmAttentionMetadata, buffers: Buffers | None) -> F f"{HP_BLOCK_SIZE}-token quantization tile, got {hp_ring_size} slots." ) hp_pool = manager.get_fp4_mla_hp_pool() + fp4_layers = manager._fp4_mla_compact_to_local if ( not isinstance(hp_pool, torch.Tensor) or hp_pool.dtype != torch.bfloat16 or hp_pool.device.type != "cuda" or hp_pool.ndim != 4 - or hp_pool.shape[1] != manager.num_local_layers + or hp_pool.shape[1] != len(fp4_layers) or hp_pool.shape[2] != manager.kv_factor - or hp_pool.shape[3] != hp_ring_size * manager.head_dim + or hp_pool.shape[3] != hp_ring_size * manager.head_dim_per_layer[fp4_layers[0]] ): raise RuntimeError( "FP4 MLA V2 HP pool must be a CUDA BF16 tensor shaped " - "[pages, local_layers, kv_factor, ring * head_dim], got " + "[pages, fp4_layers, kv_factor, ring * head_dim], got " f"{getattr(hp_pool, 'shape', None)}." ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py index d7408491b862..a377346dab9e 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py @@ -9,14 +9,23 @@ import triton import triton.language as tl +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.bindings import DataType + from .config import ( + _FP4_MLA_CUTEDSL_BACKEND, + FP4_BLOCK_SIZE, + FP4_MLA_SCALE_ROW_GROUP, FP4_MLA_TOKENS_PER_BLOCK, + HP_BLOCK_SIZE, _ceil_div, _env_enabled_default, _env_int, _fp4_mla_attention_backend, + _fp4_mla_cutedsl_fused_v_transpose_enabled, ) -from .layout import _ensure_workspace_tensor +from .fp4_mla_kernels import _fp4_mla_rebuild_v_scale_from_k_scale_kernel +from .layout import _ensure_workspace_tensor, get_fp4_mla_v_scale_pool_size def _shared_v_pack_storage_enabled() -> bool: @@ -344,7 +353,7 @@ def _update_triton_v_packed_cache( return None if page_ids.numel() == 0: return None - from .fp4_mla_triton import fp4_mla_repack_v_cache_triton + from .fp4_mla_triton import fp4_mla_repack_v_cache def _tma_alloc(size: int, alignment: int, stream): return torch.empty(size, device=kv_cache.device, dtype=torch.int8) @@ -358,7 +367,7 @@ def _tma_alloc(size: int, alignment: int, stream): dtype=torch.uint8, device=kv_cache.device, ) - fp4_mla_repack_v_cache_triton( + fp4_mla_repack_v_cache( v_packed, kv_cache, page_ids, @@ -407,3 +416,244 @@ def _maybe_update_triton_v_packed_cache( v_sf=v_sf, num_valid_pages=num_valid_pages, ) + + +def _rebuild_fp4_mla_v_scales_from_k_scales( + sf_cache: torch.Tensor, + v_scale_pool: torch.Tensor, + page_ids: torch.Tensor, + page_valid_tokens: torch.Tensor, + *, + local_layer: int, + v_head_dim: int, + page_size: int, +) -> None: + """Bit-exactly rebuild imported MLA V scales from transferred K scales.""" + if page_ids.numel() == 0: + return + if page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + "FP4 MLA imported V-scale rebuild requires " + f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." + ) + for name, tensor in ( + ("sf_cache", sf_cache), + ("v_scale_pool", v_scale_pool), + ("page_ids", page_ids), + ("page_valid_tokens", page_valid_tokens), + ): + if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor.") + if page_ids.dtype != torch.int32 or page_valid_tokens.dtype != torch.int32: + raise TypeError("FP4 MLA imported page IDs and valid-token counts must use int32.") + if page_ids.ndim != 1 or page_valid_tokens.ndim != 1: + raise ValueError("FP4 MLA imported page metadata must be one-dimensional.") + if page_ids.numel() != page_valid_tokens.numel(): + raise ValueError("FP4 MLA imported page IDs and valid-token counts must have equal length.") + if not page_ids.is_contiguous() or not page_valid_tokens.is_contiguous(): + raise ValueError("FP4 MLA imported page metadata must be contiguous.") + if not (sf_cache.device == v_scale_pool.device == page_ids.device == page_valid_tokens.device): + raise ValueError("FP4 MLA imported cache tensors must be on the same device.") + + k_sf_bytes = sf_cache.view(torch.uint8) + v_sf_bytes = v_scale_pool.view(torch.uint8) + if k_sf_bytes.ndim < 2 or v_sf_bytes.ndim < 3: + raise ValueError( + "FP4 MLA imported scale pools require per-page K storage and " + "per-layer/per-page V storage." + ) + num_layers = int(v_sf_bytes.shape[0]) + num_pages = int(v_sf_bytes.shape[1]) + if not 0 <= local_layer < num_layers: + raise IndexError( + f"local_layer={local_layer} is outside the V-scale pool with {num_layers} layers." + ) + if int(k_sf_bytes.shape[0]) != num_pages: + raise ValueError( + "FP4 MLA K/V scale pools disagree on their physical page count: " + f"{int(k_sf_bytes.shape[0])} != {num_pages}." + ) + sf_per_token = int(k_sf_bytes.shape[-1]) + required_sf_per_token = _ceil_div(v_head_dim, FP4_BLOCK_SIZE) + if sf_per_token < required_sf_per_token: + raise ValueError( + "FP4 MLA K-scale storage is too narrow for the compressed V head: " + f"{sf_per_token} < {required_sf_per_token}." + ) + required_v_page_elems = get_fp4_mla_v_scale_pool_size(v_head_dim, page_size) + if int(v_sf_bytes.shape[-1]) < required_v_page_elems: + raise ValueError( + "FP4 MLA V-scale page storage is too small for import rebuild: " + f"{int(v_sf_bytes.shape[-1])} < {required_v_page_elems}." + ) + + token_groups = page_size // HP_BLOCK_SIZE + _fp4_mla_rebuild_v_scale_from_k_scale_kernel[ + (page_ids.numel(), triton.cdiv(v_head_dim, FP4_BLOCK_SIZE)) + ]( + k_sf_bytes, + v_sf_bytes, + page_ids, + page_valid_tokens, + page_ids.numel(), + num_pages, + num_layers, + local_layer, + page_size, + k_sf_bytes.stride(0), + v_sf_bytes.stride(0), + v_sf_bytes.stride(1), + V_HEAD_D=v_head_dim, + HP_BLOCK=HP_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=token_groups, + BLOCK_TOKEN_GROUPS=triton.next_power_of_2(token_groups), + num_warps=4, + ) + + +def _stage_fp4_mla_import_page_metadata( + prompt_block_ids: list[int], + *, + prompt_len: int, + page_size: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Stage imported page metadata without synchronizing the CUDA stream.""" + if device.type != "cuda": + raise ValueError("FP4 MLA disaggregated import requires a CUDA device.") + num_prompt_pages = len(prompt_block_ids) + pin_memory = prefer_pinned() + page_ids_host = torch.tensor( + prompt_block_ids, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + page_valid_tokens_host = torch.full( + (num_prompt_pages,), + page_size, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + page_valid_tokens_host[-1] = prompt_len - (num_prompt_pages - 1) * page_size + # Constructing a CUDA tensor directly from a Python list synchronizes the + # current stream. During overlap scheduling that stream already waits for + # the previous forward, exposing the import rebuild as an inter-step gap. + # Explicit non-blocking copies keep the CPU free to enqueue every rebuild + # and the next forward while preserving their existing stream order. + page_ids = torch.empty_like(page_ids_host, device=device) + page_valid_tokens = torch.empty_like(page_valid_tokens_host, device=device) + page_ids.copy_(page_ids_host, non_blocking=True) + page_valid_tokens.copy_(page_valid_tokens_host, non_blocking=True) + return page_ids, page_valid_tokens + + +def rebuild_fp4_mla_disagg_imported_cache( + kv_cache_manager: Any, + request_id: int, + prompt_len: int, +) -> bool: + """Rebuild GEN-local FP4 MLA sidecars after a disaggregated KV import. + + The disaggregated payload carries the native V2 K, K-scale, and BF16 HP + roles. V scales and CuTeDSL's V-packed layout are deterministic + process-local views, so rebuilding them here avoids transfer bandwidth and + guarantees they are ready before a first decode step that may execute + through a pre-captured CUDA graph. + """ + if ( + kv_cache_manager is None + or getattr(kv_cache_manager, "dtype", None) != DataType.NVFP4 + or getattr(kv_cache_manager, "kv_factor", None) != 1 + or getattr(kv_cache_manager, "mla_v_scale_head_dim", None) is None + or not callable(getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None)) + ): + return False + if not isinstance(prompt_len, int) or prompt_len <= 0: + raise ValueError( + f"FP4 MLA disaggregated import needs a positive prompt_len, got {prompt_len}." + ) + + page_size = int(kv_cache_manager.tokens_per_block) + if page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + "FP4 MLA disaggregated import requires " + f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." + ) + pp_layers = list(getattr(kv_cache_manager, "pp_layers", ())) + num_local_layers = int(getattr(kv_cache_manager, "num_local_layers", len(pp_layers))) + if len(pp_layers) != num_local_layers: + raise RuntimeError( + "FP4 MLA disaggregated import cannot map local to global layers: " + f"{len(pp_layers)} PP layers for {num_local_layers} local layers." + ) + fp4_local_layers = list( + getattr(kv_cache_manager, "_fp4_mla_compact_to_local", range(num_local_layers)) + ) + if not fp4_local_layers: + raise RuntimeError("FP4 MLA disaggregated import found no local MLA layers.") + if any(local_layer < 0 or local_layer >= num_local_layers for local_layer in fp4_local_layers): + raise RuntimeError( + "FP4 MLA disaggregated import has invalid compact-to-local layer mapping: " + f"{fp4_local_layers}." + ) + + num_prompt_pages = _ceil_div(prompt_len, page_size) + first_attention_layer = pp_layers[fp4_local_layers[0]] + block_ids_per_seq = kv_cache_manager.get_batch_cache_indices( + [int(request_id)], layer_idx=first_attention_layer + ) + if len(block_ids_per_seq) != 1 or len(block_ids_per_seq[0]) < num_prompt_pages: + available = len(block_ids_per_seq[0]) if block_ids_per_seq else 0 + raise RuntimeError( + "FP4 MLA disaggregated import is missing prompt pages for request " + f"{request_id}: need {num_prompt_pages}, have {available}." + ) + prompt_block_ids = [int(block_id) for block_id in block_ids_per_seq[0][:num_prompt_pages]] + + v_scale_pool = kv_cache_manager.get_mla_v_scale_pool() + if not isinstance(v_scale_pool, torch.Tensor): + raise RuntimeError("FP4 MLA disaggregated import requires the manager V-scale pool.") + page_ids, page_valid_tokens = _stage_fp4_mla_import_page_metadata( + prompt_block_ids, + prompt_len=prompt_len, + page_size=page_size, + device=v_scale_pool.device, + ) + + v_scale_head_dim = int(kv_cache_manager.mla_v_scale_head_dim) + cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND + for compact_layer, local_layer in enumerate(fp4_local_layers): + layer_idx = pp_layers[local_layer] + kv_cache, sf_cache = kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) + _rebuild_fp4_mla_v_scales_from_k_scales( + sf_cache, + v_scale_pool, + page_ids, + page_valid_tokens, + local_layer=compact_layer, + v_head_dim=v_scale_head_dim, + page_size=page_size, + ) + if cutedsl_backend and not _fp4_mla_cutedsl_fused_v_transpose_enabled(): + v_head_dim = getattr(kv_cache_manager, "mla_v_head_dim", None) + if v_head_dim is None: + raise RuntimeError( + "CuTeDSL FP4 MLA disaggregated import requires a persistent V head dimension." + ) + v_packed = kv_cache_manager.get_mla_v_packed_pool(compact_layer) + if not isinstance(v_packed, torch.Tensor): + raise RuntimeError( + "CuTeDSL FP4 MLA disaggregated import requires the persistent V-packed pool." + ) + _repack_cutedsl_v_packed_cache( + v_packed, + kv_cache, + page_ids, + v_head_dim=int(v_head_dim), + page_size=page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + ) + return True diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 8b413cf23af0..4ccb63c37581 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -255,10 +255,6 @@ def effective_beam_width(self) -> int: _mla_ctx_cu_seqlens_valid: bool = field(default=False, init=False, repr=False) - _fp4_mla_fp8_context_state: Optional[Tuple[Any, Any]] = field(init=False, - default=None, - repr=False, - compare=False) # `DSAtrtllmAttentionMetadata` overrides this; the dense path keeps 0. num_sparse_topk: int = 0 @@ -2742,7 +2738,7 @@ def _fp4_mla_rope_generation( self.layer_idx, token_offset=getattr(metadata, "num_ctx_tokens", 0), phase="generation", - local_layer=self.get_local_layer_idx(metadata), + local_layer=self.get_fp4_mla_local_layer_idx(metadata), v_head_dim=self.kv_lora_rank, rotary_cos_sin=self.rotary_cos_sin, q_pe=q_pe, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 2b47f3f5982d..cb8d99a9b8ea 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2675,6 +2675,22 @@ def _copy_batch_block_offsets_per_layer( non_blocking=True, ) + def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: + """Return the primary cache roles allocated for one local layer.""" + roles = [Role.KEY] + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + roles.append(Role.VALUE) + if self.dtype == DataType.NVFP4: + head_dim = self.head_dim_per_layer[local_layer_idx] + assert head_dim % 2 == 0, ( + f"head_dim must be divisible by 2 for nvfp4 kv cache, " + f"but layer {local_layer_idx} has head_dim={head_dim}" + ) + roles.append(Role.KEY_BLOCK_SCALE) + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + roles.append(Role.VALUE_BLOCK_SCALE) + return roles + def _build_base_config( self, kv_cache_config: KvCacheConfig, @@ -2770,18 +2786,6 @@ def _build_base_config( ) ) - buffer_type = [Role.KEY] - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: - buffer_type.append(Role.VALUE) - if self.dtype == DataType.NVFP4: - for layer_idx, hd in enumerate(self.head_dim_per_layer): - assert hd % 2 == 0, ( - f"head_dim must be divisible by 2 for nvfp4 kv cache, but layer {layer_idx} has head_dim={hd}" - ) - buffer_type.append(Role.KEY_BLOCK_SCALE) - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: - buffer_type.append(Role.VALUE_BLOCK_SCALE) - # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional # per-layer BufferConfig entries — for example a sparse index-K # buffer — without overriding the K/V/NVFP4 scale wiring above. @@ -2795,6 +2799,7 @@ def _build_base_config( layer_configs: List[AttentionLayerConfig] = [] for layer_id in typed_range(LayerId(self.num_local_layers)): + buffer_type = self._get_buffer_roles_for_layer(layer_id) buffers = [ BufferConfig( role=role, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 295e8f30d89b..3d4984140c1b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -54,7 +54,8 @@ maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope -from .config_utils import (is_hybrid_linear, is_mla, resolve_auto_ssm_cache_dtype, +from .config_utils import (is_hybrid_linear, is_mla, + resolve_auto_ssm_cache_dtype, supports_fp4_mla_attention, uses_fp4_mla_attention, validate_kimi_kda_state_dtype) diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index fe6dde0fdf67..ae9f332dd670 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -12,6 +12,8 @@ import tensorrt_llm._torch.attention.backends.fp4_mla as fp4_mla_backend import tensorrt_llm._torch.attention.backends.fp4_mla.cache_update as fp4_mla_cache_update import tensorrt_llm._torch.attention.backends.fp4_mla.metadata as fp4_mla_metadata +import tensorrt_llm._torch.attention.backends.fp4_mla.v_cache as fp4_mla_v_cache +from tensorrt_llm._torch.attention.backends.fmha.fp4_mla import Fp4MlaFmha from tensorrt_llm._torch.attention.backends.fp4_mla import ( FP4_BLOCK_SIZE, FP4_MLA_ATTENTION_BACKEND_ENV, @@ -26,11 +28,20 @@ _cutedsl_backend_available, _fp4_mla_attention_backend, _get_fp4_mla_global_scale, + load_fp4_mla_chunked_kv_cache, run_fp4_mla_attention_decode, scatter_fp4_mla_kv_cache, ) from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 +from tensorrt_llm._torch.attention.backends.fp4_mla.fp4_mla_context import ( + _build_fp8_mla_context_metadata, +) from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, +) +from tensorrt_llm._torch.kimi_k3_cache_policy import KIMI_K3_BF16_KV_LAYERS_ENV from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, MTPDecodingConfig from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig @@ -61,6 +72,74 @@ def _is_cutedsl_unavailable() -> bool: ) +def test_fp4_mla_request_validation_uses_sparse_runtime_params() -> None: + metadata = SimpleNamespace( + num_sparse_topk=0, + kv_cache_manager=SimpleNamespace(dtype=_DataType.NVFP4, kv_factor=1), + beam_width=1, + fp4_mla_state=Fp4MlaState(hp_pool=object(), v_scale_pool=object()), + ) + forward_args = AttentionForwardArgs(attention_input_type=AttentionInputType.generation_only) + + Fp4MlaFmha._is_supported( + SimpleNamespace(attn=SimpleNamespace(sparse_params=None)), + torch.empty(1, 4), + None, + None, + metadata, + forward_args, + ) + + forward_args.sparse_runtime_params.sparse_attn_indices = torch.tensor([0]) + with pytest.raises(NotImplementedError, match="does not support sparse attention"): + Fp4MlaFmha._is_supported( + SimpleNamespace(attn=SimpleNamespace(sparse_params=None)), + torch.empty(1, 4), + None, + None, + metadata, + forward_args, + ) + + +def test_fp8_mla_context_partition_metadata_separates_q_and_kv_lengths() -> None: + prompt_lens_cuda = torch.tensor([32, 48, 1], dtype=torch.int32) + prompt_lens_cpu = prompt_lens_cuda.clone() + kv_lens_cuda = torch.tensor([128, 64], dtype=torch.int32) + kv_lens_cpu = kv_lens_cuda.clone() + meta = SimpleNamespace( + num_contexts=2, + num_ctx_tokens=80, + prompt_lens_cuda_runtime=prompt_lens_cuda, + prompt_lens_cpu_runtime=prompt_lens_cpu, + host_request_types_runtime=torch.tensor([0, 0, 1], dtype=torch.int32), + fp4_mla_state=Fp4MlaState( + positions=torch.arange(80, dtype=torch.int32), fp8_context_state=object() + ), + ) + scratch = SimpleNamespace( + cache_manager_view=object(), + block_offsets=object(), + block_ids_per_seq=object(), + host_total_kv_lens=torch.tensor([192, 0], dtype=torch.int64), + ) + + fp8_meta = _build_fp8_mla_context_metadata( + meta, + scratch, + kv_lens_cuda=kv_lens_cuda, + kv_lens_cpu=kv_lens_cpu, + ) + + assert fp8_meta.prompt_lens_cuda_runtime.data_ptr() == prompt_lens_cuda.data_ptr() + assert fp8_meta.prompt_lens_cpu_runtime.data_ptr() == prompt_lens_cpu.data_ptr() + assert fp8_meta.kv_lens_cuda_runtime is kv_lens_cuda + assert fp8_meta.kv_lens_runtime is kv_lens_cpu + assert fp8_meta.host_total_kv_lens is scratch.host_total_kv_lens + assert fp8_meta.helix_position_offsets.shape == (80,) + assert fp8_meta.fp4_mla_state is None + + def _reset_triton_allocator() -> None: import triton @@ -173,6 +252,35 @@ def test_fp4_mla_v2_cache_size_accounts_for_hp_intercept(monkeypatch) -> None: assert intercept == 3 * 2 * (HP_BLOCK_SIZE + 3) * 576 * 2 * 2 +def test_fp4_mla_v2_cache_size_accounts_for_bf16_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") + monkeypatch.setenv(KIMI_K3_BF16_KV_LAYERS_ENV, "3") + model_config = SimpleNamespace( + pretrained_config=SimpleNamespace( + model_type="kimi_linear", + num_hidden_layers=4, + kv_lora_rank=512, + qk_rope_head_dim=64, + linear_attn_config={"full_attn_layers": [2, 4]}, + ) + ) + + slope, intercept = Fp4MlaKVCacheManagerV2.get_cache_size_per_token( + model_config, + Mapping(world_size=1, tp_size=1, pp_size=1, rank=0), + num_layers=2, + tokens_per_block=FP4_MLA_TOKENS_PER_BLOCK, + max_batch_size=3, + spec_config=SimpleNamespace(tokens_per_gen_step=4), + ) + + fp4_bytes_per_token = 320 + 40 + 32 + assert slope == fp4_bytes_per_token + 576 * torch.bfloat16.itemsize + assert intercept == 3 * (HP_BLOCK_SIZE + 3) * 576 * 2 + + def test_fp4_mla_v2_runtime_sizing_accounts_for_pipeline_slots() -> None: manager = object.__new__(Fp4MlaKVCacheManagerV2) manager.max_batch_size = 3 @@ -229,6 +337,201 @@ def capture_base_init(manager, *args, **kwargs) -> None: assert captured["manager"].mla_v_head_dim == expected_v_head_dim +@pytest.mark.parametrize( + "fused_v_transpose", + [False, True], + ids=["mufu16-packed-v", "fused-v-canonical-cache"], +) +def test_fp4_mla_disagg_import_rebuilds_variant_sidecars( + monkeypatch, + fused_v_transpose: bool, +) -> None: + page_ids = torch.tensor([17, 29], dtype=torch.int32) + page_valid_tokens = torch.tensor([FP4_MLA_TOKENS_PER_BLOCK, 1], dtype=torch.int32) + v_scale_pool = torch.empty((2, 32, 1), dtype=torch.uint8) + kv_caches = { + 7: torch.empty(1, dtype=torch.uint8), + 11: torch.empty(1, dtype=torch.uint8), + } + sf_caches = { + 7: torch.empty(1, dtype=torch.uint8), + 11: torch.empty(1, dtype=torch.uint8), + } + v_packed_pools = { + 0: torch.empty(1, dtype=torch.uint8), + 1: torch.empty(1, dtype=torch.uint8), + } + staged = [] + scale_rebuilds = [] + packed_rebuilds = [] + + def stage_page_metadata(block_ids, *, prompt_len, page_size, device): + staged.append((block_ids, prompt_len, page_size, device)) + return page_ids, page_valid_tokens + + def rebuild_v_scales( + sf_cache, + observed_v_scale_pool, + observed_page_ids, + observed_page_valid_tokens, + **kwargs, + ) -> None: + scale_rebuilds.append( + ( + sf_cache, + observed_v_scale_pool, + observed_page_ids, + observed_page_valid_tokens, + kwargs, + ) + ) + + def rebuild_packed_v( + v_packed, + kv_cache, + observed_page_ids, + **kwargs, + ) -> None: + packed_rebuilds.append((v_packed, kv_cache, observed_page_ids, kwargs)) + + def get_v_packed_pool(local_layer: int) -> torch.Tensor: + if fused_v_transpose: + raise AssertionError("fused-V import must not request a packed-V pool") + return v_packed_pools[local_layer] + + manager = SimpleNamespace( + dtype=_DataType.NVFP4, + kv_factor=1, + mla_v_scale_head_dim=512, + mla_v_head_dim=None if fused_v_transpose else 512, + tokens_per_block=FP4_MLA_TOKENS_PER_BLOCK, + pp_layers=[7, 11], + num_local_layers=2, + get_fp4_mla_page_table_spec=lambda *_: SimpleNamespace(), + get_batch_cache_indices=lambda request_ids, layer_idx=None: [[17, 29, 31]], + get_mla_v_scale_pool=lambda: v_scale_pool, + get_fp4_mla_cache_buffers=lambda layer_idx: ( + kv_caches[layer_idx], + sf_caches[layer_idx], + ), + get_mla_v_packed_pool=get_v_packed_pool, + ) + monkeypatch.setattr(fp4_mla_v_cache, "_fp4_mla_attention_backend", lambda: "cutedsl") + monkeypatch.setattr( + fp4_mla_v_cache, + "_fp4_mla_cutedsl_fused_v_transpose_enabled", + lambda: fused_v_transpose, + ) + monkeypatch.setattr( + fp4_mla_v_cache, + "_stage_fp4_mla_import_page_metadata", + stage_page_metadata, + ) + monkeypatch.setattr( + fp4_mla_v_cache, + "_rebuild_fp4_mla_v_scales_from_k_scales", + rebuild_v_scales, + ) + monkeypatch.setattr( + fp4_mla_v_cache, + "_repack_cutedsl_v_packed_cache", + rebuild_packed_v, + ) + + assert fp4_mla_backend.rebuild_fp4_mla_disagg_imported_cache( + manager, + request_id=41, + prompt_len=FP4_MLA_TOKENS_PER_BLOCK + 1, + ) + + assert staged == [ + ( + [17, 29], + FP4_MLA_TOKENS_PER_BLOCK + 1, + FP4_MLA_TOKENS_PER_BLOCK, + v_scale_pool.device, + ) + ] + assert len(scale_rebuilds) == 2 + for local_layer, layer_idx in enumerate(manager.pp_layers): + sf_cache, observed_pool, observed_ids, observed_valid, kwargs = scale_rebuilds[local_layer] + assert sf_cache is sf_caches[layer_idx] + assert observed_pool is v_scale_pool + assert observed_ids is page_ids + assert observed_valid is page_valid_tokens + assert kwargs == { + "local_layer": local_layer, + "v_head_dim": 512, + "page_size": FP4_MLA_TOKENS_PER_BLOCK, + } + assert len(packed_rebuilds) == (0 if fused_v_transpose else 2) + for local_layer, layer_idx in enumerate(manager.pp_layers[: len(packed_rebuilds)]): + v_packed, kv_cache, observed_ids, kwargs = packed_rebuilds[local_layer] + assert v_packed is v_packed_pools[local_layer] + assert kv_cache is kv_caches[layer_idx] + assert observed_ids is page_ids + assert kwargs == { + "v_head_dim": 512, + "page_size": FP4_MLA_TOKENS_PER_BLOCK, + "block_v": fp4_mla_backend.FP4_MLA_SCALE_ROW_GROUP, + } + + +def test_fp4_mla_disagg_import_skips_hybrid_linear_attention_layers(monkeypatch) -> None: + page_ids = torch.tensor([17, 29], dtype=torch.int32) + page_valid_tokens = torch.tensor([FP4_MLA_TOKENS_PER_BLOCK, 1], dtype=torch.int32) + v_scale_pool = torch.empty((2, 32, 1), dtype=torch.uint8) + requested_page_layers = [] + rebuilt_cache_layers = [] + rebuilt_compact_layers = [] + + def get_batch_cache_indices(request_ids, *, layer_idx): + requested_page_layers.append((request_ids, layer_idx)) + return [[17, 29]] + + def get_fp4_mla_cache_buffers(layer_idx): + rebuilt_cache_layers.append(layer_idx) + return torch.empty(1, dtype=torch.uint8), torch.empty(1, dtype=torch.uint8) + + def rebuild_v_scales(_sf_cache, _v_scale_pool, _page_ids, _page_valid_tokens, **kwargs): + rebuilt_compact_layers.append(kwargs["local_layer"]) + + manager = SimpleNamespace( + dtype=_DataType.NVFP4, + kv_factor=1, + mla_v_scale_head_dim=512, + tokens_per_block=FP4_MLA_TOKENS_PER_BLOCK, + pp_layers=[3, 7, 11], + num_local_layers=3, + _fp4_mla_compact_to_local=[1, 2], + get_fp4_mla_page_table_spec=lambda *_: SimpleNamespace(), + get_batch_cache_indices=get_batch_cache_indices, + get_mla_v_scale_pool=lambda: v_scale_pool, + get_fp4_mla_cache_buffers=get_fp4_mla_cache_buffers, + ) + monkeypatch.setattr(fp4_mla_v_cache, "_fp4_mla_attention_backend", lambda: "triton") + monkeypatch.setattr( + fp4_mla_v_cache, + "_stage_fp4_mla_import_page_metadata", + lambda *args, **kwargs: (page_ids, page_valid_tokens), + ) + monkeypatch.setattr( + fp4_mla_v_cache, + "_rebuild_fp4_mla_v_scales_from_k_scales", + rebuild_v_scales, + ) + + assert fp4_mla_backend.rebuild_fp4_mla_disagg_imported_cache( + manager, + request_id=41, + prompt_len=FP4_MLA_TOKENS_PER_BLOCK + 1, + ) + + assert requested_page_layers == [([41], 7)] + assert rebuilt_cache_layers == [7, 11] + assert rebuilt_compact_layers == [0, 1] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_fp4_mla_manager_v2_registers_native_cache_and_hp_roles(monkeypatch) -> None: monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") @@ -559,6 +862,101 @@ def _materialize_reference_cache_tokens( return torch.stack(tokens, dim=0) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="requires Rubin SM107", +) +def test_fp4_mla_chunked_cache_gather_dequantizes_k_residual(monkeypatch) -> None: + _reset_triton_allocator() + monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") + seq_lens = [300, 180] + kv_lora_rank = 512 + qk_rope_head_dim = 64 + head_dim = kv_lora_rank + qk_rope_head_dim + kv_cache_manager = _create_fp4_mla_v2_manager( + max_tokens=640, + max_seq_len=max(seq_lens), + max_batch_size=len(seq_lens), + ) + try: + kv_cache_manager.add_dummy_requests(list(range(len(seq_lens))), seq_lens) + metadata = _build_multi_seq_metadata( + kv_cache_manager, + seq_lens=seq_lens, + page_size=FP4_MLA_TOKENS_PER_BLOCK, + ) + latent = ( + torch.randn(sum(seq_lens), head_dim, dtype=torch.bfloat16, device="cuda") * 0.25 + ).clamp_(-1.0, 1.0) + scatter_fp4_mla_kv_cache( + metadata, + latent, + layer_idx=0, + token_offset=0, + phase="context", + local_layer=0, + v_head_dim=kv_lora_rank, + ) + + chunk_lens = [96, 80] + chunk_offsets = torch.tensor([128, 64], dtype=torch.int64, device="cuda") + cu_chunk_lens = torch.tensor([0, 96, 176], dtype=torch.int64, device="cuda") + gathered_compressed_kv, gathered_k_pe = load_fp4_mla_chunked_kv_cache( + metadata, + layer_idx=0, + num_ctx_cached_tokens=sum(chunk_lens), + cu_chunked_seq_len=cu_chunk_lens, + chunked_global_offset=chunk_offsets, + chunked_max_seq_len=max(chunk_lens), + out_dtype=torch.bfloat16, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + ) + batch_indices = torch.tensor( + [0] * chunk_lens[0] + [1] * chunk_lens[1], + dtype=torch.int32, + device="cuda", + ) + positions = torch.cat( + [ + torch.arange( + offset, + offset + length, + dtype=torch.int32, + device="cuda", + ) + for offset, length in zip(chunk_offsets.tolist(), chunk_lens) + ] + ) + reference = _materialize_reference_cache_tokens( + metadata, + layer_idx=0, + batch_indices=batch_indices, + positions=positions, + head_dim=head_dim, + include_k_residual=True, + ).to(torch.bfloat16) + + torch.testing.assert_close( + gathered_compressed_kv, + reference[:, :kv_lora_rank], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + gathered_k_pe, + reference[:, kv_lora_rank:], + rtol=0, + atol=0, + ) + finally: + torch.cuda.synchronize() + _reset_triton_allocator() + kv_cache_manager.shutdown() + torch.cuda.synchronize() + torch.cuda.empty_cache() + + def _build_fp4_mla_attention_decode_case( *, seq_lens: list[int], @@ -1102,6 +1500,126 @@ def test_fp4_mla_context_tail_uses_draft_slack_ring( torch.cuda.empty_cache() +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="requires Rubin SM107", +) +def test_fp4_mla_context_scatter_fuses_qk_rope(monkeypatch) -> None: + _reset_triton_allocator() + monkeypatch.setenv(FP4_MLA_ATTENTION_BACKEND_ENV, "triton") + seq_len = 17 + kv_lora_rank = 512 + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + num_heads = 4 + head_dim = kv_lora_rank + qk_rope_head_dim + kv_cache_manager = _create_fp4_mla_v2_manager( + max_tokens=FP4_MLA_TOKENS_PER_BLOCK, + max_seq_len=FP4_MLA_TOKENS_PER_BLOCK, + max_batch_size=1, + ) + try: + kv_cache_manager.add_dummy_requests([0], [seq_len]) + metadata = _build_multi_seq_metadata( + kv_cache_manager, + seq_lens=[seq_len], + page_size=FP4_MLA_TOKENS_PER_BLOCK, + ) + latent = torch.randn( + seq_len, + head_dim, + dtype=torch.bfloat16, + device="cuda", + ) + q = torch.randn( + seq_len, + num_heads * (qk_nope_head_dim + qk_rope_head_dim), + dtype=torch.bfloat16, + device="cuda", + ) + original_latent = latent.clone() + original_q = q.clone().view( + seq_len, + num_heads, + qk_nope_head_dim + qk_rope_head_dim, + ) + rotary_cos_sin = torch.zeros( + (FP4_MLA_TOKENS_PER_BLOCK, qk_rope_head_dim, 2), + dtype=torch.float32, + device="cuda", + ) + rotary_cos_sin[..., 1] = 1.0 + + scatter_fp4_mla_kv_cache( + metadata, + latent, + layer_idx=0, + token_offset=0, + phase="context", + local_layer=0, + v_head_dim=kv_lora_rank, + rotary_cos_sin=rotary_cos_sin, + q_context=q, + q_nope_head_dim=qk_nope_head_dim, + ) + torch.cuda.synchronize() + + expected_k_pe = original_latent[:, kv_lora_rank:].reshape( + seq_len, + qk_rope_head_dim // 2, + 2, + ) + expected_k_pe = torch.stack( + (-expected_k_pe[..., 1], expected_k_pe[..., 0]), + dim=-1, + ).flatten(1) + q_view = q.view( + seq_len, + num_heads, + qk_nope_head_dim + qk_rope_head_dim, + ) + expected_q_pe = original_q[..., qk_nope_head_dim:].reshape( + seq_len, + num_heads, + qk_rope_head_dim // 2, + 2, + ) + expected_q_pe = torch.stack( + (-expected_q_pe[..., 1], expected_q_pe[..., 0]), + dim=-1, + ).flatten(-2) + torch.testing.assert_close( + latent[:, :kv_lora_rank], + original_latent[:, :kv_lora_rank], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + latent[:, kv_lora_rank:], + expected_k_pe, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + q_view[..., :qk_nope_head_dim], + original_q[..., :qk_nope_head_dim], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + q_view[..., qk_nope_head_dim:], + expected_q_pe, + rtol=0, + atol=0, + ) + finally: + torch.cuda.synchronize() + _reset_triton_allocator() + kv_cache_manager.shutdown() + torch.cuda.synchronize() + torch.cuda.empty_cache() + + _V_REPACK_PAGE_SIZE = 128 _V_REPACK_HEAD_DIM = 512 _V_REPACK_PACKED_DIM = _V_REPACK_HEAD_DIM // 2 From 80e9543775b5d326bee57f2bb5c8b17e964099a5 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:42:25 -0700 Subject: [PATCH 20/21] bench: add FP4 MLA decode microbenchmark Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- tests/microbenchmarks/bench_fp4_mla_decode.py | 1873 +++++++++++++++++ 1 file changed, 1873 insertions(+) create mode 100644 tests/microbenchmarks/bench_fp4_mla_decode.py diff --git a/tests/microbenchmarks/bench_fp4_mla_decode.py b/tests/microbenchmarks/bench_fp4_mla_decode.py new file mode 100644 index 000000000000..ad4e50a795f5 --- /dev/null +++ b/tests/microbenchmarks/bench_fp4_mla_decode.py @@ -0,0 +1,1873 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +r"""Standalone benchmark for the FP4 MLA decode kernel. + +Compares the Triton and Rubin CuTeDSL FP4 backends against the +trtllm-gen fp8 baselines. ``trtllm_fp8`` uses FlashInfer's wrapper, while +``trtllm_fp8_rubin`` uses TensorRT-LLM's native attention op. +Run with: + python tests/microbenchmarks/bench_fp4_mla_decode.py [--batch B] [--seq S] [--heads H] [--q-len Q] + +--q-len (alias --mtp-len) sets the number of query tokens per sequence (>1 for +MTP / speculative decoding); it defaults to 1 (plain decode). + +Examples (no model weights required):: + + python tests/microbenchmarks/bench_fp4_mla_decode.py --backend cutedsl --batch 16 --seq 32768 --cuda-graph + export TRTLLM_FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE=1 + python tests/microbenchmarks/bench_fp4_mla_decode.py --backend cutedsl --batch 16 \ + --seq 32768 --q-len 4 --cuda-graph --generation-step + +CuTeDSL requires Rubin SM107 and its CTM/CuTeDSL runtime. Triton requires +FP4-capable GPU kernels. ``full`` times decode (with prequantized Q), or the +scatter + fused Q/RoPE update + decode when ``--generation-step`` is set. +Byte-rate estimates are diagnostic, not measured HBM bandwidth. This is a +standalone developer benchmark, not a registered CI performance gate. +""" + +import argparse +import os +import time +from collections.abc import Callable +from types import SimpleNamespace + +import torch + +import tensorrt_llm +from tensorrt_llm._torch.attention.backends import fp4_mla as fp4_mla_backend +from tensorrt_llm._torch.attention.backends.fp4_mla import ( + FP4_BLOCK_SIZE, + FP4_MLA_ATTENTION_BACKEND_ENV, + FP4_MLA_K_RESIDUAL_DIM, + FP4_MLA_KV_GLOBAL_SCALE, + FP4_MLA_P_GLOBAL_SCALE, + FP4_MLA_Q_GLOBAL_SCALE, + FP4_MLA_Q_RESIDUAL_DIM, + FP4_MLA_TOKENS_PER_BLOCK, + _cutedsl_backend_available, + _fp4_mla_attention_backend, + _fp4_mla_cutedsl_fused_v_transpose_enabled, + _get_fp4_mla_global_scale, + run_fp4_mla_attention_decode, + scatter_fp4_mla_kv_cache, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 +from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig +from tensorrt_llm.mapping import Mapping + +_DataType = tensorrt_llm.bindings.DataType +_CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType +_BENCH_KV_GLOBAL_SCALE = FP4_MLA_KV_GLOBAL_SCALE + + +def _swizzled_sf_offset(row_idx: int, col_idx: int, sf_per_token: int) -> int: + padded_cols = ((sf_per_token + 3) // 4) * 4 + return ( + col_idx % 4 + + (col_idx // 4) * (4 * 128) + + (row_idx % 32) * 16 + + ((row_idx % 128) // 32) * 4 + + (row_idx // 128) * (128 * padded_cols) + ) + + +def _dequant_fp4_swizzled( + fp4_tensor: torch.Tensor, + sf_tensor: torch.Tensor, + *, + logical_dim: int, + sf_per_token: int, + global_scale: float, +) -> torch.Tensor: + """Decode E2M1 values and swizzled scales with tensor operations per page/batch.""" + fp4_values = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, + device=fp4_tensor.device, + ) + fp4_bytes = fp4_tensor.view(torch.uint8)[:, : logical_dim // 2] + num_rows = fp4_bytes.shape[0] + codes = torch.stack((fp4_bytes & 0x0F, fp4_bytes >> 4), dim=-1).reshape(num_rows, logical_dim) + values = fp4_values[codes.long()].reshape(num_rows, sf_per_token, FP4_BLOCK_SIZE) + + # Broadcast the scalar reference layout over all rows and scale groups. + row = torch.arange(num_rows, device=fp4_tensor.device)[:, None] + col = torch.arange(sf_per_token, device=fp4_tensor.device)[None, :] + padded_cols = ((sf_per_token + 3) // 4) * 4 + sf_offsets = ( + col % 4 + + (col // 4) * (4 * 128) + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + (row // 128) * (128 * padded_cols) + ) + sf_flat = sf_tensor.view(torch.float8_e4m3fn).reshape(-1).float() + scales = sf_flat[sf_offsets] + return (values * scales[..., None] / global_scale).reshape(num_rows, logical_dim) + + +def _duplicate_tail_groups(tensor: torch.Tensor, residual_dim: int) -> torch.Tensor: + prefix = tensor[..., :-residual_dim] + tail = tensor[..., -residual_dim:].reshape( + *tensor.shape[:-1], residual_dim // FP4_BLOCK_SIZE, FP4_BLOCK_SIZE + ) + duplicated_tail = tail.repeat_interleave(2, dim=-2).reshape( + *tensor.shape[:-1], + residual_dim * 2, + ) + return torch.cat((prefix, duplicated_tail), dim=-1) + + +def _expand_qk_residual_terms( + q: torch.Tensor, k: torch.Tensor, k_residual: torch.Tensor, residual_dim: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Build [Q, Q_r, Q] x [K, K, K_r] after the shared prefix.""" + prefix_dim = k.shape[-1] - residual_dim + residual_groups = residual_dim // FP4_BLOCK_SIZE + q_tail = q[..., prefix_dim:].reshape(*q.shape[:-1], residual_groups, 2, FP4_BLOCK_SIZE) + q_main = q_tail[..., 0, :].reshape(*q.shape[:-1], residual_dim) + q_residual = q_tail[..., 1, :].reshape(*q.shape[:-1], residual_dim) + k_main = k[..., prefix_dim:] + return ( + torch.cat((q[..., :prefix_dim], q_main, q_residual, q_main), dim=-1), + torch.cat((k[..., :prefix_dim], k_main, k_main, k_residual), dim=-1), + ) + + +def _build_multi_seq_metadata( + kv_cache_manager: Fp4MlaKVCacheManagerV2, *, seq_lens: list[int], page_size: int, layer_idx: int +) -> SimpleNamespace: + device = torch.device("cuda") + num_seqs = len(seq_lens) + request_ids = list(range(num_seqs)) + block_ids_per_seq = kv_cache_manager.get_batch_cache_indices( + request_ids, + layer_idx=layer_idx, + ) + page_spec = kv_cache_manager.get_fp4_mla_page_table_spec(layer_idx) + hp_block_ids_per_seq = kv_cache_manager._get_batch_cache_indices_by_pool_id( + request_ids, + pool_id=page_spec.hp_pool_id, + is_kv_aggregate=False, + ) + num_blocks = [(seq_len + page_size - 1) // page_size for seq_len in seq_lens] + + max_blocks_per_seq = max(num_blocks) + page_rows = [] + hp_page_rows = [] + for seq_idx, seq_blocks in enumerate(block_ids_per_seq): + active_blocks = seq_blocks[: num_blocks[seq_idx]] + page_rows.extend(active_blocks + [0] * (max_blocks_per_seq - len(active_blocks))) + active_hp_blocks = hp_block_ids_per_seq[seq_idx][: num_blocks[seq_idx]] + hp_page_rows.extend(active_hp_blocks + [0] * (max_blocks_per_seq - len(active_hp_blocks))) + paged_kv_indices = torch.tensor( + page_rows, + dtype=torch.int32, + device=device, + ) + paged_kv_indptr = ( + torch.arange(num_seqs + 1, dtype=torch.int32, device=device) * max_blocks_per_seq + ) + hp_page_indices = torch.tensor(hp_page_rows, dtype=torch.int32, device=device) + batch_indices = torch.cat( + [ + torch.full( + (seq_len,), + seq_idx, + dtype=torch.int32, + device=device, + ) + for seq_idx, seq_len in enumerate(seq_lens) + ] + ) + positions = torch.cat( + [torch.arange(seq_len, dtype=torch.int32, device=device) for seq_len in seq_lens] + ) + + hp_pool = kv_cache_manager.get_fp4_mla_hp_pool() + kv_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + prompt_lens_cuda = torch.tensor(seq_lens, dtype=torch.int32, device=device) + prompt_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32) + kv_global_scale = torch.tensor([_BENCH_KV_GLOBAL_SCALE], dtype=torch.float32, device=device) + q_global_scale = torch.tensor([FP4_MLA_Q_GLOBAL_SCALE], dtype=torch.float32, device=device) + + return SimpleNamespace( + kv_cache_manager=kv_cache_manager, + page_size=page_size, + num_contexts=num_seqs, + num_seqs=num_seqs, + kv_lens_cuda_runtime=kv_lens, + prompt_lens_cuda_runtime=prompt_lens_cuda, + prompt_lens_cpu_runtime=prompt_lens_cpu, + request_ids=request_ids, + runtime_features=SimpleNamespace(has_speculative_draft_tokens=False), + is_cuda_graph=False, + is_warmup=False, + fp4_mla_state=Fp4MlaState( + batch_indices=batch_indices, + positions=positions, + _paged_kv_indices=paged_kv_indices, + hp_page_indices=hp_page_indices, + _paged_kv_indptr=paged_kv_indptr, + paged_kv_indptr_decode=paged_kv_indptr.clone(), + device_page_table=True, + device_page_table_valid=True, + page_table_stride=max_blocks_per_seq, + num_context_blocks=num_seqs * max_blocks_per_seq, + num_generation_blocks=0, + num_sequences=num_seqs, + num_blocks=None, + hp_pool=hp_pool, + v_scale_pool=kv_cache_manager.get_mla_v_scale_pool(), + q_global_scale=q_global_scale, + kv_global_scale=kv_global_scale, + ), + ) + + +def _materialize_reference_cache_storage( + metadata: SimpleNamespace, layer_idx: int, head_dim: int +) -> torch.Tensor: + kv_cache, sf_cache = metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) + sf_cache = sf_cache.view(torch.float8_e4m3fn) + storage_head_dim = kv_cache.shape[-1] * 2 + static_global_scale = float(_get_fp4_mla_global_scale(metadata, kv_cache.device).item()) + pages = [] + dequantized_pages = {} + page_rows = metadata.fp4_mla_state.paged_kv_indices.view( + metadata.num_seqs, + metadata.fp4_mla_state.page_table_stride, + ) + # Preserve the fixed row stride used by paged_kv_indptr_decode, including + # padding for shorter requests. The reference masks padding by KV length. + src_page_ids = page_rows.reshape(-1) + for page_id in src_page_ids.tolist(): + if page_id not in dequantized_pages: + fp4_page = kv_cache[page_id, 0, :, 0, :] + sf_page = sf_cache[page_id] + dequantized_pages[page_id] = _dequant_fp4_swizzled( + fp4_page, + sf_page, + logical_dim=storage_head_dim, + sf_per_token=storage_head_dim // FP4_BLOCK_SIZE, + global_scale=static_global_scale, + ) + pages.append(dequantized_pages[page_id]) + if not pages: + return torch.empty( + (0, metadata.page_size, storage_head_dim), + dtype=torch.float32, + device=kv_cache.device, + ) + return torch.stack(pages, dim=0) + + +def _materialize_reference_cache( + metadata: SimpleNamespace, layer_idx: int, head_dim: int +) -> torch.Tensor: + return _materialize_reference_cache_storage(metadata, layer_idx, head_dim)[..., :head_dim] + + +def _materialize_reference_k_residual( + metadata: SimpleNamespace, layer_idx: int, *, head_dim: int +) -> torch.Tensor: + storage = _materialize_reference_cache_storage(metadata, layer_idx, head_dim) + return storage[..., head_dim : head_dim + FP4_MLA_K_RESIDUAL_DIM] + + +def _build_fp4_mla_attention_decode_case( + *, + seq_lens: list[int], + num_heads: int, + seed: int, + query_len_per_seq: int = 1, + num_layers: int = 1, + layer_idx: int = 0, + local_layer: int | None = None, +) -> tuple[Fp4MlaKVCacheManagerV2, SimpleNamespace, torch.Tensor, torch.Tensor, int, int]: + local_layer = layer_idx if local_layer is None else local_layer + torch.manual_seed(seed) + device = torch.device("cuda") + + kv_lora_rank = 512 + qk_rope_head_dim = 64 + head_dim = kv_lora_rank + qk_rope_head_dim + page_size = FP4_MLA_TOKENS_PER_BLOCK + num_blocks = [(seq_len + page_size - 1) // page_size for seq_len in seq_lens] + num_pages = sum(num_blocks) + max_seq_len = max(page_size, max(seq_lens)) + max_tokens = max(page_size, num_pages * page_size) + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + spec_config = ( + MTPDecodingConfig(max_draft_len=query_len_per_seq - 1) if query_len_per_seq > 1 else None + ) + kv_cache_manager = Fp4MlaKVCacheManagerV2( + KvCacheConfig( + max_tokens=max_tokens, + dtype="nvfp4", + enable_block_reuse=False, + host_cache_size=0, + ), + _CacheType.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=head_dim, + tokens_per_block=page_size, + max_seq_len=max_seq_len, + max_batch_size=len(seq_lens), + mapping=mapping, + dtype=_DataType.NVFP4, + spec_config=spec_config, + max_num_tokens=max_tokens, + pretrained_config=SimpleNamespace(kv_lora_rank=kv_lora_rank), + ) + kv_cache_manager.add_dummy_requests(list(range(len(seq_lens))), seq_lens) + expected_storage_head_dim = head_dim + ( + FP4_MLA_K_RESIDUAL_DIM if _fp4_mla_attention_backend() in ("triton", "cutedsl") else 0 + ) + for current_layer in range(num_layers): + kv_cache, sf_cache = kv_cache_manager.get_fp4_mla_cache_buffers(current_layer) + kv_cache.zero_() + sf_cache.zero_() + assert kv_cache.shape[-1] * 2 == expected_storage_head_dim + assert sf_cache.shape[-1] == expected_storage_head_dim // FP4_BLOCK_SIZE + + metadata = _build_multi_seq_metadata( + kv_cache_manager, + seq_lens=seq_lens, + page_size=page_size, + layer_idx=layer_idx, + ) + assert metadata.fp4_mla_state.v_scale_pool is not None + persistent_pool_base = kv_cache_manager.get_mla_v_packed_pool_base() + if persistent_pool_base is not None: + persistent_pool_base.zero_() + metadata.fp4_mla_state.v_scale_pool.zero_() + + latent = ( + torch.randn(sum(seq_lens), head_dim, dtype=torch.bfloat16, device=device) * 0.25 + ).clamp_(-1.0, 1.0) + scatter_fp4_mla_kv_cache( + metadata, + latent, + layer_idx=layer_idx, + token_offset=0, + phase="context", + local_layer=local_layer, + v_head_dim=kv_lora_rank, + ) + # Decode reference exercises the quantized cache without an HP overlay. + metadata.fp4_mla_state.hp_pool.zero_() + torch.cuda.synchronize() + + metadata.num_contexts = 0 + metadata.fp4_mla_state.num_context_blocks = 0 + metadata.fp4_mla_state.num_generation_blocks = ( + len(seq_lens) * metadata.fp4_mla_state.page_table_stride + ) + metadata.prompt_lens_cuda_runtime = torch.full( + (len(seq_lens),), query_len_per_seq, dtype=torch.int32, device=device + ) + metadata.prompt_lens_cpu_runtime = torch.full( + (len(seq_lens),), query_len_per_seq, dtype=torch.int32 + ) + num_queries = len(seq_lens) * query_len_per_seq + q_nope = ( + torch.randn(num_queries, num_heads, kv_lora_rank, dtype=torch.bfloat16, device=device) + * 0.25 + ).clamp_(-1.0, 1.0) + q_pe = ( + torch.randn( + num_queries, + num_heads, + qk_rope_head_dim, + dtype=torch.bfloat16, + device=device, + ) + * 0.25 + ).clamp_(-1.0, 1.0) + + return kv_cache_manager, metadata, q_nope, q_pe, kv_lora_rank, qk_rope_head_dim + + +def _fp4_mla_attention_decode_reference( + metadata: SimpleNamespace, + q_nope: torch.Tensor, + q_pe: torch.Tensor, + *, + sm_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, + layer_idx: int = 0, + local_layer: int = 0, +) -> tuple[torch.Tensor, list[torch.Tensor], list[torch.Tensor]]: + head_dim = kv_lora_rank + qk_rope_head_dim + dequant_cache = _materialize_reference_cache(metadata, layer_idx, head_dim) + dequant_k_residual = None + if _fp4_mla_attention_backend() in ("triton", "cutedsl"): + dequant_k_residual = _materialize_reference_k_residual( + metadata, + layer_idx, + head_dim=head_dim, + ) + num_heads = q_nope.shape[1] + q_full = torch.cat((q_nope, q_pe), dim=-1).reshape(-1, head_dim) + global_scale = metadata.fp4_mla_state.q_global_scale + q_fp4, q_sf = torch.ops.trtllm.fp4_quantize_with_residual( + q_full, + global_scale, + FP4_MLA_Q_RESIDUAL_DIM, + is_act=True, + ) + q_logical_dim = head_dim + FP4_MLA_Q_RESIDUAL_DIM + q_dequant = _dequant_fp4_swizzled( + q_fp4, + q_sf.view(torch.float8_e4m3fn), + logical_dim=q_logical_dim, + sf_per_token=q_logical_dim // FP4_BLOCK_SIZE, + global_scale=float(global_scale.item()), + ) + + p_dequant = None + if "_fp4_mla_attention_p_buf" in metadata.fp4_mla_state.workspaces: + p_dequant = _dequant_fp4_swizzled( + metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_buf"], + metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_sf_buf"], + logical_dim=metadata.page_size, + sf_per_token=metadata.page_size // FP4_BLOCK_SIZE, + global_scale=FP4_MLA_P_GLOBAL_SCALE, + ) + + indptr = metadata.fp4_mla_state.paged_kv_indptr_decode.cpu().tolist() + kv_lens = metadata.kv_lens_cuda_runtime.cpu().tolist() + num_seqs = metadata.num_seqs - metadata.num_contexts + query_len_per_seq = q_nope.shape[0] // num_seqs + max_pages = max(indptr[seq_idx + 1] - indptr[seq_idx] for seq_idx in range(num_seqs)) + outputs = [] + exact_probs = [] + quantized_probs = [] + for seq_idx in range(num_seqs): + kv_len = kv_lens[seq_idx] + full_cache = dequant_cache[indptr[seq_idx] : indptr[seq_idx + 1]].reshape(-1, head_dim) + full_v_cache = full_cache[:, :kv_lora_rank] + for query_offset in range(query_len_per_seq): + query_idx = seq_idx * query_len_per_seq + query_offset + effective_kv_len = kv_len - (query_len_per_seq - 1 - query_offset) + cache = full_cache[:effective_kv_len] + v_cache = full_v_cache[:effective_kv_len] + q_start = query_idx * num_heads + q = q_dequant[q_start : q_start + num_heads] + if dequant_k_residual is None: + logical_q = q + logical_k = _duplicate_tail_groups(cache.float(), FP4_MLA_Q_RESIDUAL_DIM) + else: + full_k_residual = dequant_k_residual[indptr[seq_idx] : indptr[seq_idx + 1]].reshape( + -1, FP4_MLA_K_RESIDUAL_DIM + ) + logical_q, logical_k = _expand_qk_residual_terms( + q, + cache.float(), + full_k_residual[:effective_kv_len], + FP4_MLA_Q_RESIDUAL_DIM, + ) + probs = torch.softmax( + torch.matmul(logical_q, logical_k.transpose(0, 1)) * sm_scale, + dim=-1, + ) + + if p_dequant is None: + p = probs + else: + p_pages = [] + for page_rel in range(indptr[seq_idx + 1] - indptr[seq_idx]): + page_start = page_rel * metadata.page_size + valid_tokens = max(min(effective_kv_len - page_start, metadata.page_size), 0) + if valid_tokens == 0: + continue + p_page = query_idx * max_pages + page_rel + p_start = p_page * num_heads + p_pages.append(p_dequant[p_start : p_start + num_heads, :valid_tokens]) + p = torch.cat(p_pages, dim=-1) + + exact_probs.append(probs) + quantized_probs.append(p) + outputs.append(torch.matmul(p, v_cache.float())) + return torch.stack(outputs, dim=0), exact_probs, quantized_probs + + +BACKEND_CHOICES = ( + "trtllm_fp8", + "trtllm_fp8_rubin", + "triton", + "cutedsl", +) + + +def _bench(fn: Callable[[], object], warmup: int = 10, iters: int = 50) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def _bench_with_untimed_setup( + fn: Callable[[], object], + setup: Callable[[], object], + warmup: int = 10, + iters: int = 50, + queue_delay_cycles: int = 0, +) -> float: + """Benchmark ``fn`` after a same-stream setup excluded from event timing.""" + for _ in range(warmup): + setup() + fn() + torch.cuda.synchronize() + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + for start, end in zip(starts, ends): + setup() + if queue_delay_cycles: + torch.cuda._sleep(queue_delay_cycles) + start.record() + fn() + end.record() + torch.cuda.synchronize() + return sum(start.elapsed_time(end) for start, end in zip(starts, ends)) / iters + + +def _capture_cuda_graph( + fn: Callable[[], object], device: torch.device +) -> tuple[torch.cuda.CUDAGraph, torch.cuda.Stream]: + """Warm and capture the GPU work submitted by ``fn`` on a side stream.""" + current_stream = torch.cuda.current_stream(device) + capture_stream = torch.cuda.Stream(device=device) + capture_stream.wait_stream(current_stream) + with torch.cuda.stream(capture_stream): + fn() + current_stream.wait_stream(capture_stream) + torch.cuda.synchronize(device) + + cuda_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph, stream=capture_stream): + fn() + torch.cuda.synchronize(device) + return cuda_graph, capture_stream + + +def _bench_cuda_graph( + replay: Callable[[], object], device: torch.device, warmup: int = 10, iters: int = 50 +) -> tuple[float, float]: + """Return per-replay CUDA-event and host-wall times in milliseconds.""" + for _ in range(warmup): + replay() + torch.cuda.synchronize(device) + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + replay() + end.record() + torch.cuda.synchronize(device) + event_ms = start.elapsed_time(end) / iters + + wall_start_ns = time.perf_counter_ns() + for _ in range(iters): + replay() + torch.cuda.synchronize(device) + wall_ms = (time.perf_counter_ns() - wall_start_ns) / 1_000_000.0 / iters + return event_ms, wall_ms + + +PAGE_SIZE = FP4_MLA_TOKENS_PER_BLOCK + + +def _fp4_mla_lazy_rebase_stats(exact_probs: list[torch.Tensor]) -> SimpleNamespace: + """Classify lazy resident-O rebases from exact quantized-QK probabilities.""" + kernel = fp4_mla_backend._fp4_mla_cutedsl_kernel_module() + kv_tile = kernel.KV_TILE + rebase_threshold = kernel.SMEM_P4_LAZY_ANCHOR_REBASE_LOG2 + + rebase_counts = [] + max_anchor_delta_log2 = 0.0 + for probabilities in exact_probs: + log2_probabilities = torch.log2(probabilities.float()) + num_rows = log2_probabilities.shape[0] + true_row_max = torch.full( + (num_rows,), -torch.inf, dtype=torch.float32, device=probabilities.device + ) + row_anchor = torch.full_like(true_row_max, -torch.inf) + row_rebases = torch.zeros_like(true_row_max, dtype=torch.int32) + for tile_start in range(0, log2_probabilities.shape[1], kv_tile): + tile_row_max = ( + log2_probabilities[:, tile_start : tile_start + kv_tile].max(dim=1).values + ) + candidate_row_max = torch.maximum(true_row_max, tile_row_max) + finite_pair = torch.isfinite(row_anchor) & torch.isfinite(candidate_row_max) + first_finite = ~torch.isfinite(row_anchor) & torch.isfinite(candidate_row_max) + anchor_delta_log2 = torch.where( + finite_pair, + candidate_row_max - row_anchor, + torch.zeros_like(candidate_row_max), + ) + finite_delta = anchor_delta_log2[finite_pair] + if finite_delta.numel(): + max_anchor_delta_log2 = max(max_anchor_delta_log2, float(finite_delta.max().item())) + rebase = finite_pair & (anchor_delta_log2 > rebase_threshold) + row_rebases += rebase.to(torch.int32) + row_anchor = torch.where(first_finite | rebase, candidate_row_max, row_anchor) + true_row_max = candidate_row_max + rebase_counts.append(row_rebases.cpu()) + + if not rebase_counts: + raise ValueError("exact_probs must contain at least one query") + counts = torch.cat(rebase_counts) + return SimpleNamespace( + total_rows=int(counts.numel()), + rebase_rows=int((counts > 0).sum().item()), + rebase_total=int(counts.sum().item()), + rebase_max_per_row=int(counts.max().item()), + max_anchor_delta_log2=max_anchor_delta_log2, + threshold_log2=float(rebase_threshold), + ) + + +_CUTEDSL_LAUNCHER_TENSOR_NAMES = ( + "q_fp4", + "q_sf", + "kv_cache", + "sf_cache", + "v_packed", + "v_sf", + "global_scale", + "src_page_ids", + "paged_kv_indptr_decode", + "kv_lens", + "output", +) + + +def _launcher_tensor_pointers(tensors: tuple[object, ...]) -> dict[str, int | None]: + """Record stable launch addresses, including fused-V's absent sidecar.""" + if len(tensors) != len(_CUTEDSL_LAUNCHER_TENSOR_NAMES): + raise RuntimeError( + f"CuTeDSL launcher requires {len(_CUTEDSL_LAUNCHER_TENSOR_NAMES)} " + f"positional tensor arguments, got {len(tensors)}" + ) + pointers = {} + for name, tensor in zip(_CUTEDSL_LAUNCHER_TENSOR_NAMES, tensors): + if name == "v_packed" and tensor is None: + pointers[name] = None + elif isinstance(tensor, torch.Tensor): + pointers[name] = int(tensor.data_ptr()) + else: + raise TypeError(f"CuTeDSL launcher argument {name} must be a tensor") + return pointers + + +class RetainedBenchmarkState: + """Keep one benchmark allocation and captured Graph alive for A/B checks.""" + + def __init__( + self, + kv_cache_manager: Fp4MlaKVCacheManagerV2, + launcher_tensors: tuple[torch.Tensor | None, ...], + keepalive: dict[str, object], + ) -> None: + self._kv_cache_manager = kv_cache_manager + self._launcher_tensors = tuple(launcher_tensors) + self._keepalive = keepalive + self.launcher_tensor_ptrs = _launcher_tensor_pointers(self._launcher_tensors) + + def shutdown(self) -> None: + """Synchronize, destroy Graph references, and release the retained manager.""" + manager = self._kv_cache_manager + if manager is None: + return + output = self._keepalive.get("output") + if isinstance(output, torch.Tensor) and output.is_cuda: + torch.cuda.synchronize(output.device) + self._keepalive["cuda_graph"] = None + self._keepalive["capture_stream"] = None + self._launcher_tensors = () + self._kv_cache_manager = None + self._keepalive.clear() + manager.shutdown() + + +def _seq_lens_for_batch(batch: int, seq: int) -> list[int]: + return [seq] * batch + + +def _seq_label(seq_lens: list[int]) -> str: + return f"{seq_lens[0]}-{seq_lens[-1]}" if len(seq_lens) > 1 else str(seq_lens[0]) + + +def _causal_token_pairs(seq_lens: list[int], q_len: int) -> int: + return sum( + max(seq_len - (q_len - 1 - query_offset), 0) + for seq_len in seq_lens + for query_offset in range(q_len) + ) + + +def _kernel_io_bytes( + seq_lens: list[int], + heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + q_len: int = 1, + k_residual_dim: int = 0, +) -> int: + """Estimate logical input/output bytes; this is not measured HBM traffic.""" + batch = len(seq_lens) + token_pairs = _causal_token_pairs(seq_lens, q_len) + q_head_dim = kv_lora_rank + qk_rope_head_dim + FP4_MLA_Q_RESIDUAL_DIM + k_head_dim = kv_lora_rank + qk_rope_head_dim + k_residual_dim + q_sf_per_token = q_head_dim // FP4_BLOCK_SIZE + k_sf_per_token = k_head_dim // FP4_BLOCK_SIZE + pages = sum((seq_len + PAGE_SIZE - 1) // PAGE_SIZE for seq_len in seq_lens) + sf_per_page = PAGE_SIZE // FP4_BLOCK_SIZE + + q_fp4 = batch * q_len * heads * q_head_dim // 2 + q_sf = batch * q_len * heads * q_sf_per_token + kv_cache = token_pairs * k_head_dim // 2 + k_sf_cache = token_pairs * k_sf_per_token + v_packed = token_pairs * kv_lora_rank // 2 + v_sf_cache = q_len * pages * kv_lora_rank * sf_per_page + out = batch * q_len * heads * kv_lora_rank * 2 # bf16/half + return q_fp4 + q_sf + kv_cache + k_sf_cache + v_packed + v_sf_cache + out + + +def _fixed_tile_request_bytes( + seq_lens: list[int], + heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + q_len: int = 1, + k_residual_dim: int = 0, +) -> int: + """Estimate static global requests made by the current K256 kernel. + + This implementation-specific diagnostic rounds every query to the launch + maximum K256 tile count. It is neither logical algorithm traffic nor a + hardware DRAM counter, and therefore must not be used as a stable gate. + """ + batch = len(seq_lens) + queries = batch * q_len + q_head_dim = kv_lora_rank + qk_rope_head_dim + FP4_MLA_Q_RESIDUAL_DIM + k_head_dim = kv_lora_rank + qk_rope_head_dim + k_residual_dim + physical_k = ((max(seq_lens) + 255) // 256) * 256 + + q_fp4 = queries * heads * q_head_dim // 2 + # Both CTAs independently request the complete Q scale vector. + q_sf = 2 * queries * heads * (q_head_dim // FP4_BLOCK_SIZE) + k_fp4 = queries * physical_k * k_head_dim // 2 + k_sf = queries * physical_k * (k_head_dim // FP4_BLOCK_SIZE) + v_fp4 = queries * physical_k * kv_lora_rank // 2 + v_sf = queries * physical_k * (kv_lora_rank // FP4_BLOCK_SIZE) + out = queries * heads * kv_lora_rank * 2 + return q_fp4 + q_sf + k_fp4 + k_sf + v_fp4 + v_sf + out + + +def _tma_smem_completion_bytes( + seq_lens: list[int], + heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + q_len: int = 1, + k_residual_dim: int = 0, +) -> int: + """Estimate TMA shared-memory destination bytes for diagnostics only. + + Relative to static global requests, K and V scale multicast contributes a + destination in each CTA. This is not a global-memory or HBM byte count. + """ + global_requests = _fixed_tile_request_bytes( + seq_lens, + heads, + kv_lora_rank, + qk_rope_head_dim, + q_len, + k_residual_dim, + ) + queries = len(seq_lens) * q_len + physical_k = ((max(seq_lens) + 255) // 256) * 256 + multicast_scale_destination = ( + queries + * physical_k + * ( + (kv_lora_rank + qk_rope_head_dim + k_residual_dim) // FP4_BLOCK_SIZE + + kv_lora_rank // FP4_BLOCK_SIZE + ) + ) + return global_requests + multicast_scale_destination + + +def _trtllm_mla_io_bytes( + seq_lens: list[int], + heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + q_len: int = 1, + elem_bytes: int = 2, +) -> int: + """Estimate logical bytes for the trtllm-gen MLA decode baseline. + + ``elem_bytes`` is the byte width of the Q and KV-cache elements (2 for bf16, + 1 for fp8). The output is always written as bf16 (2 B/elem). + """ + batch = len(seq_lens) + token_pairs = _causal_token_pairs(seq_lens, q_len) + head_dim = kv_lora_rank + qk_rope_head_dim + q = batch * q_len * heads * head_dim * elem_bytes + kv = token_pairs * head_dim * elem_bytes # ckv + kpe paged caches + out = batch * q_len * heads * kv_lora_rank * 2 + return q + kv + out + + +def run_one_trtllm( + batch: int, seq: int, heads: int, q_len: int = 1, warmup: int = 0, iters: int = 1 +) -> float: + """Fp8 baseline using the FlashInfer trtllm-gen MLA decode kernel. + + Feeds fp8 (e4m3) Q and KV cache so the kernel uses fp8 tensor cores + (output stays bf16). + """ + import flashinfer + + device = torch.device("cuda") + label = "trtllm_fp8" + io_dtype = torch.float8_e4m3fn + elem_bytes = 1 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + qk_nope_head_dim = 128 # DeepSeek-V3 default; only used for fused scale convention. + head_dim_qk = kv_lora_rank + qk_rope_head_dim + page_size = 64 # trtllm-gen MLA decode only supports page_size of 32 or 64. + seq_lens_list = _seq_lens_for_batch(batch, seq) + max_seq = max(seq_lens_list) + blocks_per_seq = [(seq_len + page_size - 1) // page_size for seq_len in seq_lens_list] + max_blocks_per_seq = max(blocks_per_seq) + total_pages = sum(blocks_per_seq) + + torch.manual_seed(8) + # query layout: [batch, q_len, heads, kv_lora_rank + qk_rope_head_dim] + query = torch.randn(batch, q_len, heads, head_dim_qk, dtype=torch.bfloat16, device=device) + # kv_cache layout: [num_pages, page_size, head_dim_ckv + head_dim_kpe] + kv_cache = torch.randn(total_pages, page_size, head_dim_qk, dtype=torch.bfloat16, device=device) + # Quantize to fp8 e4m3 (randn ~ N(0,1) is well within e4m3 range). + query = query.to(io_dtype) + kv_cache = kv_cache.to(io_dtype) + + block_tables = torch.zeros((batch, max_blocks_per_seq), dtype=torch.int32, device=device) + page_start = 0 + for batch_idx, num_blocks in enumerate(blocks_per_seq): + block_tables[batch_idx, :num_blocks] = torch.arange( + page_start, page_start + num_blocks, dtype=torch.int32, device=device + ) + page_start += num_blocks + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device=device) + workspace = torch.zeros(128 * 1024 * 1024, dtype=torch.int8, device=device).view(-1, 4) + + output = torch.empty(batch, q_len, heads, kv_lora_rank, dtype=torch.bfloat16, device=device) + + # bmm1_scale folds q_scale * k_scale * sm_scale / sqrt(head_dim_qk); a + # representative value (q_scale = k_scale = 1.0 for the fp8 unit-scale tensors). + def run() -> None: + flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla( + query=query, + kv_cache=kv_cache, + workspace_buffer=workspace, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + block_tables=block_tables, + seq_lens=seq_lens, + max_seq_len=max_seq, + out=output, + bmm1_scale=0.1, + bmm2_scale=1.0, + backend="trtllm-gen", + ) + + run() + torch.cuda.synchronize() + avg_ms = _bench(run, warmup=warmup, iters=iters) + qk_dim = kv_lora_rank + qk_rope_head_dim + pv_dim = kv_lora_rank + flops = 2 * heads * _causal_token_pairs(seq_lens_list, q_len) * (qk_dim + pv_dim) + tflops = flops / avg_ms / 1e9 + bytes_per_call = _trtllm_mla_io_bytes( + seq_lens_list, heads, kv_lora_rank, qk_rope_head_dim, q_len, elem_bytes + ) + gb_s = bytes_per_call / (avg_ms * 1e-3) / 1e9 + print( + f"backend={label:>12s} bs={batch:>3d} seq={_seq_label(seq_lens_list):>11s} " + f"heads={heads:>3d} qlen={q_len:>2d}: " + f"{avg_ms:>7.3f} ms {tflops:>6.2f} TFLOP/s " + f"Est.IO {gb_s:>6.1f} GB/s", + flush=True, + ) + return avg_ms + + +def run_one_trtllm_rubin( + batch: int, + seq: int, + heads: int, + q_len: int = 1, + warmup: int = 0, + iters: int = 1, + queue_delay_cycles: int = 0, + use_cuda_graph: bool = False, + graph_timing: str = "event", +) -> float: + """FP8 baseline using TensorRT-LLM's native SM107 trtllm-gen path.""" + import tensorrt_llm + from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha + from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + RopeParams, + ) + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager + from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str + from tensorrt_llm.functional import PositionEmbeddingType + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.quantization.mode import QuantAlgo + + if q_len > seq: + raise ValueError(f"q_len ({q_len}) must not exceed seq ({seq})") + if queue_delay_cycles < 0: + raise ValueError("queue delay cycles must be non-negative") + if graph_timing not in ("event", "wall"): + raise ValueError(f"unsupported CUDA Graph timing mode: {graph_timing}") + + device = torch.device("cuda") + label = "trtllm_fp8_rubin" + io_dtype = torch.float8_e4m3fn + kv_lora_rank = 512 + qk_rope_head_dim = 64 + qk_nope_head_dim = 128 + q_lora_rank = 1536 + v_head_dim = 128 + head_dim_qk = kv_lora_rank + qk_rope_head_dim + page_size = 32 + seq_lens_list = _seq_lens_for_batch(batch, seq) + request_ids = list(range(batch)) + past_seq = seq - q_len + max_tokens = batch * ((seq + page_size - 1) // page_size) * page_size + mapping = Mapping(world_size=1, tp_size=1, rank=0) + + kv_cache_manager = KVCacheManager( + KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=head_dim_qk, + tokens_per_block=page_size, + max_seq_len=seq, + max_batch_size=batch, + mapping=mapping, + dtype=str_dtype_to_binding(torch_dtype_to_str(io_dtype)), + ) + try: + kv_cache_manager.add_dummy_requests(request_ids, [seq] * batch) + kv_cache_manager.get_buffers(0).zero_() + + metadata = TrtllmAttention.Metadata( + seq_lens=torch.full((batch,), q_len, dtype=torch.int32), + request_ids=request_ids, + max_num_requests=batch, + num_contexts=0, + prompt_lens=[past_seq] * batch, + max_num_tokens=batch * q_len, + kv_cache_manager=kv_cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[past_seq] * batch, + ), + mapping=mapping, + ) + metadata.prepare() + + attention = TrtllmAttention( + layer_idx=0, + num_heads=heads, + head_dim=head_dim_qk, + num_kv_heads=1, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value), + q_scaling=1.0, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams( + dim=qk_rope_head_dim, + max_positions=seq, + original_max_positions=seq, + duplicate_data=True, + ), + is_neox=False, + ), + mla_params=MLAParams( + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + qk_nope_head_dim=qk_nope_head_dim, + v_head_dim=v_head_dim, + predicted_tokens_per_seq=q_len, + ), + ) + # Force the THOP adapter. Its SM107 AttentionOp selects the native + # TRTLLM-gen runner; other adapters may route back through FlashInfer. + attention._fmha_manager.fmha_libs = [FallbackFmha(attention)] + + torch.manual_seed(8) + num_tokens = batch * q_len + fused_q = torch.randn( + num_tokens, + heads * head_dim_qk, + dtype=torch.bfloat16, + device=device, + ) + q_pe = torch.randn( + num_tokens, + heads, + qk_rope_head_dim, + dtype=torch.bfloat16, + device=device, + ) + latent_cache = torch.randn( + num_tokens, + head_dim_qk, + dtype=torch.bfloat16, + device=device, + ) + cu_q_seqlens = torch.empty(batch + 1, dtype=torch.int32, device=device) + cu_kv_seqlens = torch.empty(batch + 1, dtype=torch.int32, device=device) + fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=device) + mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=device) + mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=device) + quant_q_buffer = torch.empty( + num_tokens, + heads * head_dim_qk, + dtype=torch.uint8, + device=device, + ) + output = torch.empty( + num_tokens, + heads * kv_lora_rank, + dtype=torch.bfloat16, + device=device, + ) + + def prepare() -> None: + attention.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + ) + + def run() -> None: + attention.forward( + fused_q, + None, + None, + metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=latent_cache, + q_pe=q_pe, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + mla_bmm1_scale=mla_bmm1_scale, + mla_bmm2_scale=mla_bmm2_scale, + quant_q_buffer=quant_q_buffer, + output=output, + ) + + graph_event_ms = None + graph_wall_ms = None + prepare() + run() + torch.cuda.synchronize() + if not bool(torch.isfinite(output).all().item()): + raise RuntimeError("native TRTLLM FP8 MLA output contains non-finite values") + if use_cuda_graph: + prepare() + torch.cuda.synchronize(device) + cuda_graph, capture_stream = _capture_cuda_graph(run, device) + + def replay_graph() -> None: + cuda_graph.replay() + + output.fill_(float("nan")) + torch.cuda.synchronize(device) + replay_graph() + torch.cuda.synchronize(device) + if not bool(torch.isfinite(output).all().item()): + raise RuntimeError("native TRTLLM FP8 CUDA Graph output contains non-finite values") + graph_event_ms, graph_wall_ms = _bench_cuda_graph( + replay_graph, + device, + warmup=warmup, + iters=iters, + ) + avg_ms = graph_event_ms if graph_timing == "event" else graph_wall_ms + del capture_stream + else: + avg_ms = _bench_with_untimed_setup( + run, + prepare, + warmup=warmup, + iters=iters, + queue_delay_cycles=queue_delay_cycles, + ) + finally: + kv_cache_manager.shutdown() + + qk_dim = kv_lora_rank + qk_rope_head_dim + pv_dim = kv_lora_rank + flops = 2 * heads * _causal_token_pairs(seq_lens_list, q_len) * (qk_dim + pv_dim) + tflops = flops / avg_ms / 1e9 + bytes_per_call = _trtllm_mla_io_bytes( + seq_lens_list, + heads, + kv_lora_rank, + qk_rope_head_dim, + q_len, + elem_bytes=1, + ) + gb_s = bytes_per_call / (avg_ms * 1e-3) / 1e9 + graph_label = f" cuda_graph=True graph_timing={graph_timing}" if use_cuda_graph else "" + print( + f"backend={label:>18s} bs={batch:>3d} seq={_seq_label(seq_lens_list):>11s} " + f"heads={heads:>3d} qlen={q_len:>2d}{graph_label}: " + f"{avg_ms:>7.3f} ms {tflops:>6.2f} TFLOP/s " + f"Est.IO {gb_s:>6.1f} GB/s", + flush=True, + ) + if graph_event_ms is not None and graph_wall_ms is not None: + print( + f"cuda_graph event_us={graph_event_ms * 1000.0:.6f} " + f"host_wall_us={graph_wall_ms * 1000.0:.6f}", + flush=True, + ) + elif queue_delay_cycles: + print( + f"fused queue delay={queue_delay_cycles} cycles (excluded from the event interval)", + flush=True, + ) + return avg_ms + + +class _FusedLaunchTimer: + """Time only compiled CuTeDSL fused launches on the active CUDA stream.""" + + def __init__(self, device: torch.device, iters: int, queue_delay_cycles: int = 0) -> None: + self.device = device + self.queue_delay_cycles = queue_delay_cycles + self.enabled = False + self.count = 0 + self.starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + self.ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + + def wrap(self, compiled: Callable[..., object]) -> Callable[..., object]: + def timed_launch(*args: object, **kwargs: object) -> object: + if not self.enabled: + return compiled(*args, **kwargs) + if self.count >= len(self.starts): + raise RuntimeError("observed more timed fused launches than expected") + if self.queue_delay_cycles: + torch.cuda._sleep(self.queue_delay_cycles) + stream = torch.cuda.current_stream(self.device) + self.starts[self.count].record(stream) + result = compiled(*args, **kwargs) + self.ends[self.count].record(stream) + self.count += 1 + return result + + return timed_launch + + def mean_ms(self) -> float: + if self.count != len(self.starts): + raise RuntimeError( + f"expected {len(self.starts)} timed fused launches, got {self.count}" + ) + return ( + sum(start.elapsed_time(end) for start, end in zip(self.starts, self.ends)) / self.count + ) + + +def run_one( + batch: int, + seq: int, + heads: int, + backend: str, + q_len: int = 1, + warmup: int = 0, + iters: int = 1, + timing_scope: str = "full", + fused_queue_delay_cycles: int = 0, + report_correction: bool = False, + use_cuda_graph: bool = False, + generation_step: bool = False, + retain_state: bool = False, + graph_timing: str = "event", +) -> float | tuple[float, RetainedBenchmarkState]: + if warmup < 0: + raise ValueError("warmup must be non-negative") + if iters <= 0: + raise ValueError("iters must be positive") + if q_len <= 0: + raise ValueError("query length must be positive") + if generation_step and q_len > seq: + raise ValueError("generation-step query length must not exceed the sequence length") + if timing_scope not in ("full", "fused", "cupti"): + raise ValueError(f"unsupported timing scope: {timing_scope}") + if fused_queue_delay_cycles < 0: + raise ValueError("fused queue delay cycles must be non-negative") + if graph_timing not in ("event", "wall"): + raise ValueError(f"unsupported CUDA Graph timing mode: {graph_timing}") + if fused_queue_delay_cycles and timing_scope != "fused": + raise ValueError("fused queue delay requires fused timing") + if use_cuda_graph and fused_queue_delay_cycles: + raise ValueError("fused queue delay is not used for CUDA Graph replay") + if use_cuda_graph and backend != "cutedsl": + raise ValueError("CUDA Graph replay requires the CuTeDSL backend") + if use_cuda_graph and timing_scope == "cupti": + raise ValueError("CUDA Graph replay does not support CUPTI timing") + if generation_step and not (backend == "cutedsl" and use_cuda_graph and timing_scope == "full"): + raise ValueError( + "generation-step timing requires the CuTeDSL backend, CUDA Graph, and full timing" + ) + if timing_scope in ("fused", "cupti") and backend != "cutedsl": + raise ValueError(f"{timing_scope} timing is only supported by the CuTeDSL backend") + if retain_state and backend != "cutedsl": + raise ValueError("retained benchmark state requires the CuTeDSL backend") + os.environ[FP4_MLA_ATTENTION_BACKEND_ENV] = backend + seq_lens = _seq_lens_for_batch(batch, seq) + ( + kv_cache_manager, + metadata, + q_nope, + q_pe, + kv_lora_rank, + qk_rope_head_dim, + ) = _build_fp4_mla_attention_decode_case( + seq_lens=seq_lens, + num_heads=heads, + seed=8, + query_len_per_seq=q_len, + ) + metadata.is_cuda_graph = use_cuda_graph + # The synthetic benchmark never changes its sequence or append lengths. + # Bind the authoritative runtime tensors as an already-populated result so + # timing excludes the production MTP length-correction helper. + metadata.fp4_mla_state.generation_kv_lens = metadata.kv_lens_cuda_runtime[:batch] + metadata.fp4_mla_state.generation_append_lens = metadata.prompt_lens_cuda_runtime[:batch] + metadata.fp4_mla_state.generation_lengths_num_tokens = batch * q_len + metadata.fp4_mla_state.generation_lengths_num_seqs = batch + metadata.fp4_mla_state.generation_lengths_num_contexts = metadata.num_contexts + metadata.fp4_mla_state.generation_lengths_capture_recorded = True + if use_cuda_graph: + # Synthetic benchmark page tables are immutable and never grow on replay. + metadata.fp4_mla_state._paged_kv_indices = metadata.fp4_mla_state.paged_kv_indices + generation_latent = None + generation_rotary_cos_sin = None + persistent_v_pack = False + if generation_step: + context_lens = torch.tensor( + [seq_len - q_len for seq_len in seq_lens], + dtype=torch.int32, + device=q_nope.device, + ) + metadata.num_ctx_tokens = 0 + metadata.fp4_mla_state.batch_indices = torch.arange( + batch, dtype=torch.int32, device=q_nope.device + ).repeat_interleave(q_len) + metadata.fp4_mla_state.positions = ( + context_lens[:, None] + + torch.arange(q_len, dtype=torch.int32, device=q_nope.device)[None, :] + ).reshape(-1) + persistent_v_packed = fp4_mla_backend._get_fp4_mla_v_packed_pool(metadata, 0) + if persistent_v_packed is None and not _fp4_mla_cutedsl_fused_v_transpose_enabled(): + kv_cache_manager.shutdown() + raise RuntimeError( + "generation-step timing requires the manager-owned persistent V-packed sidecar" + ) + persistent_v_pack = persistent_v_packed is not None + generation_latent = ( + torch.randn( + batch * q_len, + kv_lora_rank + qk_rope_head_dim, + dtype=torch.bfloat16, + device=q_nope.device, + ) + * 0.25 + ).clamp_(-1.0, 1.0) + generation_rotary_cos_sin = torch.zeros( + max(seq_lens) + q_len, + qk_rope_head_dim, + 2, + dtype=torch.float32, + device=q_nope.device, + ) + generation_rotary_cos_sin[..., 0] = 1.0 + launcher_module = None + original_launcher = None + launcher_wrapper = None + core_launcher = None + original_compile_fused = None + launcher_tensors = None + launcher_tensor_ptrs = None + launcher_kwargs = None + retained_state = None + capture_stream = None + cuda_graph = None + graph_event_ms = None + graph_wall_ms = None + cupti_cuda_events = [] + cupti_fused_launches_per_iter = 1 + cupti_fused_slot_stats = [] + try: + if report_correction: + _, exact_probs, _ = _fp4_mla_attention_decode_reference( + metadata, + q_nope, + q_pe, + sm_scale=0.1, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + ) + correction_stats = _fp4_mla_lazy_rebase_stats(exact_probs) + print( + "correction " + f"total_rows={correction_stats.total_rows} " + f"row_rebase_rows={correction_stats.rebase_rows} " + f"row_rebase_total={correction_stats.rebase_total} " + f"row_rebase_max={correction_stats.rebase_max_per_row} " + f"max_anchor_delta_log2=" + f"{correction_stats.max_anchor_delta_log2:.6f} " + f"threshold_log2={correction_stats.threshold_log2:.6f}", + flush=True, + ) + if backend == "cutedsl": + cutedsl = fp4_mla_backend._fp4_mla_cutedsl_kernel_module() + + launcher_module = cutedsl + original_launcher = cutedsl.run_trtllm_fp4_mla_decode_page_native_from_raw + + def record_launcher(*args: object, **kwargs: object) -> object: + nonlocal launcher_kwargs, launcher_tensors, launcher_tensor_ptrs + current_tensors = tuple(args) + current_ptrs = _launcher_tensor_pointers(current_tensors) + if launcher_tensor_ptrs is None: + launcher_tensors = current_tensors + launcher_tensor_ptrs = current_ptrs + launcher_kwargs = dict(kwargs) + elif current_ptrs != launcher_tensor_ptrs: + raise RuntimeError( + "CuTeDSL launcher tensor pointers changed within one benchmark run" + ) + return original_launcher(*args, **kwargs) + + launcher_wrapper = record_launcher + except (ImportError, TypeError, ValueError, RuntimeError): + kv_cache_manager.shutdown() + raise + + try: + if launcher_module is not None and launcher_wrapper is not None: + launcher_module.run_trtllm_fp4_mla_decode_page_native_from_raw = launcher_wrapper + output = torch.full_like(q_nope, float("nan")) + q = torch.cat((q_nope, q_pe), dim=-1).contiguous() + baseline_q_fp4, baseline_q_sf = torch.ops.trtllm.fp4_quantize_with_residual( + q.view(-1, q.shape[-1]), + metadata.fp4_mla_state.q_global_scale, + FP4_MLA_Q_RESIDUAL_DIM, + is_act=True, + ) + baseline_q_sf = baseline_q_sf.view(torch.float8_e4m3fn).reshape(-1) + + def decode() -> None: + q_fp4 = getattr(metadata.fp4_mla_state, "prequantized_q", None) + q_sf = getattr(metadata.fp4_mla_state, "prequantized_q_sf", None) + q_batch_capacity = getattr(metadata.fp4_mla_state, "q_batch_capacity", None) + if q_fp4 is None: + q_fp4 = baseline_q_fp4 + q_sf = baseline_q_sf + q_batch_capacity = q.shape[0] + run_fp4_mla_attention_decode( + metadata, + layer_idx=0, + local_layer=0, + q=q, + output=output, + sm_scale=0.1, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + prequantized_q=q_fp4, + prequantized_q_sf=q_sf, + q_batch_capacity=q_batch_capacity, + ) + + run = decode + if generation_step: + assert generation_latent is not None + assert generation_rotary_cos_sin is not None + + def generation_run() -> None: + scatter_fp4_mla_kv_cache( + metadata, + generation_latent, + layer_idx=0, + token_offset=0, + phase="generation", + local_layer=0, + v_head_dim=kv_lora_rank, + rotary_cos_sin=generation_rotary_cos_sin, + q_pe=q[..., kv_lora_rank:], + q_rope_out=q[..., kv_lora_rank:], + q_quant_input=q, + ) + decode() + + run = generation_run + + run() + torch.cuda.synchronize() + if launcher_module is not None and original_launcher is not None: + core_launcher = original_launcher + launcher_module.run_trtllm_fp4_mla_decode_page_native_from_raw = original_launcher + original_launcher = None + if not bool(torch.isfinite(output).all().item()): + raise RuntimeError("FP4 MLA benchmark output contains non-finite values.") + if use_cuda_graph: + if timing_scope == "fused": + if core_launcher is None or launcher_tensors is None or launcher_kwargs is None: + raise RuntimeError("core-only CUDA Graph capture missed the CuTeDSL launcher") + + def eager_run() -> None: + core_launcher(*launcher_tensors, **launcher_kwargs) + + else: + eager_run = run + cuda_graph, capture_stream = _capture_cuda_graph(eager_run, output.device) + + def replay_graph() -> None: + cuda_graph.replay() + + run = replay_graph + output.fill_(float("nan")) + run() + torch.cuda.synchronize(output.device) + if not bool(torch.isfinite(output).all().item()): + raise RuntimeError( + "FP4 MLA CUDA Graph replay did not overwrite the output sentinel." + ) + if use_cuda_graph: + graph_event_ms, graph_wall_ms = _bench_cuda_graph( + run, + output.device, + warmup=warmup, + iters=iters, + ) + avg_ms = graph_event_ms if graph_timing == "event" else graph_wall_ms + elif timing_scope == "fused": + timer = _FusedLaunchTimer( + output.device, + iters, + queue_delay_cycles=fused_queue_delay_cycles, + ) + original_compile_fused = launcher_module._compile_fused + compile_cache_size = len(launcher_module._FUSED_COMPILE_CACHE) + + def compile_timed_fused(*args: object, **kwargs: object) -> Callable[..., object]: + return timer.wrap(original_compile_fused(*args, **kwargs)) + + launcher_module._compile_fused = compile_timed_fused + for _ in range(warmup): + run() + torch.cuda.synchronize(output.device) + if len(launcher_module._FUSED_COMPILE_CACHE) != compile_cache_size: + raise RuntimeError("CuTeDSL fused compile cache changed during warmup") + timer.enabled = True + try: + for _ in range(iters): + run() + finally: + timer.enabled = False + torch.cuda.synchronize(output.device) + if len(launcher_module._FUSED_COMPILE_CACHE) != compile_cache_size: + raise RuntimeError("CuTeDSL fused compile cache changed during timing") + avg_ms = timer.mean_ms() + elif timing_scope == "cupti": + compile_cache_size = len(launcher_module._FUSED_COMPILE_CACHE) + for _ in range(warmup): + run() + torch.cuda.synchronize(output.device) + from torch.profiler import ProfilerActivity, profile + + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + acc_events=True, + ) as profiler: + for _ in range(iters): + run() + torch.cuda.synchronize(output.device) + if len(launcher_module._FUSED_COMPILE_CACHE) != compile_cache_size: + raise RuntimeError("CuTeDSL fused compile cache changed during CUPTI timing") + cupti_cuda_events = [ + event + for event in profiler.events() + if event.device_type == torch.autograd.DeviceType.CUDA + ] + target_prefix = "kernel_cutlass_kernel_" + fused_events = [ + event for event in cupti_cuda_events if event.name.startswith(target_prefix) + ] + fused_events.sort(key=lambda event: event.time_range.start) + if not fused_events or len(fused_events) % iters != 0: + raise RuntimeError( + "expected a fixed positive number of CUPTI fused events " + f"per iteration, got {len(fused_events)} events for {iters} iterations" + ) + cupti_fused_launches_per_iter = len(fused_events) // iters + for slot in range(cupti_fused_launches_per_iter): + slot_events = fused_events[slot::cupti_fused_launches_per_iter] + if len(slot_events) != iters: + raise RuntimeError( + f"expected {iters} CUPTI fused events for slot {slot}, " + f"got {len(slot_events)}" + ) + cupti_fused_slot_stats.append( + ( + slot, + len(slot_events), + sum(event.self_device_time_total for event in slot_events), + ) + ) + # PyTorch profiler device times are reported in microseconds. + avg_ms = sum(event.self_device_time_total for event in fused_events) / iters / 1000.0 + if avg_ms <= 0.0: + raise RuntimeError(f"CUPTI reported non-positive fused time {avg_ms}") + else: + avg_ms = _bench(run, warmup=warmup, iters=iters) + if not bool(torch.isfinite(output).all().item()): + raise RuntimeError("FP4 MLA timed output contains non-finite values.") + k_residual_dim = FP4_MLA_K_RESIDUAL_DIM if backend in ("triton", "cutedsl") else 0 + qk_dim = kv_lora_rank + qk_rope_head_dim + FP4_MLA_Q_RESIDUAL_DIM + k_residual_dim + pv_dim = kv_lora_rank + flops = 2 * heads * _causal_token_pairs(seq_lens, q_len) * (qk_dim + pv_dim) + tflops = flops / avg_ms / 1e9 + min_tflops = float(os.environ.get("FP4_MLA_MIN_TFLOPS", "0")) + if tflops < min_tflops: + raise RuntimeError(f"FP4 MLA throughput {tflops:.2f} TFLOP/s is below {min_tflops:.2f}") + # None of these estimates is a hardware counter or measured HBM bandwidth. + logical_bytes = _kernel_io_bytes( + seq_lens, + heads, + kv_lora_rank, + qk_rope_head_dim, + q_len, + k_residual_dim, + ) + fixed_tile_bytes = _fixed_tile_request_bytes( + seq_lens, + heads, + kv_lora_rank, + qk_rope_head_dim, + q_len, + k_residual_dim, + ) + tma_smem_bytes = _tma_smem_completion_bytes( + seq_lens, + heads, + kv_lora_rank, + qk_rope_head_dim, + q_len, + k_residual_dim, + ) + logical_gb_s = logical_bytes / (avg_ms * 1e-3) / 1e9 + fixed_tile_gb_s = fixed_tile_bytes / (avg_ms * 1e-3) / 1e9 + tma_smem_gb_s = tma_smem_bytes / (avg_ms * 1e-3) / 1e9 + graph_timing_label = f" graph_timing={graph_timing}" if use_cuda_graph else "" + generation_label = ( + f" generation_step=True persistent_v_pack={persistent_v_pack}" + if generation_step + else "" + ) + print( + f"backend={backend:>12s} bs={batch:>3d} seq={_seq_label(seq_lens):>11s} " + f"heads={heads:>3d} qlen={q_len:>2d}{graph_timing_label}{generation_label} " + f"scope={timing_scope}: " + f"{avg_ms:>10.6f} ms {tflops:>6.2f} TFLOP/s " + f"Est.Logical {logical_gb_s:>6.1f} GB/s " + f"Est.GlobalReq {fixed_tile_gb_s:>6.1f} GB/s " + f"Est.TMA-SMEM {tma_smem_gb_s:>6.1f} GB/s", + flush=True, + ) + if graph_event_ms is not None and graph_wall_ms is not None: + print( + f"cuda_graph event_us={graph_event_ms * 1000.0:.6f} " + f"host_wall_us={graph_wall_ms * 1000.0:.6f}", + flush=True, + ) + elif fused_queue_delay_cycles: + print( + f"fused queue delay={fused_queue_delay_cycles} cycles " + "(excluded from the event interval)", + flush=True, + ) + if cupti_cuda_events: + if cupti_fused_launches_per_iter > 1: + print( + f"cupti fused launches/iteration={cupti_fused_launches_per_iter}", + flush=True, + ) + for slot, count, total_us in cupti_fused_slot_stats: + role = ( + ("fast_partition", "slow_partition")[slot] + if cupti_fused_launches_per_iter == 2 + else "single" + ) + print( + f"cupti fused slot={slot} role={role} count={count} " + f"total_us={total_us:.3f} avg_us={total_us / count:.3f}", + flush=True, + ) + event_totals = {} + ignored_events = {"Activity Buffer Request", "Lazy Function Loading"} + for event in cupti_cuda_events: + if event.name in ignored_events: + continue + count, total_us = event_totals.get(event.name, (0, 0.0)) + event_totals[event.name] = ( + count + 1, + total_us + event.self_device_time_total, + ) + for name, (count, total_us) in sorted( + event_totals.items(), + key=lambda item: item[1][1], + reverse=True, + )[:8]: + print( + f"cupti count={count} total_us={total_us:.3f} " + f"avg_us={total_us / count:.3f} kernel={name[:120]}", + flush=True, + ) + if retain_state: + if launcher_tensors is None or launcher_tensor_ptrs is None: + raise RuntimeError("retained state did not observe CuTeDSL launcher tensors") + retained_state = RetainedBenchmarkState( + kv_cache_manager, + launcher_tensors, + { + "metadata": metadata, + "q": q, + "q_nope": q_nope, + "q_pe": q_pe, + "output": output, + "generation_latent": generation_latent, + "cuda_graph": cuda_graph, + "capture_stream": capture_stream, + }, + ) + return avg_ms, retained_state + return avg_ms + finally: + if launcher_module is not None and original_launcher is not None: + launcher_module.run_trtllm_fp4_mla_decode_page_native_from_raw = original_launcher + if launcher_module is not None and original_compile_fused is not None: + launcher_module._compile_fused = original_compile_fused + if retained_state is None: + kv_cache_manager.shutdown() + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--batch", type=int, nargs="+", default=None) + p.add_argument("--seq", type=int, default=30080) + p.add_argument("--heads", type=int, default=128) + p.add_argument( + "--q-len", + "--mtp-len", + dest="q_len", + type=int, + default=1, + help="Query tokens per sequence (>1 for MTP / speculative decoding).", + ) + p.add_argument( + "--backend", + default=None, + choices=BACKEND_CHOICES, + help="Backend to benchmark; default runs all backends supported on the active GPU.", + ) + p.add_argument("--warmup", type=int, default=10) + p.add_argument("--iters", type=int, default=50) + p.add_argument( + "--timing-scope", + choices=("full", "fused", "cupti"), + default="full", + help=( + "full times the complete integrated decode call; fused isolates the " + "CuTeDSL or Rubin FP8 attention core; cupti uses exact CUDA activity " + "timestamps for the fused kernel (CuTeDSL backend only)." + ), + ) + p.add_argument( + "--fused-queue-delay-cycles", + "--queue-delay-cycles", + type=int, + default=0, + help=( + "enqueue a same-stream GPU delay before each timed start event to " + "hide host launch submission; the delay is excluded from timing" + ), + ) + p.add_argument( + "--cuda-graph", + action="store_true", + help="Capture the selected timed call once and benchmark Graph replay.", + ) + p.add_argument( + "--graph-timing", + choices=("event", "wall"), + default="event", + help=( + "select the primary CUDA Graph metric; event is GPU elapsed time, " + "while wall includes host Graph launch overhead" + ), + ) + p.add_argument( + "--generation-step", + action="store_true", + help=( + "Capture generation KV scatter, high-precision KV update, persistent " + "touched-page V repack, and decode together (requires CuTeDSL CUDA " + "Graph full timing)." + ), + ) + p.add_argument( + "--report-correction", + action="store_true", + help=( + "Classify lazy resident-O correction activity from the exact " + "quantized-QK reference before timing." + ), + ) + args = p.parse_args() + + if args.batch is not None and any(batch <= 0 for batch in args.batch): + p.error("--batch values must be positive") + if args.seq <= 0 or args.heads <= 0: + p.error("--seq and --heads must be positive") + if not 0 < args.q_len <= args.seq: + p.error("--q-len must be positive and must not exceed --seq") + if args.warmup < 0: + p.error("--warmup must be non-negative") + if args.iters <= 0: + p.error("--iters must be positive") + if args.fused_queue_delay_cycles < 0: + p.error("--fused-queue-delay-cycles must be non-negative") + if args.fused_queue_delay_cycles and args.timing_scope != "fused": + p.error("--fused-queue-delay-cycles requires --timing-scope fused") + if args.fused_queue_delay_cycles and args.cuda_graph: + p.error("--fused-queue-delay-cycles is not used with --cuda-graph") + graph_backends = {"cutedsl", "trtllm_fp8_rubin"} + if args.cuda_graph and args.backend not in graph_backends: + p.error("--cuda-graph requires --backend cutedsl or trtllm_fp8_rubin") + if args.cuda_graph and args.timing_scope == "cupti": + p.error("--cuda-graph does not support --timing-scope cupti") + if args.cuda_graph and args.backend == "trtllm_fp8_rubin" and args.timing_scope != "fused": + p.error("FP8 core-only --cuda-graph requires --timing-scope fused") + if args.generation_step and not ( + args.backend == "cutedsl" and args.cuda_graph and args.timing_scope == "full" + ): + p.error("--generation-step requires --backend cutedsl --cuda-graph --timing-scope full") + if args.report_correction and args.backend != "cutedsl": + p.error("--report-correction requires --backend cutedsl") + if args.timing_scope == "fused" and args.backend not in ("cutedsl", "trtllm_fp8_rubin"): + p.error("--timing-scope fused requires --backend cutedsl or trtllm_fp8_rubin") + if args.timing_scope == "cupti" and args.backend != "cutedsl": + p.error("--timing-scope cupti requires --backend cutedsl") + + batches = args.batch if args.batch else [16, 30, 60, 120, 200, 300] + if not torch.cuda.is_available(): + p.error("This benchmark requires a CUDA GPU.") + is_rubin = torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 7) + if args.backend: + if args.backend == "cutedsl": + if not is_rubin: + p.error("--backend cutedsl requires a Rubin SM107 GPU.") + if not _cutedsl_backend_available(): + p.error("--backend cutedsl requires the CTM and CuTeDSL runtime packages.") + if args.backend == "trtllm_fp8_rubin" and not is_rubin: + p.error("--backend trtllm_fp8_rubin requires a Rubin SM107 GPU.") + backends = [args.backend] + else: + excluded_backends = {"cutedsl"} + excluded_backends.add("trtllm_fp8" if is_rubin else "trtllm_fp8_rubin") + backends = [backend for backend in BACKEND_CHOICES if backend not in excluded_backends] + if is_rubin and _cutedsl_backend_available(): + backends.append("cutedsl") + + for b in batches: + for be in backends: + if be == "trtllm_fp8": + run_one_trtllm(b, args.seq, args.heads, args.q_len, args.warmup, args.iters) + elif be == "trtllm_fp8_rubin": + run_one_trtllm_rubin( + b, + args.seq, + args.heads, + args.q_len, + args.warmup, + args.iters, + queue_delay_cycles=args.fused_queue_delay_cycles, + use_cuda_graph=args.cuda_graph, + graph_timing=args.graph_timing, + ) + else: + run_one( + b, + args.seq, + args.heads, + be, + args.q_len, + args.warmup, + args.iters, + args.timing_scope, + args.fused_queue_delay_cycles, + args.report_correction, + args.cuda_graph, + args.generation_step, + graph_timing=args.graph_timing, + ) + + +if __name__ == "__main__": + main() From b0d8b984bf736e4712523c96d919fe567a72c810 Mon Sep 17 00:00:00 2001 From: Tracin <10434017+Tracin@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:13:27 -0700 Subject: [PATCH 21/21] fix: isolate dense FP4 MLA dispatch and remove stale state Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> --- .../batch_manager/kvCacheManager.cpp | 3 - .../_torch/attention/backends/fmha/fp4_mla.py | 2 +- .../attention/backends/fmha/interface.py | 6 +- .../attention/backends/fp4_mla/__init__.py | 3 - .../backends/fp4_mla/cache_update.py | 225 ++++++------------ .../attention/backends/fp4_mla/config.py | 43 ---- .../attention/backends/fp4_mla/decode.py | 53 +++++ .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 127 +++++++++- ...p4_mla_cutedsl_mufu16_fused_v_transpose.py | 127 +++++++++- .../attention/backends/fp4_mla/metadata.py | 1 - .../attention/backends/fp4_mla/v_cache.py | 8 +- .../_torch/attention/backends/trtllm.py | 22 +- tensorrt_llm/_torch/attention/mla.py | 6 +- tensorrt_llm/_torch/pyexecutor/_util.py | 13 +- .../_torch/pyexecutor/model_loader.py | 8 +- .../_torch/pyexecutor/resource_manager.py | 4 - tests/microbenchmarks/bench_fp4_mla_decode.py | 1 - .../unittest/_torch/attention/test_fp4_mla.py | 7 +- 18 files changed, 388 insertions(+), 271 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 43c7e6605ec2..255452dd654b 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -3270,9 +3270,6 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case , mEnableBlockReuse{mSinkBubbleLength > 0 ? false : enableBlockReuse} { - TLLM_CHECK_WITH_INFO(dtype != tensorrt_llm::DataType::kFP4 || cacheType != CacheType::kSELFKONLY, - "NVFP4 SELFKONLY cache storage requires Fp4MlaKVCacheManagerV2; KVCacheManager V1 is not supported."); - // When num_layers < len(maxAttentionWindowVec), not all window sizes in the // repeating pattern are used. Update mMaxAttentionWindow to the actual // maximum window size that has been allocated in the block manager. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index 9e65fa7c522a..9b47e4a4b806 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -59,7 +59,7 @@ def __init__(self, attn: "TrtllmAttention") -> None: @classmethod def _is_available(cls, attn: "TrtllmAttention") -> bool: - return attn.is_mla_enable and attn.has_fp4_kv_cache + return bool(getattr(attn, "uses_fp4_mla_attention", False)) def _is_supported( self, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index efac2712c932..2d1bd5c86742 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/interface.py @@ -77,11 +77,7 @@ def is_available(cls, attn: "TrtllmAttention") -> bool: f"{cls.__name__} is unavailable: skip-correction is enabled and unsupported." ) return False - if ( - getattr(attn, "is_mla_enable", False) - and getattr(attn, "has_fp4_kv_cache", False) - and not cls.supports_fp4_mla - ): + if getattr(attn, "uses_fp4_mla_attention", False) and not cls.supports_fp4_mla: return False return cls._is_available(attn) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py index 596b208194a7..a2a0389b6204 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -33,7 +33,6 @@ from .config import ( _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD as _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD, ) -from .config import _FP4_MLA_TRITON_PRELOAD_KEYS as _FP4_MLA_TRITON_PRELOAD_KEYS from .config import ( FP4_BLOCK_SIZE, FP4_MLA_ATTENTION_BACKEND_ENV, @@ -68,8 +67,6 @@ from .config import _fp4_mla_cutedsl_kernel_module as _fp4_mla_cutedsl_kernel_module from .config import _fp4_mla_q1_kv_blocks_per_program as _fp4_mla_q1_kv_blocks_per_program from .config import _fp4_mla_q1_prefix_blocks_per_program as _fp4_mla_q1_prefix_blocks_per_program -from .config import _fp4_mla_q1_preload_variants as _fp4_mla_q1_preload_variants -from .config import _fp4_mla_triton_preload_key_set as _fp4_mla_triton_preload_key_set from .config import _HPUpdatePhase as _HPUpdatePhase from .decode import _SM_COUNT_CACHE as _SM_COUNT_CACHE from .decode import _cutedsl_pad_q_and_sf_kernel as _cutedsl_pad_q_and_sf_kernel diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py index 0456f702fafd..d2fd2f7da0ce 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py @@ -30,8 +30,6 @@ _fp4_mla_cutedsl_fused_v_transpose_enabled, _fp4_mla_q1_kv_blocks_per_program, _fp4_mla_q1_prefix_blocks_per_program, - _fp4_mla_q1_preload_variants, - _fp4_mla_triton_preload_key_set, _HPUpdatePhase, ) from .fp4_mla_kernels import ( @@ -434,6 +432,8 @@ def _scatter_fp4_mla_kv_cache_2d_generation( q_sf_out: torch.Tensor, v_packed_base: Optional[torch.Tensor], v_page_offset: int, + helix_position_offsets: Optional[torch.Tensor], + helix_is_inactive_rank: Optional[torch.Tensor], ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: num_contexts = metadata.num_contexts num_seqs = metadata.num_seqs @@ -476,8 +476,60 @@ def _scatter_fp4_mla_kv_cache_2d_generation( num_hp_pages = pool.shape[0] max_gen_len = num_tokens // num_gen + use_helix = helix_position_offsets is not None or helix_is_inactive_rank is not None + use_helix_local_slots = use_helix and bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + if use_helix: + if helix_position_offsets is None or helix_is_inactive_rank is None: + raise RuntimeError( + "FP4 MLA Helix requires both position-offset and inactive-rank metadata." + ) + if max_gen_len != 1 and not use_helix_local_slots: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + if ( + helix_position_offsets.dtype != torch.int32 + or helix_position_offsets.device != latent_cache.device + or helix_position_offsets.ndim != 1 + or helix_position_offsets.numel() < num_tokens + or not helix_position_offsets.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix position offsets must be a contiguous same-device " + "int32 tensor covering every generation token." + ) + if ( + helix_is_inactive_rank.dtype != torch.bool + or helix_is_inactive_rank.device != latent_cache.device + or helix_is_inactive_rank.ndim != 1 + or helix_is_inactive_rank.numel() < num_gen + or not helix_is_inactive_rank.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix inactive-rank metadata must be a contiguous " + "same-device bool tensor covering every generation sequence." + ) + if use_helix_local_slots: + helix_local_slots = getattr(metadata, "helix_local_slots", None) + if ( + not isinstance(helix_local_slots, torch.Tensor) + or helix_local_slots.dtype != torch.int32 + or helix_local_slots.device != latent_cache.device + or helix_local_slots.ndim != 1 + or helix_local_slots.numel() < num_tokens + or not helix_local_slots.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix local slots must be a contiguous " + "same-device int32 tensor covering every generation token." + ) + else: + helix_local_slots = helix_position_offsets + else: + helix_position_offsets = kv_lens_gen + helix_local_slots = kv_lens_gen + helix_is_inactive_rank = gen_lens_gen _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) - max_rewind_len = hp_pool_size - HP_BLOCK_SIZE page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) rope_dim = head_dim - v_head_dim block_q_heads = 32 @@ -589,6 +641,9 @@ def launch_generation_update( q_sf_output, kv_lens_gen, gen_lens_gen, + helix_position_offsets, + helix_local_slots, + helix_is_inactive_rank, page_ids, hp_page_ids, metadata.fp4_mla_state.paged_kv_indptr_decode, @@ -627,6 +682,8 @@ def launch_generation_update( K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, STORE_K_RESIDUAL=store_k_residual, FUSE_ROPE_CACHE_STORE=True, + USE_HELIX=use_helix, + USE_HELIX_LOCAL_SLOTS=use_helix_local_slots, WRITE_V_PACKED=write_v_packed, MAX_GEN_TILES=max_gen_tiles_variant, ROPE_DIM=rope_dim, @@ -646,159 +703,6 @@ def launch_generation_update( maxnreg=56, ) - # Triton compiles and loads a CUDA module on first launch. Use runtime-zero - # work here so every reachable static tuning variant is resident before - # warmup hands the engine to serving. - if getattr(metadata, "is_warmup", False) and not torch.cuda.is_current_stream_capturing(): - configured_generation_len = int(getattr(metadata, "max_total_draft_tokens", 0) or 0) + 1 - max_preload_generation_len = max(max_gen_len, configured_generation_len) - if max_preload_generation_len - 1 > max_rewind_len: - raise NotImplementedError( - "FP4 MLA finite Triton preload exceeds the HP ring's rewind slack: " - f"max_rewind={max_rewind_len}, generation=" - f"{max_preload_generation_len}." - ) - max_num_sequences = int(getattr(metadata, "max_num_sequences", num_seqs) or num_seqs) - q1_variants = _fp4_mla_q1_preload_variants( - max_num_sequences, - v_head_dim, - ) - multi_token_tiles = tuple( - sorted( - { - _ceil_div(gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) - for gen_len in range(2, max_preload_generation_len + 1) - } - ) - ) - preload_key = ( - "generation-cache-update", - str(latent_cache.device), - hp_pool_size, - hp_head_dim, - v_head_dim, - num_q_heads, - metadata.page_size, - write_v_packed, - store_k_residual, - tuple(q1_variants), - multi_token_tiles, - tuple(kv_cache.stride()), - tuple(sf_cache.stride()), - tuple(pool.stride()), - tuple(v_sf.stride()), - tuple(v_packed_output.stride()), - tuple(q_pe_input.stride()), - tuple(q_rope_output.stride()), - str(kv_cache.dtype), - str(latent_cache.dtype), - str(q_pe_input.dtype), - str(q_fp4_output.dtype), - str(q_sf_output.dtype), - ) - preload_keys = _fp4_mla_triton_preload_key_set(metadata) - if preload_key not in preload_keys: - context_page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) - if not isinstance(context_page_ids, torch.Tensor): - context_page_ids = page_ids - context_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) - if not isinstance(context_indptr, torch.Tensor): - context_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode - context_batch_indices = getattr(metadata.fp4_mla_state, "batch_indices", None) - if not isinstance(context_batch_indices, torch.Tensor): - context_batch_indices = hp_page_ids - context_positions = getattr(metadata.fp4_mla_state, "positions", None) - if not isinstance(context_positions, torch.Tensor): - context_positions = hp_page_ids - _fp4_mla_context_cache_update_kernel[(1, num_dim_blocks)]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - latent_cache, - latent_cache, - global_scale, - rotary_table, - pool, - hp_page_ids, - context_batch_indices, - context_positions, - context_page_ids, - context_indptr, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - latent_cache.stride(0), - latent_cache.stride(1), - latent_cache.stride(0), - 0, - 0, - v_sf.stride(0), - v_sf.stride(1), - v_packed_s0, - v_packed_s1, - pool.stride(0), - pool.stride(1), - HEAD_D=head_dim, - V_HEAD_D=v_head_dim, - HP_BLOCK=FP4_BLOCK_SIZE, - HP_POOL_SIZE=hp_pool_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=sf_per_page, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - STORE_K_RESIDUAL=store_k_residual, - ROPE_DIM=rope_dim, - APPLY_K_ROPE=True, - APPLY_Q_ROPE=False, - NUM_DIM_BLOCKS=num_dim_blocks, - NUM_Q_HEADS=0, - Q_NOPE_DIM=0, - BLOCK_Q_HEADS=16, - POOL_HEAD_D=hp_head_dim, - STORE_HP_TAIL=True, - WRITE_V_PACKED=write_v_packed, - ) - q1_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM - for q1_kv_blocks, q1_prefix_blocks_per_program in q1_variants: - launch_generation_update( - (1, 1), - page_ids_len=0, - indptr_len=0, - max_gen_tiles_variant=1, - q_prefix_block_dim_variant=FP4_MLA_Q1_PREFIX_BLOCK_DIM, - q_prefix_blocks_variant=q1_prefix_blocks, - q_prefix_blocks_per_program_variant=q1_prefix_blocks_per_program, - q1_kv_blocks_per_program_variant=q1_kv_blocks, - ) - multi_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q_PREFIX_BLOCK_DIM - for multi_token_tile in multi_token_tiles: - launch_generation_update( - (1, 1), - page_ids_len=0, - indptr_len=0, - max_gen_tiles_variant=multi_token_tile, - q_prefix_block_dim_variant=FP4_MLA_Q_PREFIX_BLOCK_DIM, - q_prefix_blocks_variant=multi_prefix_blocks, - q_prefix_blocks_per_program_variant=1, - q1_kv_blocks_per_program_variant=1, - ) - torch.cuda.synchronize(latent_cache.device) - preload_keys.add(preload_key) - launch_generation_update( launch_grid, page_ids_len=page_ids.shape[0], @@ -828,6 +732,8 @@ def scatter_fp4_mla_kv_cache( q_pe: Optional[torch.Tensor] = None, q_rope_out: Optional[torch.Tensor] = None, q_quant_input: Optional[torch.Tensor] = None, + helix_position_offsets: Optional[torch.Tensor] = None, + helix_is_inactive_rank: Optional[torch.Tensor] = None, q_context: Optional[torch.Tensor] = None, q_nope_head_dim: Optional[int] = None, ) -> bool: @@ -946,6 +852,11 @@ def scatter_fp4_mla_kv_cache( else: if q_context is not None or q_nope_head_dim is not None: raise ValueError("FP4 MLA generation cache update does not accept context Q tensors.") + if (helix_position_offsets is None) != (helix_is_inactive_rank is None): + raise ValueError( + "FP4 MLA Helix position-offset and inactive-rank metadata " + "must be provided together." + ) if not all(arg is not None for arg in generation_inputs): raise ValueError( "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " @@ -1084,6 +995,8 @@ def scatter_fp4_mla_kv_cache( q_sf_out=q_sf_out, v_packed_base=v_packed_base, v_page_offset=v_page_offset, + helix_position_offsets=helix_position_offsets, + helix_is_inactive_rank=helix_is_inactive_rank, ) v_pack_page_ids = _fp4_mla_generation_page_ids( metadata, metadata.num_seqs - metadata.num_contexts diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py index 80626bb41deb..e11a0e654faf 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py @@ -102,9 +102,6 @@ _HPUpdatePhase = Literal["context", "generation"] -_FP4_MLA_TRITON_PRELOAD_KEYS = "_fp4_mla_triton_preload_keys" - - _FP4_MLA_PAGE_TABLE_TILE_SIZE = 128 @@ -200,43 +197,3 @@ def _fp4_mla_q1_prefix_blocks_per_program( return min(4, max_prefix_blocks) return min(2, max_prefix_blocks) return 1 - - -def _fp4_mla_q1_preload_variants( - max_num_sequences: int, - v_head_dim: int, -) -> tuple[tuple[int, int], ...]: - """Return every Q1 tuning variant reachable by the configured batch limit.""" - batch_sizes = [1] - for threshold in ( - _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD, - _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD, - _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD, - _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD, - ): - if threshold <= max_num_sequences: - batch_sizes.append(threshold) - - variants = [] - for batch_size in batch_sizes: - kv_blocks = _fp4_mla_q1_kv_blocks_per_program(batch_size, v_head_dim) - prefix_blocks = _fp4_mla_q1_prefix_blocks_per_program( - batch_size, - kv_blocks, - ) - variant = (kv_blocks, prefix_blocks) - if variant not in variants: - variants.append(variant) - return tuple(variants) - - -def _fp4_mla_triton_preload_key_set(metadata: Any) -> set[tuple[object, ...]]: - """Return the engine-scoped set of Triton variants loaded during warmup.""" - owner = getattr(metadata, "kv_cache_manager", None) - if owner is None: - owner = metadata - preload_keys = getattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, None) - if preload_keys is None: - preload_keys = set() - setattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, preload_keys) - return preload_keys diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py index a61381cd1ed4..424a1c95c5b9 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py @@ -777,6 +777,7 @@ def run_fp4_mla_attention_decode( prequantized_q: torch.Tensor, prequantized_q_sf: torch.Tensor, q_batch_capacity: int, + softmax_stats_tensor: torch.Tensor | None = None, ) -> None: """Run MLA decode with FP4 QK and FP4 PV tensor-core matmuls. @@ -824,6 +825,43 @@ def run_fp4_mla_attention_decode( raise ValueError("FP4 MLA attention output batch dimensions do not match.") backend = _fp4_mla_attention_backend() + helix_spec_tokens_valid = bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + helix_kv_bounds = None + if softmax_stats_tensor is not None: + if backend != _FP4_MLA_CUTEDSL_BACKEND: + raise NotImplementedError( + "FP4 MLA Helix softmax stats require the cutedsl attention backend." + ) + if query_len_per_seq != 1 and not helix_spec_tokens_valid: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + expected_stats_shape = (num_queries, num_heads, 2) + if ( + softmax_stats_tensor.shape != expected_stats_shape + or softmax_stats_tensor.dtype != torch.float32 + or softmax_stats_tensor.device != q.device + or not softmax_stats_tensor.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix requires contiguous same-device float32 softmax " + f"stats with shape {expected_stats_shape}." + ) + if helix_spec_tokens_valid: + helix_kv_bounds = getattr(metadata, "helix_kv_bounds", None) + if ( + not isinstance(helix_kv_bounds, torch.Tensor) + or helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.device != q.device + or helix_kv_bounds.ndim != 1 + or helix_kv_bounds.numel() < num_queries + or not helix_kv_bounds.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix KV bounds must be a contiguous " + "same-device int32 tensor covering every query token." + ) + helix_kv_bounds = helix_kv_bounds[:num_queries] if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: raise RuntimeError( "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." @@ -1026,6 +1064,15 @@ def run_fp4_mla_attention_decode( ) kernel_output = output + kernel_softmax_stats = None + if softmax_stats_tensor is not None: + kernel_softmax_stats = _ensure_workspace_tensor( + metadata, + "_fp4_mla_cutedsl_softmax_stats_buf", + (2, num_queries, physical_heads), + dtype=torch.float32, + device=output.device, + ) if num_heads < physical_heads: kernel_output = _ensure_workspace_tensor( metadata, @@ -1056,9 +1103,15 @@ def run_fp4_mla_attention_decode( v_page_offset=v_page_offset, q_batch_capacity=q_batch_capacity, partition_runtime_valid_k=bool(getattr(metadata, "is_cuda_graph", False)), + softmax_row_max=(None if kernel_softmax_stats is None else kernel_softmax_stats[0]), + softmax_row_sum=(None if kernel_softmax_stats is None else kernel_softmax_stats[1]), + helix_kv_bounds=helix_kv_bounds, ) if kernel_output is not output: output.copy_(kernel_output[:, :num_heads]) + if kernel_softmax_stats is not None: + softmax_stats_tensor[..., 0].copy_(kernel_softmax_stats[0, :, :num_heads]) + softmax_stats_tensor[..., 1].copy_(kernel_softmax_stats[1, :, :num_heads]) return total_p_rows = num_queries * max_pages * num_heads diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 2dbb74741e8c..e80c018415ac 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -982,6 +982,7 @@ def fused_fp4_mla_decode_ctm( sfb_ptr: cute.Pointer, page_table_ptr: cute.Pointer, valid_k_ptr: cute.Pointer, + helix_kv_bounds_ptr: cute.Pointer, c_ptr: cute.Pointer, accum_ptr: cute.Pointer, row_max_ptr: cute.Pointer, @@ -1005,9 +1006,11 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, + use_helix_kv_bounds: ctm.Constexpr = False, use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: n, k = problem_size m = runtime_m @@ -1046,6 +1049,10 @@ def fused_fp4_mla_decode_ctm( cute.recast_ptr(valid_k_ptr, dtype=cutlass.Int32), cute.make_layout((batch_size // query_len_per_seq,), stride=(1,)), ) + helix_kv_bounds_tensor = cute.make_tensor( + cute.recast_ptr(helix_kv_bounds_ptr, dtype=cutlass.Int32), + cute.make_layout((batch_size,), stride=(1,)), + ) v_tma_tensor = cute.make_tensor( b_ptr, cute.make_layout( @@ -1302,6 +1309,7 @@ def fused_fp4_mla_decode_ctm( page_table_tensor, page_indptr_tensor, valid_k_tensor, + helix_kv_bounds_tensor, q_global_scale_tensor, kv_global_scale_tensor, tma_q_desc, @@ -1329,10 +1337,12 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ).launch( grid=( cute.ceil_div(c_tensor.shape[0], SMEM_P4_CTA_GROUP_M) * CLUSTER_SHAPE_MNK[0], @@ -4542,6 +4552,7 @@ def _store_final_o_from_tmem( final_row_sum: ctm.Float32, final_stat_scale: ctm.Float32, output_normalizer: ctm.Float32, + write_softmax_stats: ctm.Constexpr = False, producer_warp_base: ctm.Constexpr = 0, ) -> None: gC_arr = ctm.make_array_view(mC_mnl) @@ -4561,6 +4572,10 @@ def _store_final_o_from_tmem( + local_row ) batch_offset = bidz * m * n + if cutlass.const_expr(write_softmax_stats): + if col_band == ctm.Int32(0) and bidy == ctm.Int32(0): + mRowMax_ml[row, bidz] = final_row_max * final_stat_scale + mRowSum_ml[row, bidz] = final_row_sum n_tile_total = n // ctm.Int32(OUT_DIM) subtile_cols: ctm.Constexpr = 32 subtiles_per_n_tile: ctm.Constexpr = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N // subtile_cols @@ -5822,6 +5837,7 @@ def _run_mla_decode_body( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_ptr, @@ -5849,10 +5865,12 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: pv_psf_rescale = ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) * pv_output_scale warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -5863,11 +5881,14 @@ def _run_mla_decode_body( page_batch = ctm.Int32(0) valid_k_for_l = ctm.Int32(0) page_batch = cute.arch.make_warp_uniform(bidz // ctm.Int32(query_len_per_seq)) - query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) - valid_k_for_l = ctm.max( - mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), - ctm.Int32(0), - ) + if cutlass.const_expr(use_helix_kv_bounds): + valid_k_for_l = mHelixKvBounds_l[bidz] + else: + query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) + valid_k_for_l = ctm.max( + mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), + ctm.Int32(0), + ) valid_k_for_l = cute.arch.make_warp_uniform(valid_k_for_l) csr_page_begin = ctm.Int32(0) csr_page_count = ctm.Int32(0) @@ -6818,8 +6839,19 @@ def _run_mla_decode_body( final_anchor_row_sum = sFinalAnchorRowSum[final_row_state_local_row] final_stat_scale = ctm.Float32(1.0) final_row_sum = final_anchor_row_sum + final_row_max = running_row_max + if cutlass.const_expr(write_softmax_stats): + final_stat_scale = softmax_scale_log2 / ctm.Float32(LOG2_E) + final_row_sum = final_anchor_row_sum * cute.exp2( + (running_row_anchor - running_row_max) * softmax_scale_log2, + fastmath=True, + ) + if valid_k_for_l == ctm.Int32(0): + final_stat_scale = ctm.Float32(1.0) + final_row_max = ctm.Float32(-ctm.Float32.inf) + final_row_sum = ctm.Float32(0.0) output_normalizer = ctm.Float32(0.0) - if final_anchor_row_sum != ctm.Float32(0.0): + if valid_k_for_l != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): output_normalizer = cute.arch.rcp_approx(final_anchor_row_sum) * pv_output_scale last_pv_li_idx = stream_li_total - ctm.Int32(1) last_pv_slot = last_pv_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -6845,10 +6877,11 @@ def _run_mla_decode_body( bidz, m, n, - final_row_max=running_row_max, + final_row_max=final_row_max, final_row_sum=final_row_sum, final_stat_scale=final_stat_scale, output_normalizer=output_normalizer, + write_softmax_stats=write_softmax_stats, producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, ) prims.barrier(barrier_id=O_STORE_BAR_ID, number_of_threads=O_STORE_BAR_THREADS) @@ -6867,6 +6900,7 @@ def kernel( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_desc: ctm.GridConstant[cuda_tma.TensorMap], @@ -6894,15 +6928,18 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, + write_softmax_stats: ctm.Constexpr, ) -> None: _run_mla_decode_body( mPageTable_pl, mPageIndptr_s, mValidK_l, + mHelixKvBounds_l, mQGlobalScale, mKvGlobalScale, tma_q_desc.get_ptr(), @@ -6930,10 +6967,12 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ) @@ -6954,6 +6993,7 @@ def _make_fused_ptrs( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -6973,6 +7013,7 @@ def _make_fused_ptrs( make_ptr(cutlass.Uint8, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=32), make_ptr(cutlass.Int32, page_table_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr(cutlass.Int32, valid_k_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(cutlass.Int32, helix_kv_bounds_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr( _cutlass_output_dtype(output_dtype), c_data_ptr, @@ -7162,6 +7203,7 @@ def _compile_fused( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7183,6 +7225,8 @@ def _compile_fused( ksf_page_stride_bytes: int = 0, vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, + write_softmax_stats: bool = False, + use_helix_kv_bounds: bool = False, ) -> Callable: if type(kv) is not int: raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") @@ -7209,6 +7253,8 @@ def _compile_fused( query_len_per_seq, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, + use_helix_kv_bounds, ) cached = _FUSED_COMPILE_CACHE.get(cache_key) if cached is not None: @@ -7223,6 +7269,7 @@ def _compile_fused( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, accum_data_ptr, row_max_data_ptr, @@ -7251,9 +7298,11 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, + use_helix_kv_bounds=use_helix_kv_bounds, use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, + write_softmax_stats=write_softmax_stats, options="--opt-level 2 --ptxas-options '--uumn'", ) _FUSED_COMPILE_CACHE[cache_key] = compiled @@ -7285,6 +7334,9 @@ def run_trtllm_fp4_mla_decode_page_native( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7331,6 +7383,24 @@ def run_trtllm_fp4_mla_decode_page_native( f"got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" ) physical_m, q_bytes, l_batch = q_internal.shape + if (softmax_row_max is None) != (softmax_row_sum is None): + raise ValueError("softmax_row_max and softmax_row_sum must be provided together") + write_softmax_stats = softmax_row_max is not None + if write_softmax_stats: + expected_stats_shape = (l_batch, physical_m) + for name, tensor in ( + ("softmax_row_max", softmax_row_max), + ("softmax_row_sum", softmax_row_sum), + ): + if ( + tensor.dtype != torch.float32 + or tensor.shape != expected_stats_shape + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous float32 with shape {expected_stats_shape}" + ) + _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) if l_batch <= 0 or l_batch > CUDA_GRID_Z_MAX: raise ValueError(f"queries must be in [1, {CUDA_GRID_Z_MAX}], got {l_batch}") if q_batch_capacity is None: @@ -7410,6 +7480,16 @@ def run_trtllm_fp4_mla_decode_page_native( f"got dtype={valid_k.dtype} shape={tuple(valid_k.shape)} " f"stride={valid_k.stride()}" ) + if helix_kv_bounds is not None and ( + helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.shape != (l_batch,) + or helix_kv_bounds.stride(0) != 1 + ): + raise ValueError( + "helix_kv_bounds must be contiguous int32 with shape " + f"[{l_batch}], got dtype={helix_kv_bounds.dtype} " + f"shape={tuple(helix_kv_bounds.shape)} stride={helix_kv_bounds.stride()}" + ) cache_layout = _validate_v_packed_cache_args( v_packed, kv_cache, @@ -7526,6 +7606,10 @@ def run_trtllm_fp4_mla_decode_page_native( q_global_scale, kv_global_scale, ) + if write_softmax_stats: + tensors += (softmax_row_max, softmax_row_sum) + if helix_kv_bounds is not None: + tensors += (helix_kv_bounds,) if device.type != "cuda" or any((tensor.device != device for tensor in tensors)): raise ValueError("all page-native decode tensors must share one CUDA device") if q_global_scale.dtype != torch.float32 or q_global_scale.numel() != 1: @@ -7545,9 +7629,14 @@ def run_trtllm_fp4_mla_decode_page_native( k_sf_data_ptr = sf_cache.data_ptr() b_data_ptr = v_packed.data_ptr() scratch_ptr = output.data_ptr() + row_max_data_ptr = softmax_row_max.data_ptr() if write_softmax_stats else scratch_ptr + row_sum_data_ptr = softmax_row_sum.data_ptr() if write_softmax_stats else scratch_ptr sfb_data_ptr = v_sf.data_ptr() page_table_data_ptr = src_page_ids.data_ptr() valid_k_data_ptr = valid_k.data_ptr() + helix_kv_bounds_data_ptr = ( + helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else valid_k_data_ptr + ) c_data_ptr = output.data_ptr() page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() q_global_scale_data_ptr = q_global_scale.data_ptr() @@ -7566,10 +7655,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -7587,6 +7677,8 @@ def run_trtllm_fp4_mla_decode_page_native( ksf_page_stride_bytes=ksf_page_stride_bytes, vsf_page_stride_bytes=vsf_page_stride_bytes, use_consecutive_page_pair=use_consecutive_page_pair, + write_softmax_stats=write_softmax_stats, + use_helix_kv_bounds=helix_kv_bounds is not None, ) supports_prepared = _class_defines_callables( fused, "to", "generate_execution_args", "run_compiled_program" @@ -7604,7 +7696,10 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -7637,10 +7732,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -7670,7 +7766,8 @@ def run_trtllm_fp4_mla_decode_page_native( if executor is None: if torch.cuda.is_current_stream_capturing(): raise RuntimeError( - "FP4 MLA CUDA Graph capture requires an eager warmup for the compiled kernel, device, and CUDA context." + "FP4 MLA CUDA Graph capture requires an eager warmup for the " + "compiled kernel, device, and CUDA context." ) candidate = fused.to(device_index) if not _class_defines_callables( @@ -7718,6 +7815,9 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7769,4 +7869,7 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles=assume_consecutive_page_prefix_tiles, partition_runtime_valid_k=partition_runtime_valid_k, enable_mxi_imlp=enable_mxi_imlp, + softmax_row_max=softmax_row_max, + softmax_row_sum=softmax_row_sum, + helix_kv_bounds=helix_kv_bounds, ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py index 6e6ded7d99be..03f360e80605 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16_fused_v_transpose.py @@ -976,6 +976,7 @@ def fused_fp4_mla_decode_ctm( sfb_ptr: cute.Pointer, page_table_ptr: cute.Pointer, valid_k_ptr: cute.Pointer, + helix_kv_bounds_ptr: cute.Pointer, c_ptr: cute.Pointer, accum_ptr: cute.Pointer, row_max_ptr: cute.Pointer, @@ -999,9 +1000,11 @@ def fused_fp4_mla_decode_ctm( page_size: ctm.Constexpr = KV_TILE, use_mixed_imlp: ctm.Constexpr = False, query_len_per_seq: ctm.Constexpr = 1, + use_helix_kv_bounds: ctm.Constexpr = False, use_smem_page_plan: ctm.Constexpr = True, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: n, k = problem_size m = runtime_m @@ -1048,6 +1051,10 @@ def fused_fp4_mla_decode_ctm( cute.recast_ptr(valid_k_ptr, dtype=cutlass.Int32), cute.make_layout((l // query_len_per_seq,), stride=(1,)), ) + helix_kv_bounds_tensor = cute.make_tensor( + cute.recast_ptr(helix_kv_bounds_ptr, dtype=cutlass.Int32), + cute.make_layout((l,), stride=(1,)), + ) v_tma_tensor = cute.make_tensor( b_ptr, cute.make_layout( @@ -1310,6 +1317,7 @@ def fused_fp4_mla_decode_ctm( page_table_tensor, page_indptr_tensor, valid_k_tensor, + helix_kv_bounds_tensor, q_global_scale_tensor, kv_global_scale_tensor, tma_q_desc, @@ -1338,10 +1346,12 @@ def fused_fp4_mla_decode_ctm( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ).launch( grid=( cute.ceil_div(c_tensor.shape[0], SMEM_P4_CTA_GROUP_M) * CLUSTER_SHAPE_MNK[0], @@ -4978,6 +4988,7 @@ def _store_final_o_from_tmem( final_row_sum: ctm.Float32, final_stat_scale: ctm.Float32, output_normalizer: ctm.Float32, + write_softmax_stats: ctm.Constexpr = False, producer_warp_base: ctm.Constexpr = 0, ) -> None: gC_arr = ctm.make_array_view(mC_mnl) @@ -4999,6 +5010,10 @@ def _store_final_o_from_tmem( + local_row ) batch_offset = bidz * m * n + if cutlass.const_expr(write_softmax_stats): + if col_band == ctm.Int32(0) and bidy == ctm.Int32(0): + mRowMax_ml[row, bidz] = final_row_max * final_stat_scale + mRowSum_ml[row, bidz] = final_row_sum n_tile_total = n // ctm.Int32(OUT_DIM) subtile_cols: ctm.Constexpr = 32 subtiles_per_n_tile: ctm.Constexpr = SMEM_P4_BMM2_N // SMEM_P4_TMEM_WARP_N // subtile_cols @@ -6282,6 +6297,7 @@ def _run_mla_decode_body( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_ptr, @@ -6310,10 +6326,12 @@ def _run_mla_decode_body( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr = False, use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, + write_softmax_stats: ctm.Constexpr = False, ) -> None: pv_psf_rescale = ctm.Float32(FP4_MLA_P_GLOBAL_SCALE) * pv_output_scale warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -6324,11 +6342,14 @@ def _run_mla_decode_body( page_batch = ctm.Int32(0) valid_k_for_l = ctm.Int32(0) page_batch = cute.arch.make_warp_uniform(bidz // ctm.Int32(query_len_per_seq)) - query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) - valid_k_for_l = ctm.max( - mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), - ctm.Int32(0), - ) + if cutlass.const_expr(use_helix_kv_bounds): + valid_k_for_l = mHelixKvBounds_l[bidz] + else: + query_offset = bidz - page_batch * ctm.Int32(query_len_per_seq) + valid_k_for_l = ctm.max( + mValidK_l[page_batch] - (ctm.Int32(query_len_per_seq - 1) - query_offset), + ctm.Int32(0), + ) valid_k_for_l = cute.arch.make_warp_uniform(valid_k_for_l) csr_page_begin = ctm.Int32(0) csr_page_count = ctm.Int32(0) @@ -7317,8 +7338,19 @@ def _run_mla_decode_body( final_anchor_row_sum = sFinalAnchorRowSum[final_row_state_local_row] final_stat_scale = ctm.Float32(1.0) final_row_sum = final_anchor_row_sum + final_row_max = running_row_max + if cutlass.const_expr(write_softmax_stats): + final_stat_scale = softmax_scale_log2 / ctm.Float32(LOG2_E) + final_row_sum = final_anchor_row_sum * cute.exp2( + (running_row_anchor - running_row_max) * softmax_scale_log2, + fastmath=True, + ) + if valid_k_for_l == ctm.Int32(0): + final_stat_scale = ctm.Float32(1.0) + final_row_max = ctm.Float32(-ctm.Float32.inf) + final_row_sum = ctm.Float32(0.0) output_normalizer = ctm.Float32(0.0) - if final_anchor_row_sum != ctm.Float32(0.0): + if valid_k_for_l != ctm.Int32(0) and final_anchor_row_sum != ctm.Float32(0.0): output_normalizer = cute.arch.rcp_approx(final_anchor_row_sum) * pv_output_scale last_pv_li_idx = stream_li_total - ctm.Int32(1) last_pv_slot = last_pv_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) @@ -7344,10 +7376,11 @@ def _run_mla_decode_body( bidz, m, n, - final_row_max=running_row_max, + final_row_max=final_row_max, final_row_sum=final_row_sum, final_stat_scale=final_stat_scale, output_normalizer=output_normalizer, + write_softmax_stats=write_softmax_stats, producer_warp_base=SMEM_P4_CORRECTION_WARP_ID_BEGIN, ) prims.barrier(barrier_id=O_STORE_BAR_ID, number_of_threads=O_STORE_BAR_THREADS) @@ -7366,6 +7399,7 @@ def kernel( mPageTable_pl: cute.Tensor, mPageIndptr_s: cute.Tensor, mValidK_l: cute.Tensor, + mHelixKvBounds_l: cute.Tensor, mQGlobalScale: cute.Tensor, mKvGlobalScale: cute.Tensor, tma_q_desc: ctm.GridConstant[cuda_tma.TensorMap], @@ -7394,15 +7428,18 @@ def kernel( pv_output_scale: ctm.Float32, page_size: ctm.Constexpr, query_len_per_seq: ctm.Constexpr, + use_helix_kv_bounds: ctm.Constexpr, use_smem_page_plan: ctm.Constexpr, use_mixed_imlp: ctm.Constexpr, use_consecutive_page_pair: ctm.Constexpr, use_ksf_gather4: ctm.Constexpr, + write_softmax_stats: ctm.Constexpr, ) -> None: _run_mla_decode_body( mPageTable_pl, mPageIndptr_s, mValidK_l, + mHelixKvBounds_l, mQGlobalScale, mKvGlobalScale, tma_q_desc.get_ptr(), @@ -7431,10 +7468,12 @@ def kernel( pv_output_scale, page_size, query_len_per_seq, + use_helix_kv_bounds, use_smem_page_plan, use_mixed_imlp, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, ) @@ -7455,6 +7494,7 @@ def _make_fused_ptrs( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7474,6 +7514,7 @@ def _make_fused_ptrs( make_ptr(cutlass.Uint8, sfb_data_ptr, cute.AddressSpace.gmem, assumed_align=32), make_ptr(cutlass.Int32, page_table_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr(cutlass.Int32, valid_k_data_ptr, cute.AddressSpace.gmem, assumed_align=4), + make_ptr(cutlass.Int32, helix_kv_bounds_data_ptr, cute.AddressSpace.gmem, assumed_align=4), make_ptr( _cutlass_output_dtype(output_dtype), c_data_ptr, @@ -7609,6 +7650,7 @@ def _compile_fused( sfb_data_ptr: int, page_table_data_ptr: int, valid_k_data_ptr: int, + helix_kv_bounds_data_ptr: int, c_data_ptr: int, accum_data_ptr: int, row_max_data_ptr: int, @@ -7630,6 +7672,8 @@ def _compile_fused( ksf_page_stride_bytes: int = 0, vsf_page_stride_bytes: int = 0, use_consecutive_page_pair: bool = False, + write_softmax_stats: bool = False, + use_helix_kv_bounds: bool = False, ) -> Callable: if type(kv) is not int: raise TypeError(f"runtime-KV compile K must be an int, got {type(kv).__name__}") @@ -7654,6 +7698,8 @@ def _compile_fused( query_len_per_seq, use_consecutive_page_pair, use_ksf_gather4, + write_softmax_stats, + use_helix_kv_bounds, ) cached = _FUSED_COMPILE_CACHE.get(cache_key) if cached is not None: @@ -7668,6 +7714,7 @@ def _compile_fused( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, accum_data_ptr, row_max_data_ptr, @@ -7696,9 +7743,11 @@ def _compile_fused( page_size=page_size, use_mixed_imlp=use_mixed_imlp, query_len_per_seq=query_len_per_seq, + use_helix_kv_bounds=use_helix_kv_bounds, use_smem_page_plan=kv == SMEM_P4_PAGE_PLAN_PROFILE_KV, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, + write_softmax_stats=write_softmax_stats, options="--opt-level 2 --ptxas-options '--uumn'", ) _FUSED_COMPILE_CACHE[cache_key] = compiled @@ -7730,6 +7779,9 @@ def run_trtllm_fp4_mla_decode_page_native( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -7775,6 +7827,24 @@ def run_trtllm_fp4_mla_decode_page_native( f"q_internal must be a uint8 [M, Q640/2, L] tensor, got dtype={q_internal.dtype} shape={tuple(q_internal.shape)}" ) physical_m, q_bytes, l_batch = q_internal.shape + if (softmax_row_max is None) != (softmax_row_sum is None): + raise ValueError("softmax_row_max and softmax_row_sum must be provided together") + write_softmax_stats = softmax_row_max is not None + if write_softmax_stats: + expected_stats_shape = (l_batch, physical_m) + for name, tensor in ( + ("softmax_row_max", softmax_row_max), + ("softmax_row_sum", softmax_row_sum), + ): + if ( + tensor.dtype != torch.float32 + or tensor.shape != expected_stats_shape + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous float32 with shape {expected_stats_shape}" + ) + _validate_tensor_pointer_alignment(name, tensor, alignment_bytes=32) if l_batch <= 0 or l_batch > CUDA_GRID_Z_MAX: raise ValueError(f"queries must be in [1, {CUDA_GRID_Z_MAX}], got {l_batch}") if q_batch_capacity is None: @@ -7847,6 +7917,16 @@ def run_trtllm_fp4_mla_decode_page_native( raise ValueError( f"valid_k must be contiguous int32 [{num_sequences}], got dtype={valid_k.dtype} shape={tuple(valid_k.shape)} stride={valid_k.stride()}" ) + if helix_kv_bounds is not None and ( + helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.shape != (l_batch,) + or helix_kv_bounds.stride(0) != 1 + ): + raise ValueError( + "helix_kv_bounds must be contiguous int32 with shape " + f"[{l_batch}], got dtype={helix_kv_bounds.dtype} " + f"shape={tuple(helix_kv_bounds.shape)} stride={helix_kv_bounds.stride()}" + ) cache_layout = _kv_cache_3d_layout(kv_cache, page_size) _validate_tensor_pointer_alignment("kv_cache", kv_cache, alignment_bytes=16) if ( @@ -7938,6 +8018,10 @@ def run_trtllm_fp4_mla_decode_page_native( q_global_scale, kv_global_scale, ) + if write_softmax_stats: + tensors += (softmax_row_max, softmax_row_sum) + if helix_kv_bounds is not None: + tensors += (helix_kv_bounds,) if device.type != "cuda" or any((tensor.device != device for tensor in tensors)): raise ValueError("all page-native decode tensors must share one CUDA device") if q_global_scale.dtype != torch.float32 or q_global_scale.numel() != 1: @@ -7959,9 +8043,14 @@ def run_trtllm_fp4_mla_decode_page_native( # canonical KV pointer so the compiled call carries no sidecar allocation. b_data_ptr = kv_cache.data_ptr() scratch_ptr = output.data_ptr() + row_max_data_ptr = softmax_row_max.data_ptr() if write_softmax_stats else scratch_ptr + row_sum_data_ptr = softmax_row_sum.data_ptr() if write_softmax_stats else scratch_ptr sfb_data_ptr = v_sf.data_ptr() page_table_data_ptr = src_page_ids.data_ptr() valid_k_data_ptr = valid_k.data_ptr() + helix_kv_bounds_data_ptr = ( + helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else valid_k_data_ptr + ) c_data_ptr = output.data_ptr() page_indptr_data_ptr = paged_kv_indptr_decode.data_ptr() q_global_scale_data_ptr = q_global_scale.data_ptr() @@ -7980,10 +8069,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -8001,6 +8091,8 @@ def run_trtllm_fp4_mla_decode_page_native( ksf_page_stride_bytes=ksf_page_stride_bytes, vsf_page_stride_bytes=vsf_page_stride_bytes, use_consecutive_page_pair=use_consecutive_page_pair, + write_softmax_stats=write_softmax_stats, + use_helix_kv_bounds=helix_kv_bounds is not None, ) supports_prepared = _class_defines_callables( fused, "to", "generate_execution_args", "run_compiled_program" @@ -8018,7 +8110,10 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -8051,10 +8146,11 @@ def run_trtllm_fp4_mla_decode_page_native( sfb_data_ptr, page_table_data_ptr, valid_k_data_ptr, + helix_kv_bounds_data_ptr, c_data_ptr, scratch_ptr, - scratch_ptr, - scratch_ptr, + row_max_data_ptr, + row_sum_data_ptr, page_indptr_data_ptr, q_global_scale_data_ptr, kv_global_scale_data_ptr, @@ -8084,7 +8180,8 @@ def run_trtllm_fp4_mla_decode_page_native( if executor is None: if torch.cuda.is_current_stream_capturing(): raise RuntimeError( - "FP4 MLA CUDA Graph capture requires an eager warmup for the compiled kernel, device, and CUDA context." + "FP4 MLA CUDA Graph capture requires an eager warmup for the " + "compiled kernel, device, and CUDA context." ) candidate = fused.to(device_index) if not _class_defines_callables( @@ -8132,6 +8229,9 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles: int = 0, partition_runtime_valid_k: bool = False, enable_mxi_imlp: bool = True, + softmax_row_max: torch.Tensor | None = None, + softmax_row_sum: torch.Tensor | None = None, + helix_kv_bounds: torch.Tensor | None = None, ) -> None: if type(v_pack_block) is not int: raise TypeError(f"v_pack_block must be an int, got {type(v_pack_block).__name__}") @@ -8183,4 +8283,7 @@ def run_trtllm_fp4_mla_decode_page_native_from_raw( assume_consecutive_page_prefix_tiles=assume_consecutive_page_prefix_tiles, partition_runtime_valid_k=partition_runtime_valid_k, enable_mxi_imlp=enable_mxi_imlp, + softmax_row_max=softmax_row_max, + softmax_row_sum=softmax_row_sum, + helix_kv_bounds=helix_kv_bounds, ) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py index 3527338f8bc6..3c149fd35207 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py @@ -427,7 +427,6 @@ def configure_fp4_mla_device_page_table( and int(getattr(metadata, "beam_width", 1)) == 1 and not bool(getattr(metadata, "is_spec_dec_tree", False)) and not bool(getattr(metadata, "locality_domain_enabled", False)) - and not bool(getattr(metadata, "enable_helix", False)) and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK and max_page_capacity > 0 and page_index_scale > 0 diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py index a377346dab9e..db53f3e69663 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py @@ -571,10 +571,14 @@ def rebuild_fp4_mla_disagg_imported_cache( or not callable(getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None)) ): return False - if not isinstance(prompt_len, int) or prompt_len <= 0: + if not isinstance(prompt_len, int) or prompt_len < 0: raise ValueError( - f"FP4 MLA disaggregated import needs a positive prompt_len, got {prompt_len}." + f"FP4 MLA disaggregated import needs a nonnegative prompt_len, got {prompt_len}." ) + # Helix assigns whole pages round-robin, so a rank may own no prompt pages. + # There are no process-local V sidecars to rebuild on that rank. + if prompt_len == 0: + return True page_size = int(kv_cache_manager.tokens_per_block) if page_size != FP4_MLA_TOKENS_PER_BLOCK: diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 4ccb63c37581..44f4ef79e8a0 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -30,7 +30,6 @@ from ...speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned -from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType from tensorrt_llm.logger import logger @@ -224,9 +223,6 @@ def effective_beam_width(self) -> int: draft_block_ids_per_seq: Optional[torch.Tensor] = None draft_kv_block_ids_per_seq: Optional[torch.Tensor] = None - # True during warmup forward passes (dummy requests, no real data). - is_warmup: bool = False - # Batch-shared FP4 state; other attention paths allocate none of it. fp4_mla_state: Optional[Fp4MlaState] = field(init=False, default=None, @@ -583,9 +579,9 @@ def _post_init_with_buffers(self, buffers) -> None: pin_memory=prefer_pinned(), ) - if (self.kv_cache_manager is not None - and self.kv_cache_manager.kv_factor == 1 - and self.kv_cache_manager.dtype == DataType.NVFP4): + if callable( + getattr(self.kv_cache_manager, "get_fp4_mla_page_table_spec", + None)): self.fp4_mla_state = Fp4MlaState.create(self, buffers) # Allocate static buffers for helix parallelism support. @@ -1695,6 +1691,12 @@ def __init__( if not skip_create_weights_in_init: self.update_quant_config(self.quant_config) + @property + def uses_fp4_mla_attention(self) -> bool: + """Whether this layer executes dense FP4 MLA, not sparse NVFP4 storage.""" + return (self.is_mla_enable and self.has_fp4_kv_cache + and self.sparse_params is None) + def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.quant_config = new_quant_config or QuantConfig() self.quant_mode = int(self.quant_config.layer_quant_mode) @@ -2622,7 +2624,7 @@ def mla_rope_generation( # kernel reads it. self._ensure_rope_table_size(metadata.max_seq_len) - if self.has_fp4_kv_cache: + if self.uses_fp4_mla_attention: self._fp4_mla_rope_generation( fused_q, q_pe, @@ -2744,6 +2746,8 @@ def _fp4_mla_rope_generation( q_pe=q_pe, q_rope_out=fused_q[..., self.kv_lora_rank:], q_quant_input=fused_q, + helix_position_offsets=metadata.helix_position_offsets, + helix_is_inactive_rank=metadata.helix_is_inactive_rank, ) if not hp_pool_updated: raise RuntimeError( @@ -2757,5 +2761,5 @@ def can_fuse_fp4_mla_q_quant( metadata: TrtllmAttentionMetadata, ) -> bool: return bool( - self.has_fp4_kv_cache + self.uses_fp4_mla_attention and can_fuse_fp4_mla_q_quant(metadata, fused_q, q_pe, latent_cache)) diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index 49903165372d..2445c8aa6bdd 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -698,7 +698,8 @@ def create_weights(self): ) mla_weight_dtype = torch.float8_e4m3fn if has_fp8_block_scales else self.dtype if ( - self.mqa.support_fp4_kv_cache() + self.sparse_params is None + and self.mqa.support_fp4_kv_cache() and self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp4_kv_cache() ): @@ -1598,7 +1599,8 @@ def forward_absorption_generation( ) fp4_mla = ( - self.mqa.support_fp4_kv_cache() + self.sparse_params is None + and self.mqa.support_fp4_kv_cache() and self.quant_config is not None and self.quant_config.layer_quant_mode.has_fp4_kv_cache() ) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 2cbe479c6254..e08afe5f4632 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -189,7 +189,8 @@ def get_kv_cache_manager_cls( sparse_attn_config = model_config.sparse_attention_config sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) quant_config = getattr(model_config, "quant_config", None) - if (is_mla(config) and quant_config is not None + if (sparse_attn_config is None and is_mla(config) + and quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache()): if kv_cache_config.use_kv_cache_manager_v2 is False: raise ValueError("FP4 MLA requires use_kv_cache_manager_v2=True.") @@ -203,13 +204,6 @@ def get_kv_cache_manager_cls( raise NotImplementedError( "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " "support hybrid linear-attention models.") - if sparse_attn_config is not None: - sparse_attn_algorithm = (sparse_attn_algorithm - or type(sparse_attn_config).__name__) - raise NotImplementedError( - "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " - f"support sparse attention algorithm {sparse_attn_algorithm!r}." - ) from ..attention.backends.fp4_mla.cache_manager import \ Fp4MlaKVCacheManagerV2 @@ -913,7 +907,8 @@ def _validate_or_fallback_kv_cache_manager_v2( f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") quant_config = getattr(model_config, "quant_config", None) - if (is_mla(config) and quant_config is not None + if (sparse_attn_config is None and is_mla(config) + and quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache()): raise NotImplementedError( "FP4 MLA requires Fp4MlaKVCacheManagerV2, which is " diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 3d4984140c1b..69392781064c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -212,13 +212,13 @@ def validate_and_set_kv_cache_quant(model_config: ModelConfig, def validate_fp4_mla_config(model_config: ModelConfig, llm_args: TorchLlmArgs) -> None: """Validate FP4 MLA before model construction and KV-cache allocation.""" - if not (is_mla(model_config.pretrained_config) + if not (model_config.sparse_attention_config is None + and is_mla(model_config.pretrained_config) and model_config.quant_config.quant_mode.has_fp4_kv_cache()): return if not supports_fp4_mla_attention(model_config): - raise ValueError( - "FP4 MLA requires the TRTLLM attention backend with dense MLA; " - "sparse and hybrid linear attention are not supported.") + raise ValueError("Dense FP4 MLA requires the TRTLLM attention backend; " + "hybrid linear attention is not supported.") if model_config.mapping.cp_size != 1: raise ValueError("FP4 MLA does not support context parallelism.") if llm_args.kv_cache_config.use_kv_cache_manager_v2 is False: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index fe97e51ae823..c755dda60358 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -398,10 +398,6 @@ def __init__( self.mapping = mapping self.dtype = dtype self.kv_cache_type = kv_cache_type - if dtype == DataType.NVFP4 and kv_cache_type == CacheTypeCpp.SELFKONLY: - raise ValueError( - "NVFP4 SELFKONLY cache storage requires " - "Fp4MlaKVCacheManagerV2; KVCacheManager V1 is not supported.") # Consumed by the disaggregation page-table builder to expose the DSA # indexer K cache pool as a REPLICATED pool view. self.enable_indexer_k_cache = enable_indexer_k_cache diff --git a/tests/microbenchmarks/bench_fp4_mla_decode.py b/tests/microbenchmarks/bench_fp4_mla_decode.py index ad4e50a795f5..98c957475d99 100644 --- a/tests/microbenchmarks/bench_fp4_mla_decode.py +++ b/tests/microbenchmarks/bench_fp4_mla_decode.py @@ -204,7 +204,6 @@ def _build_multi_seq_metadata( request_ids=request_ids, runtime_features=SimpleNamespace(has_speculative_draft_tokens=False), is_cuda_graph=False, - is_warmup=False, fp4_mla_state=Fp4MlaState( batch_indices=batch_indices, positions=positions, diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index ae9f332dd670..0715202baa3c 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -737,7 +737,6 @@ def _build_multi_seq_metadata( request_ids=request_ids, runtime_features=SimpleNamespace(has_speculative_draft_tokens=False), is_cuda_graph=False, - is_warmup=False, fp4_mla_state=Fp4MlaState( batch_indices=batch_indices, positions=positions, @@ -1174,10 +1173,10 @@ def _fp4_mla_attention_decode_reference( ) p_dequant = None - if hasattr(metadata, "_fp4_mla_attention_p_buf"): + if "_fp4_mla_attention_p_buf" in metadata.fp4_mla_state.workspaces: p_dequant = _dequant_fp4_swizzled( - metadata._fp4_mla_attention_p_buf, - metadata._fp4_mla_attention_p_sf_buf, + metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_buf"], + metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_sf_buf"], logical_dim=metadata.page_size, sf_per_token=metadata.page_size // FP4_BLOCK_SIZE, global_scale=FP4_MLA_P_GLOBAL_SCALE,