diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index aca96dcd2d47..a89c39e5a73f 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 @@ -427,6 +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 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/__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/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index b2703f364708..9b47e4a4b806 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 @@ -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, _build_fp8_mla_context_metadata, @@ -39,7 +38,7 @@ ) from tensorrt_llm.bindings import DataType -from .fallback import FallbackFmha +from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha if TYPE_CHECKING: @@ -52,49 +51,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.uses_fp4_mla_attention + return bool(getattr(attn, "uses_fp4_mla_attention", False)) - 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: + # 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.") 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_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 - ): + sparse_params = self.attn.sparse_params + if getattr(sparse_params, "uses_spcompress", False): raise NotImplementedError("FP4 MLA does not support sparse attention.") kv_cache_manager = metadata.kv_cache_manager @@ -104,39 +91,89 @@ 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.") + 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.") + 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 - k = params.k_input - v = params.v_input - output = params.context_buf + q = params.query_input + k = params.key_input + v = params.value_input + 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: 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.") @@ -148,18 +185,12 @@ 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) 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 - attn._ensure_rope_table_size(metadata.max_seq_len) latent_cache = forward_args.latent_cache[:num_tokens] def update_fp4_cache() -> None: @@ -176,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 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) - - 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 + 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, @@ -230,7 +238,7 @@ def forward_context_partition( *, kv_lens_cuda: torch.Tensor, kv_lens_cpu: torch.Tensor, - ) -> Tuple[torch.Tensor, Optional[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 @@ -281,33 +289,7 @@ def forward_context_partition( 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 + 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: @@ -339,8 +321,8 @@ 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: @@ -355,18 +337,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, @@ -383,9 +365,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..2d1bd5c86742 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,8 @@ def is_available(cls, attn: "TrtllmAttention") -> bool: f"{cls.__name__} is unavailable: skip-correction is enabled and unsupported." ) return False + if getattr(attn, "uses_fp4_mla_attention", False) and not cls.supports_fp4_mla: + return False return cls._is_available(attn) @classmethod 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..a2a0389b6204 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/__init__.py @@ -1,4362 +1,177 @@ # 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, prefer_pinned -from tensorrt_llm.bindings import DataType - -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, +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, ) - -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" -# 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) -_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 V2 cache-layout 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. - - 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. - """ - 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 - - 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, "_paged_kv_indices", None) - paged_kv_indptr = getattr(metadata, "_paged_kv_indptr", None) - paged_kv_indptr_decode = getattr(metadata, "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 - 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 - ) - # 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) - 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 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_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_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) - - 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 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 = ( - 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_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_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 - 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_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)): - 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) - 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_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) - 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._paged_kv_indptr, - metadata.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_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_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._paged_kv_indptr, - metadata.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_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_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) - - -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 - - -def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata, "_fp4_mla_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_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 = int(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], - q_context: Optional[torch.Tensor], - q_nope_head_dim: Optional[int], - *, - 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 - - 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.") - 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_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 + q_head_blocks, - ) - ]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - latent_cache, - q_context_view, - global_scale, - 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], - 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), - 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, - 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, - 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, - ) - 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.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_generation_kv_lens", None) - generation_lens = getattr(metadata, "fp4_mla_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_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 - ) - 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_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 - if record_for_capture: - metadata._fp4_mla_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_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)): - 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, - 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 - 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, "high_precision_kv_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 - 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 - 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, - helix_position_offsets, - helix_local_slots, - helix_is_inactive_rank, - page_ids, - hp_page_ids, - metadata.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, - 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, - 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, - use_helix, - use_helix_local_slots, - 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, "_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) - if not isinstance(context_indptr, torch.Tensor): - context_indptr = metadata.paged_kv_indptr_decode - context_batch_indices = getattr(metadata, "batch_indices", None) - if not isinstance(context_batch_indices, torch.Tensor): - context_batch_indices = hp_page_ids - context_positions = getattr(metadata, "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], - indptr_len=metadata.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 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, - 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, - 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. - - 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 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. - 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_prequantized_q = None - metadata._fp4_mla_prequantized_q_sf = None - 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] - 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.batch_indices.shape[0] - positions_len = metadata.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_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.") - 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, " - "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_prequantized_q = q_fp4_out - metadata._fp4_mla_prequantized_q_sf = q_sf_out - metadata._fp4_mla_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, - q_context, - q_nope_head_dim, - 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.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, - 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 - ) - 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.paged_kv_indices[: metadata.num_context_blocks] - cutedsl_repack_page_indptr = metadata.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 - ) - 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_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: - tensor = getattr(metadata, attr_name, None) - 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) - setattr(metadata, 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) - ) - 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, - ), - ) - - -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 getattr(metadata, valid_attr, None) == _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 = getattr(metadata, _triton_v_packed_attr(layer_idx), None) - 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 - - 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( - 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_device_page_table", False): - raise RuntimeError("FP4 MLA generation requires fixed-stride device page metadata.") - max_pages = int(metadata.fp4_mla_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, "_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_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, "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_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)) - 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.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 = getattr(metadata, "_fp4_mla_attention_group_max_buf", None) - gsum = getattr(metadata, "_fp4_mla_attention_group_sum_buf", None) - 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.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 = getattr(metadata, "_fp4_mla_attention_pv_partial_buf", None) - 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.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.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.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.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, - softmax_stats_tensor: Optional[torch.Tensor] = None, -) -> 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() - 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." - ) - - 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_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 - 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, - "_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.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)), - 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 - 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_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 _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 +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", + "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_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 55983fd1572c..5aa8429afe24 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -5,7 +5,7 @@ import math from dataclasses import dataclass, replace from types import MethodType -from typing import Iterable, List, Optional +from typing import TYPE_CHECKING, Iterable, List, Optional import torch @@ -15,15 +15,21 @@ 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, ) -from . import ( +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + +from .config import ( _FP4_MLA_CUTEDSL_BACKEND, _FP4_MLA_K_RESIDUAL_BACKENDS, FP4_BLOCK_SIZE, @@ -32,8 +38,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) @@ -172,6 +178,7 @@ 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", @@ -224,6 +231,11 @@ def finalize_manager(manager) -> None: 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) @@ -343,7 +355,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: 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) @@ -390,7 +402,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) bf16_local_layers = self._bf16_mla_local_layer_indices() fp4_local_layers = self._fp4_mla_local_layer_indices() @@ -905,7 +917,7 @@ def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: raise ValueError(f"Unknown FP4 MLA V2 internal layer {internal_layer}.") return result - def _get_runtime_cache_size_layer_components(self): + 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 = [ @@ -935,7 +947,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() @@ -946,7 +967,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) @@ -969,10 +989,9 @@ def get_cache_size_per_token(model_config, mapping, num_layers=None, **kwargs): 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 - max_batch_size = int(kwargs.get("max_batch_size") or 0) + max_batch_size = int(max_batch_size or 0) cache_bytes = ( per_layer * num_fp4_layers + logical_head_dim 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..d2fd2f7da0ce --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py @@ -0,0 +1,1063 @@ +# 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, + _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], + q_context: Optional[torch.Tensor], + q_nope_head_dim: Optional[int], + *, + 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 + + 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.") + 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 + q_head_blocks, + ) + ]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + latent_cache, + q_context_view, + 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), + 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, + 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, + 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, + ) + 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, + 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 + 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 + 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) + 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, + helix_position_offsets, + helix_local_slots, + helix_is_inactive_rank, + 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, + 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, + 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, + ) + + 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, + 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. + + 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 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. + 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.") + 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] + 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.") + 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, " + "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, + q_context, + q_nope_head_dim, + 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, + 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 + ) + 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..e11a0e654faf --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py @@ -0,0 +1,199 @@ +# 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_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 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..424a1c95c5b9 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py @@ -0,0 +1,1183 @@ +# 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, + softmax_stats_tensor: torch.Tensor | None = None, +) -> 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() + 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." + ) + + 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 + 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, + "_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)), + 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 + 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 b0c05229c288..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 @@ -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 @@ -34,7 +35,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: @@ -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) @@ -84,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 @@ -299,18 +309,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.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 = [] 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 @@ -327,7 +338,7 @@ def _build_fp8_mla_context_metadata( 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.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 @@ -343,7 +354,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 @@ -352,11 +363,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/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 7187a8601215..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 @@ -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) @@ -498,43 +449,6 @@ 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 - - -@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 - - 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) @dataclass(frozen=True) @@ -7702,177 +7616,177 @@ def run_trtllm_fp4_mla_decode_page_native( 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() + 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: + 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( 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..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 @@ -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) @@ -500,43 +451,6 @@ 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 - - -@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 - - 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) @dataclass(frozen=True) @@ -8114,179 +8028,179 @@ def run_trtllm_fp4_mla_decode_page_native( 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() + 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: + 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( 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..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,107 +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() - - -_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 - - 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) - @dataclass(frozen=True) class KvCache3DLayout: @@ -801,44 +707,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_triton.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_triton.py index 38814cc4d902..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 @@ -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 @@ -491,7 +476,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 +575,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/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..3c149fd35207 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py @@ -0,0 +1,885 @@ +# 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 V2 cache-layout 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. + + 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. + """ + 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 + 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 + ) + # 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) + 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 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 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 = ( + 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 new file mode 100644 index 000000000000..49fc36d9cc6e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py @@ -0,0 +1,204 @@ +# 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 .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 + 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() + 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] != len(fp4_layers) + or hp_pool.shape[2] != manager.kv_factor + 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, fp4_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: + # 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( + "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/fp4_mla/v_cache.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py new file mode 100644 index 000000000000..db53f3e69663 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py @@ -0,0 +1,663 @@ +# 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 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 .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: + 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 + + 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( + 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 _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 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 7ab975a02a33..44f4ef79e8a0 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 @@ -40,6 +40,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 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, @@ -144,6 +146,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 @@ -218,6 +223,12 @@ 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 + # 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. flash_mla_tile_scheduler_metadata: Optional[torch.Tensor] = None @@ -240,10 +251,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 @@ -301,6 +308,21 @@ 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. + """ + 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 host_kv_cache_pool_pointers(self) -> Optional[torch.Tensor]: """ @@ -557,6 +579,11 @@ def _post_init_with_buffers(self, buffers) -> None: pin_memory=prefer_pinned(), ) + 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. if self.enable_helix: self.helix_position_offsets = self.get_empty( @@ -620,19 +647,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() + 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() + 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 @@ -640,6 +669,11 @@ 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: + super().restore_from_spec_dec() + if self.fp4_mla_state is not None: + self.fp4_mla_state.restore_from_spec_dec(self) + def update_helix_param( self, helix_position_offsets: List[int], @@ -830,9 +864,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. - self._fp4_mla_fp8_context_state = None super().prepare() # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers. self._invalidate_mla_scheduler_buffers() @@ -977,14 +1008,28 @@ 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.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. + 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.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, context_kv_tokens: int, generation_kv_tokens: int, @@ -1646,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) @@ -1857,8 +1908,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 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.rotary_inv_freq, self.rotary_cos_sin = ( self.rope_params.create_rope_const_params()) @@ -2300,6 +2358,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, @@ -2531,6 +2593,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 +2614,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 +2624,16 @@ def mla_rope_generation( # kernel reads it. self._ensure_rope_table_size(metadata.max_seq_len) + if self.uses_fp4_mla_attention: + 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 +2692,74 @@ 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.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( + "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.") + + 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_fp4_mla_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, + helix_position_offsets=metadata.helix_position_offsets, + helix_is_inactive_rank=metadata.helix_is_inactive_rank, + ) + if not hp_pool_updated: + raise RuntimeError( + "Fused FP4 MLA RoPE/cache scatter did not update the HP pool.") + + 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.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 e331d9ebcf24..2445c8aa6bdd 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -697,6 +697,25 @@ 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 ( + 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() + ): + 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), @@ -1579,6 +1598,20 @@ def forward_absorption_generation( device=q.device, ) + fp4_mla = ( + 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() + ) + if fp4_mla and latent_cache is None: + raise RuntimeError( + "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 {} + def _mla_gen_rope(): if self.apply_rotary_emb: # Non-fused backends (Vanilla / FlashInfer) do not fuse RoPE @@ -1608,9 +1641,13 @@ 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 + 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 cb6466258562..e08afe5f4632 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 = getattr(model_config, "quant_config", 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.") + 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 " + "Python NIXL integration.") + if is_hybrid_linear(config): + raise NotImplementedError( + "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " + "support hybrid linear-attention models.") + 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,14 @@ 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 = getattr(model_config, "quant_config", 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 " + 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/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/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 8642943c4ab5..69392781064c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -54,7 +54,9 @@ 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 +163,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 +209,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 (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("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: + 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 +831,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 +1965,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 a5ddf0b68d43..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 @@ -59,6 +59,7 @@ _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS_STR = "/".join( f"SM{sm_version}" for sm_version in _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) +FP4_MLA_TOKENS_PER_BLOCK = 128 class _ExecutorMemoryMonitor: @@ -642,7 +643,20 @@ def allocation_scope(current_stage: ExecutorMemoryType): enable_overlap_headroom=getattr(model_engine, "_enable_overlap_headroom", False)) if is_mla(config): - if model_engine.model.model_config.enable_flash_mla: + 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" + ) + 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 diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index f2cb1e00d237..7363ace81da9 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -945,6 +945,8 @@ 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, "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: for i in range(num_contexts, batch_size): 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 8d5e4ae62676..b3a8eb5f368c 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -298,6 +298,53 @@ 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(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, + 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: + 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]) def test_bfloat16_flashinfer(self, enable_chunked_prefill): diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 27921e20429d..02ba64e2f551 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 + # 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/microbenchmarks/bench_fp4_mla_decode.py b/tests/microbenchmarks/bench_fp4_mla_decode.py new file mode 100644 index 000000000000..98c957475d99 --- /dev/null +++ b/tests/microbenchmarks/bench_fp4_mla_decode.py @@ -0,0 +1,1872 @@ +# 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, + 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() diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index 521ad2554009..0715202baa3c 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -10,12 +10,14 @@ 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 +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, 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, @@ -34,25 +36,24 @@ 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 ( - 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 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 @@ -75,17 +76,30 @@ 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, + fp4_mla_state=Fp4MlaState(hp_pool=object(), v_scale_pool=object()), ) forward_args = AttentionForwardArgs(attention_input_type=AttentionInputType.generation_only) - Fp4MlaFmha._validate_request(None, None, None, metadata, forward_args) + 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._validate_request(None, None, None, metadata, forward_args) + 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: @@ -99,8 +113,9 @@ def test_fp8_mla_context_partition_metadata_separates_q_and_kv_lengths() -> None 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(), + fp4_mla_state=Fp4MlaState( + positions=torch.arange(80, dtype=torch.int32), fp8_context_state=object() + ), ) scratch = SimpleNamespace( cache_manager_view=object(), @@ -122,7 +137,7 @@ def test_fp8_mla_context_partition_metadata_separates_q_and_kv_lengths() -> None 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 + assert fp8_meta.fp4_mla_state is None def _reset_triton_allocator() -> None: @@ -138,9 +153,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) @@ -157,13 +175,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 = [] @@ -174,7 +195,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, ) @@ -189,7 +210,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() @@ -209,11 +230,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, @@ -271,17 +289,14 @@ def test_fp4_mla_v2_runtime_sizing_accounts_for_pipeline_slots() -> None: 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._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) - # 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 + # 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 @@ -401,24 +416,24 @@ def get_v_packed_pool(local_layer: int) -> torch.Tensor: ), get_mla_v_packed_pool=get_v_packed_pool, ) - monkeypatch.setattr(fp4_mla_backend, "_fp4_mla_attention_backend", lambda: "cutedsl") + monkeypatch.setattr(fp4_mla_v_cache, "_fp4_mla_attention_backend", lambda: "cutedsl") monkeypatch.setattr( - fp4_mla_backend, + fp4_mla_v_cache, "_fp4_mla_cutedsl_fused_v_transpose_enabled", lambda: fused_v_transpose, ) monkeypatch.setattr( - fp4_mla_backend, + fp4_mla_v_cache, "_stage_fp4_mla_import_page_metadata", stage_page_metadata, ) monkeypatch.setattr( - fp4_mla_backend, + fp4_mla_v_cache, "_rebuild_fp4_mla_v_scales_from_k_scales", rebuild_v_scales, ) monkeypatch.setattr( - fp4_mla_backend, + fp4_mla_v_cache, "_repack_cutedsl_v_packed_cache", rebuild_packed_v, ) @@ -494,14 +509,14 @@ def rebuild_v_scales(_sf_cache, _v_scale_pool, _page_ids, _page_valid_tokens, ** 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_v_cache, "_fp4_mla_attention_backend", lambda: "triton") monkeypatch.setattr( - fp4_mla_backend, + fp4_mla_v_cache, "_stage_fp4_mla_import_page_metadata", lambda *args, **kwargs: (page_ids, page_valid_tokens), ) monkeypatch.setattr( - fp4_mla_backend, + fp4_mla_v_cache, "_rebuild_fp4_mla_v_scales_from_k_scales", rebuild_v_scales, ) @@ -564,9 +579,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() @@ -586,37 +598,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: @@ -655,7 +647,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( @@ -680,7 +672,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)) @@ -734,41 +728,41 @@ def _build_multi_seq_metadata(kv_cache_manager, *, seq_lens, page_size): 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, + ), ) @@ -778,13 +772,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) @@ -836,16 +832,14 @@ 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) 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 = [] @@ -863,14 +857,7 @@ 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) @@ -971,12 +958,12 @@ def test_fp4_mla_chunked_cache_gather_dequantizes_k_residual(monkeypatch) -> Non 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") @@ -991,7 +978,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, @@ -1015,11 +1004,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, @@ -1028,13 +1017,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 = ( @@ -1052,8 +1041,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 @@ -1062,12 +1053,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) @@ -1121,16 +1112,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( @@ -1142,20 +1133,20 @@ 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 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) @@ -1165,7 +1156,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, @@ -1182,16 +1173,16 @@ 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, ) - 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 = ( @@ -1207,7 +1198,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 @@ -1268,7 +1259,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)), @@ -1301,9 +1291,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() @@ -1401,7 +1391,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, @@ -1490,9 +1480,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], 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,