# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import functools import math import os import weakref from contextlib import nullcontext from typing import Dict, Optional, cast import torch from torch import nn import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm._utils import get_sm_version, is_sm_100f, is_sm_120f from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.quantization.utils.fp4_utils import NVFP4_SF_VEC_SIZE from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE from ..distributed import AllReduceParams from ..locality_domain.policy import LocalityDomainExecutionPlanner, LocalityDomainPolicy from ..locality_domain.runtime import LocalityDomainRuntime from ..model_config import ModelConfig from ..modules.linear import Linear, TensorParallelMode, is_static_nvfp4_input_eligible from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..pyexecutor.breakable_cuda_graph import eager_on_graph from ..utils import ( AuxStreamType, Fp4QuantizedTensor, compute_swizzled_sf_shape, maybe_compiled_cat, maybe_compiled_copy_, ) from .attention import ( _helix_cp_allgather_input, _helix_cp_output_projection, _helix_post_process, _helix_zero_kv_mask, extract_extra_attrs, ) from .backends import ( AttentionForwardArgs, AttentionInputType, AttentionMetadata, FlashInferAttentionMetadata, TrtllmAttention, TrtllmAttentionMetadata, ) from .backends.interface import AttentionBackend, PositionalEmbeddingParams, PredefinedAttentionMask from .backends.sparse.hooks import get_sparse_mla_hooks from .backends.utils import create_attention from .rotary_embedding import RotaryEmbedding def _slice_hidden_states_to_num_tokens(hidden_states, num_tokens: int): """Drop CUDA-graph padding by slicing hidden_states to num_tokens rows. For a plain tensor this is the usual row slice. For an Fp4QuantizedTensor (produced when the previous layer's boundary fusion pre-quantized this layer's input) the slice must act on the packed FP4 form: - fp4_tensor: row-major [m, k//2], so [:num_tokens] is a plain row slice; - scaling_factor: 1D swizzled buffer of padUp(m,128)*padUp(cols,4). NVFP4 quant is per-row independent and the swizzle is laid out in independent 128-row tiles, so the leading padUp(num_tokens,128)*padUp(cols,4) bytes are exactly the scale factors for the first num_tokens rows (verified by tests/unittest/_torch/modules/test_fp4_num_tokens_slice.py). - unquantized_hidden_states (when present): plain row slice. """ if not isinstance(hidden_states, Fp4QuantizedTensor): return hidden_states[:num_tokens, ...] fp4 = hidden_states.fp4_tensor if fp4.shape[0] == num_tokens: return hidden_states if not hidden_states.is_sf_swizzled: raise ValueError( "_slice_hidden_states_to_num_tokens only supports swizzled FP4 scaling factors" ) sf = hidden_states.scaling_factor sf_cols = fp4.shape[-1] * 2 // NVFP4_SF_VEC_SIZE padded_rows, padded_cols = compute_swizzled_sf_shape(num_tokens, sf_cols) sf_len = padded_rows * padded_cols sliced_unquant = ( hidden_states.unquantized_hidden_states[:num_tokens] if hidden_states.unquantized_hidden_states is not None else None ) return Fp4QuantizedTensor( fp4_tensor=fp4[:num_tokens].contiguous(), scaling_factor=sf.view(-1)[:sf_len].contiguous(), is_sf_swizzled=True, unquantized_hidden_states=sliced_unquant, ) def _extract_mla_extra_attrs(layer_idx: str): metadata, mla_layer = extract_extra_attrs(layer_idx, "mla") assert isinstance(mla_layer, MLA), "MLA layer must be a subclass of MLA or an instance of MLA" return metadata, mla_layer def create_mla_outputs_impl(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) outputs = mla_layer._create_outputs(hidden_states, metadata) if len(outputs) != 1: raise RuntimeError("MLA custom ops require exactly one output tensor.") return outputs[0] @torch.library.custom_op("trtllm::create_mla_outputs", mutates_args=()) def create_mla_outputs(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @create_mla_outputs.register_fake def _create_mla_outputs_fake(hidden_states: torch.Tensor, layer_idx: str) -> torch.Tensor: return create_mla_outputs_impl(hidden_states, layer_idx) @torch.library.custom_op( "trtllm::mla_custom_op_inplace", mutates_args=("output",), ) def mla_custom_op_inplace( hidden_states: torch.Tensor, position_ids: Optional[torch.Tensor], layer_idx: str, output: torch.Tensor, latent_cache_gen: Optional[torch.Tensor], hidden_states_fp4: Optional[torch.Tensor] = None, hidden_states_sf: Optional[torch.Tensor] = None, ) -> None: metadata, mla_layer = _extract_mla_extra_attrs(layer_idx) if hidden_states_fp4 is not None or hidden_states_sf is not None: assert hidden_states_fp4 is not None and hidden_states_sf is not None, ( "hidden_states_fp4 and hidden_states_sf must be passed together" ) hidden_states = Fp4QuantizedTensor( fp4_tensor=hidden_states_fp4, scaling_factor=hidden_states_sf, unquantized_hidden_states=hidden_states, ) mla_layer.forward_impl( position_ids, hidden_states, metadata, attn_output=[output], latent_cache_gen=latent_cache_gen, ) maybe_bcg_mla_custom_op_inplace = eager_on_graph(mla_custom_op_inplace) def _is_cute_dsl_fp8_bmm_available(sm_version: Optional[int] = None) -> bool: if sm_version is None: sm_version = get_sm_version() if sm_version == 107: return IS_CUTLASS_DSL_RUBIN_AVAILABLE return sm_version in (100, 103) and IS_CUTLASS_DSL_AVAILABLE def _cute_dsl_fp8_bmm_out( mat1_fp8: torch.Tensor, mat2_fp8: torch.Tensor, mat1_scale: torch.Tensor, mat2_scale: torch.Tensor, out: torch.Tensor, ) -> None: """FP8 block-scaled BMM through the CuTe DSL op for the current SM. Callers gate on SM before reaching here; on every SM other than 107 this matches main's unconditional cute_dsl_fp8_bmm_blackwell call exactly. """ if get_sm_version() == 107: torch.ops.trtllm.cute_dsl_fp8_bmm_rubin(mat1_fp8, mat2_fp8, mat1_scale, mat2_scale, out) else: torch.ops.trtllm.cute_dsl_fp8_bmm_blackwell(mat1_fp8, mat2_fp8, mat1_scale, mat2_scale, out) def fp8_block_scaling_bmm_out( mat1: torch.Tensor, mat2_fp8: torch.Tensor, mat2_scale: torch.Tensor, out: torch.Tensor, mat2_dequant: Optional[torch.Tensor] = None, use_cute_dsl_blockscaling_bmm: bool = False, ) -> torch.Tensor: sm_version = get_sm_version() if sm_version == 90 or sm_version == 89: mat1_fp8, mat1_scale = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(mat1) output = out.new_empty(out.shape, dtype=out.dtype, device=out.device) torch.ops.trtllm.fp8_block_scaling_bmm_out( mat1_fp8, mat2_fp8, mat1_scale, mat2_scale, output ) out.copy_(output) elif is_sm_120f(sm_version): mat1_fp8, mat1_scale = fp8_utils.per_token_quant_and_transform(mat1, need_permute102=True) output = out.new_empty(out.shape, dtype=out.dtype, device=out.device) torch.ops.trtllm.fp8_block_scaling_bmm_out( mat1_fp8, mat2_fp8, mat1_scale, mat2_scale, output ) out.copy_(output) elif is_sm_100f(sm_version): if use_cute_dsl_blockscaling_bmm: mat1_fp8, mat1_scale = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(mat1) _cute_dsl_fp8_bmm_out(mat1_fp8, mat2_fp8, mat1_scale, mat2_scale, out) mat1_scale = None else: torch.bmm(mat1.transpose(0, 1), mat2_dequant.transpose(1, 2), out=out) else: raise NotImplementedError(f"SM{sm_version} is not supported") class MLA(nn.Module): def __init__( self, *, hidden_size: int, num_attention_heads: int, num_key_value_heads: int, qk_nope_head_dim: int, qk_rope_head_dim: int, v_head_dim: int, q_lora_rank: int, kv_lora_rank: int, predicted_tokens_per_seq: int, max_position_embeddings: int, bias: bool, aux_stream_dict: Optional[Dict[AuxStreamType, torch.cuda.Stream]] = None, pos_embd_params: Optional[PositionalEmbeddingParams] = None, layer_idx: Optional[int] = None, dtype: torch.dtype = None, dense_bias: Optional[bool] = None, config: Optional[ModelConfig] = None, mapping_with_cp: Optional[Mapping] = None, reduce_output: bool = True, num_groups: int = 1, o_lora_rank: int = 1024, fuse_qkv_a_proj: bool = True, rms_norm_eps: Optional[float] = None, enable_locality_domain_bf16_linear: bool = False, ) -> None: """ Initialize the MLA module. Args: hidden_size (int): The size of the hidden dimension. num_attention_heads (int): The number of attention heads. num_key_value_heads (int): The number of key value heads. qk_nope_head_dim (int): The dimension of the query and key without Rope. qk_rope_head_dim (int): The dimension of the Rope of query and key. v_head_dim (int): The dimension of the value. q_lora_rank (int): The dimension of the compressed query. kv_lora_rank (int): The dimension of the compressed key and value. predicted_tokens_per_seq (int): The number of predicted tokens per sequence. max_position_embeddings (int): The maximum position embeddings. bias (bool): Whether to use bias in the linear layers. aux_stream_dict (Optional[Dict[AuxStreamType, torch.cuda.Stream]]): Auxiliary CUDA streams by role. pos_embd_params (PositionalEmbeddingParams): The positional embedding parameters. layer_idx (int): The layer index. dtype (torch.dtype): The data type. dense_bias (bool): Whether to use bias in the output projection layer. config (ModelConfig): The model configuration. num_groups (int): The number of groups. o_lora_rank (int): The dimension of the compressed output. fuse_qkv_a_proj (bool): Whether q_a and kv_a share one fused projection. Set to ``False`` for checkpoints that store a separate ``q_a_proj``. rms_norm_eps (Optional[float]): Override the RMSNorm epsilon from the pretrained config. If neither source provides a value (e.g. config.pretrained_config is None), falls back to 1e-6. """ super().__init__() self.layer_idx = layer_idx self.layer_idx_str = str(layer_idx) self.dtype = dtype self._weights_transformed = False # Context-phase fused FP8-Q gate and its Q placeholder; both are # resolved once in create_weights(), never on the forward path. self._fused_q_fp8_ctx_ok = False self.register_buffer("_fused_q_placeholder", None, persistent=False) self.hidden_size = hidden_size self.num_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.qk_nope_head_dim = qk_nope_head_dim self.qk_rope_head_dim = qk_rope_head_dim self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim self.v_head_dim = v_head_dim self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank self.predicted_tokens_per_seq = predicted_tokens_per_seq self.max_position_embeddings = max_position_embeddings self.pos_embd_params = pos_embd_params self.dense_bias = dense_bias self.num_groups = num_groups self.o_lora_rank = o_lora_rank self.fuse_qkv_a_proj = fuse_qkv_a_proj if dense_bias is None: self.dense_bias = bias if self.q_lora_rank is None: self.q_lora_rank = hidden_size self.is_lite = True else: self.is_lite = False assert pos_embd_params is not None, "pos_embd_params must be provided in MLA" self.register_to_config = False if config is not None: if "mla_layers" not in config.extra_attrs: config.extra_attrs["mla_layers"] = {} suffix = 0 # ``layer_idx`` is local to an attention stack, while this registry is shared # by target and draft modules in one-model speculative decoding. Keep the first # registration under ``""`` and suffix later collisions as # ``"_"`` so custom ops resolve the originating module without # overwriting another stack's weak reference. while self.layer_idx_str in config.extra_attrs["mla_layers"]: self.layer_idx_str = str(layer_idx) + f"_{suffix}" suffix += 1 config.extra_attrs["mla_layers"][self.layer_idx_str] = weakref.ref(self) self.register_to_config = True config = config or ModelConfig() self.kv_cache_dtype = config.extra_attrs.get("kv_cache_dtype", "auto") sparse_attn_cfg = config.sparse_attention_config sparse_params = ( sparse_attn_cfg.to_sparse_params( pretrained_config=config.pretrained_config, layer_idx=self.layer_idx, ) if sparse_attn_cfg is not None else None ) self.sparse_params = sparse_params self.sparse_attn_hooks = get_sparse_mla_hooks(self) if self.sparse_attn_hooks and not fuse_qkv_a_proj: # forward_dsa_proj assumes the fused [q_a | kv_a | k_pe] # projection layout; the separate q_a_proj layout (Kimi K3 MLA) # is not wired into the sparse MLA paths. raise NotImplementedError( "Sparse MLA requires " "fuse_qkv_a_proj=True; the separate q_a_proj layout is not " "supported with sparse MLA." ) if self.sparse_attn_hooks is not None and type(self)._apply_output_gate is not ( MLA._apply_output_gate ): # _apply_output_gate is defined against the dense output path, where # attn_output[0] is o_proj's input. Sparse hooks own _create_outputs # and _project_output, so the tensor's rank, shape and projection # differ and the gate would not compose. raise NotImplementedError( f"{type(self).__name__} overrides _apply_output_gate, which is " "only defined for the dense MLA output path; it cannot be " "combined with sparse MLA hooks." ) # Fold the residual-less q_a_layernorm -> q_b_proj NVFP4 input # quantization into one fused RMSNorm + FP4-quantize kernel. Resolve # lazily because q_b_proj.input_scale is finalized after weight loading. self._qa_fused_scale = None # tensor parallel if mapping_with_cp is not None: logger.warning_once( "[MLA::__init__] Overriding mapping with CP detected.", key="mla_init_mapping_with_cp", ) self.mapping = mapping_with_cp else: self.mapping = config.mapping tp_size = self.mapping.tp_size pp_size = self.mapping.pp_size cp_size = self.mapping.cp_size dp_size = 1 if self.mapping.enable_attention_dp: dp_size = tp_size tp_size = 1 if self.mapping.has_cp_ulysses(): raise NotImplementedError("MLA doesn't support CP Ulysses yet") if self.mapping.cp_size > 1: assert self.mapping.has_cp_helix(), ( f"CP type must be HELIX for MLA, but got {self.mapping.cp_config['cp_type']}." ) mapping = Mapping( world_size=pp_size * dp_size * tp_size * cp_size, tp_size=tp_size, pp_size=pp_size * dp_size, cp_size=cp_size, cp_config=self.mapping.cp_config, rank=self.mapping.rank, gpus_per_node=self.mapping.gpus_per_node, enable_attention_dp=self.mapping.enable_attention_dp, ) assert self.num_heads % (tp_size * cp_size) == 0 self.num_heads_tp = self.num_heads // tp_size self.num_heads_tp_cp = self.num_heads_tp // cp_size self.num_context_heads = ( self.num_heads_tp_cp if self.mapping.has_cp_helix() and self.sparse_attn_hooks is None else self.num_heads_tp ) if self.num_context_heads != self.num_heads_tp and config.attn_backend.upper() != "TRTLLM": raise NotImplementedError( "Dense MLA with Helix context parallelism requires the TRTLLM attention backend; " f"got {config.attn_backend}." ) self.num_key_value_heads_tp = (self.num_key_value_heads + tp_size - 1) // tp_size if rms_norm_eps is None: rms_norm_eps = getattr(config.pretrained_config, "rms_norm_eps", None) if rms_norm_eps is None: # No explicit value and pretrained_config is None or lacks # rms_norm_eps (e.g. unit tests constructing MLA directly): # keep the historical default. rms_norm_eps = 1e-6 quant_config = config.get_quant_config() self.quant_config = quant_config # Construction settings used by sparse module hooks. self.rms_norm_eps = rms_norm_eps self.reduce_output = reduce_output self.skip_create_weights_in_init = config.skip_create_weights_in_init self.allreduce_strategy = config.allreduce_strategy self.force_dynamic_quantization = config.force_dynamic_quantization self.num_hidden_layers = getattr(config.pretrained_config, "num_hidden_layers", None) self.pp_mapping = config.mapping self.use_cute_dsl_blockscaling_mm = config.use_cute_dsl_blockscaling_mm self.use_cute_dsl_blockscaling_bmm = config.use_cute_dsl_blockscaling_bmm self.use_cute_dsl_bf16_bmm = config.use_cute_dsl_bf16_bmm self._kvbproj_head_weights_cache: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = ( None ) self.use_cute_dsl_bf16_gemm = config.use_cute_dsl_bf16_gemm locality_domain_policy = config.extra_attrs.get("locality_domain_policy") self._locality_domain_policy = ( locality_domain_policy if locality_domain_policy is not None else LocalityDomainPolicy() ) self._locality_domain_runtime: Optional[LocalityDomainRuntime] = None self._locality_domain_k_b_proj_trans_shards: Optional[tuple[torch.Tensor, torch.Tensor]] = ( None ) self._locality_domain_v_b_proj_shards: Optional[tuple[torch.Tensor, torch.Tensor]] = None if not self.is_lite: kv_a_out_features = self.kv_lora_rank + self.qk_rope_head_dim if self.fuse_qkv_a_proj: kv_a_out_features += self.q_lora_rank self.kv_a_proj_with_mqa = Linear( hidden_size, kv_a_out_features, bias=bias, dtype=dtype, quant_config=quant_config, skip_create_weights_in_init=config.skip_create_weights_in_init, use_custom_cublas_mm=True, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, ) if not self.fuse_qkv_a_proj: self.q_a_proj = Linear( hidden_size, self.q_lora_rank, bias=bias, dtype=dtype, quant_config=quant_config, skip_create_weights_in_init=config.skip_create_weights_in_init, use_custom_cublas_mm=True, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, ) self.q_a_layernorm = RMSNorm( hidden_size=self.q_lora_rank, eps=rms_norm_eps, dtype=dtype ) self.q_b_proj = Linear( self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=bias, dtype=dtype, mapping=mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, quant_config=quant_config, skip_create_weights_in_init=config.skip_create_weights_in_init, allreduce_strategy=config.allreduce_strategy, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear, locality_domain_policy=self._locality_domain_policy, ) else: self.kv_a_proj_with_mqa = Linear( hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=bias, dtype=dtype, quant_config=quant_config, skip_create_weights_in_init=config.skip_create_weights_in_init, use_custom_cublas_mm=True, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, ) self.q_proj = Linear( self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=bias, dtype=dtype, mapping=mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, quant_config=quant_config, skip_create_weights_in_init=config.skip_create_weights_in_init, allreduce_strategy=config.allreduce_strategy, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear, locality_domain_policy=self._locality_domain_policy, ) self.q_b_proj = self.q_proj mapping_o = Mapping( world_size=pp_size * dp_size * tp_size * cp_size, tp_size=tp_size * cp_size, pp_size=pp_size * dp_size, cp_size=1, rank=self.mapping.rank, gpus_per_node=self.mapping.gpus_per_node, enable_attention_dp=self.mapping.enable_attention_dp, ) self.mapping_o = mapping_o def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 mscale_all_dim = pos_embd_params.rope.mscale_all_dim scaling_factor = pos_embd_params.rope.scale mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) q_scaling = 1.0 / (mscale * mscale) self.aux_stream_dict = aux_stream_dict or {} self.aux_stream = self.aux_stream_dict.get(AuxStreamType.Attention) mqa_aux_stream = ( self.sparse_attn_hooks.get_mqa_aux_stream(self) if self.sparse_attn_hooks is not None else self.aux_stream ) self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] need_absorption = self.sparse_attn_hooks is None or self.sparse_attn_hooks.need_absorption if need_absorption: self.kv_a_layernorm = RMSNorm( hidden_size=self.kv_lora_rank, dtype=dtype, eps=rms_norm_eps ) self.kv_b_proj = Linear( self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=bias, dtype=dtype, mapping=mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, quant_config=quant_config, skip_create_weights_in_init=config.skip_create_weights_in_init, allreduce_strategy=config.allreduce_strategy, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear, locality_domain_policy=self._locality_domain_policy, ) self.v_b_proj = nn.Parameter( torch.empty( (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank), dtype=dtype, ), requires_grad=False, ) # Helix CP only: the BF16 V up-projection weight of every tp-local head, # [num_heads_tp, v_head_dim, kv_lora_rank], built from kv_b_proj in # transform_weights. With it, generation applies the V up-projection to # this rank's partial outputs before the CP all-to-all (see # _helix_project_v_partials), so the exchange and the combine carry # v_head_dim instead of kv_lora_rank values per head. self.register_buffer("helix_v_b_proj_full", None, persistent=False) self._helix_v_up_before_exchange = False self.mqa = create_attention( config.attn_backend, self.layer_idx, self.num_heads_tp, head_dim=self.kv_lora_rank + self.qk_rope_head_dim, num_kv_heads=1, pos_embd_params=self.pos_embd_params, quant_config=quant_config, q_scaling=q_scaling, is_mla_enable=True, q_lora_rank=self.q_lora_rank, kv_lora_rank=self.kv_lora_rank, qk_nope_head_dim=self.qk_nope_head_dim, qk_rope_head_dim=self.qk_rope_head_dim, hidden_size=self.hidden_size, predicted_tokens_per_seq=self.predicted_tokens_per_seq, skip_create_weights_in_init=config.skip_create_weights_in_init, v_head_dim=self.kv_lora_rank, sparse_params=self.sparse_params, dtype=dtype, aux_stream=mqa_aux_stream, rope_append=(self.sparse_attn_hooks is None or self.sparse_attn_hooks.mqa_rope_append), kv_cache_dtype=self.kv_cache_dtype, skip_correction_threshold=config.skip_correction_threshold, ) if self.mqa is None: raise RuntimeError("MLA requires a non-null MQA attention backend") if self.sparse_attn_hooks is not None: self.sparse_attn_hooks.initialize(self) create_dense_mha = self.sparse_attn_hooks is None or self.sparse_attn_hooks.need_dense_mha if create_dense_mha: mha_sparse_params = self.sparse_params if self.sparse_attn_hooks is None else None self.mha = create_attention( config.attn_backend, self.layer_idx, self.num_context_heads, head_dim=self.qk_head_dim, num_kv_heads=( self.num_context_heads if self.num_context_heads != self.num_heads_tp else self.num_key_value_heads_tp ), pos_embd_params=self.pos_embd_params, quant_config=quant_config, q_scaling=q_scaling, is_mla_enable=True, q_lora_rank=self.q_lora_rank, kv_lora_rank=self.kv_lora_rank, qk_nope_head_dim=self.qk_nope_head_dim, qk_rope_head_dim=self.qk_rope_head_dim, v_head_dim=self.v_head_dim, predicted_tokens_per_seq=self.predicted_tokens_per_seq, skip_create_weights_in_init=config.skip_create_weights_in_init, sparse_params=mha_sparse_params, skip_correction_threshold=config.skip_correction_threshold, ) else: self.mha = None need_default_o_proj = ( self.sparse_attn_hooks is None or self.sparse_attn_hooks.need_default_o_proj ) if need_default_o_proj: self.o_proj = Linear( self.num_key_value_heads * self.v_head_dim, self.hidden_size, bias=self.dense_bias, dtype=dtype, mapping=mapping_o, tensor_parallel_mode=TensorParallelMode.ROW, quant_config=quant_config, skip_create_weights_in_init=config.skip_create_weights_in_init, reduce_output=reduce_output, allreduce_strategy=config.allreduce_strategy, force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear, locality_domain_policy=self._locality_domain_policy, ) self.attention_output_hidden_size = self.o_proj.in_features self.softmax_scale = 1.0 / (math.sqrt(self.qk_head_dim) * q_scaling) self.rope_fusion = self.mqa.support_fused_rope() self.rotary_emb = None self.apply_rotary_emb = not self.rope_fusion if self.apply_rotary_emb: self.rotary_emb = RotaryEmbedding( pos_embd_params.rope, head_dim=self.qk_rope_head_dim, is_neox=pos_embd_params.is_neox, ) self.llama_4_scaling = False if hasattr(config.pretrained_config, "llama_4_scaling"): self.llama_4_scaling = True self.floor_scale = getattr( config.pretrained_config.llama_4_scaling, "original_max_position_embeddings", 8192, ) self.attn_scale = getattr(config.pretrained_config.llama_4_scaling, "beta", 0.1) if not config.skip_create_weights_in_init: self.create_weights() def create_weights(self): # self.mha/mqa has no weights but has states that are related to # quant_config, which could be modified after __init__. if self.mha is not None: self.mha.update_quant_config(self.quant_config) self.mqa.update_quant_config(self.quant_config) self._fused_q_fp8_ctx_ok = False # Although we use FP8 MLA for context/generation phase, the output is still in BF16 self.out_scale = None if self.sparse_attn_hooks is not None: self.sparse_attn_hooks.create_weights(self) if self.sparse_attn_hooks is not None and not self.sparse_attn_hooks.need_absorption: self._weights_transformed = False return # k_b_proj_trans's dtype must be consistent with self.kv_b_proj, # which can be modified after __init__ has_fp8_block_scales = bool( self.kv_b_proj.quant_config 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() ): # The absorption weights follow kv_b_proj: BF16, or FP8 with # 128x128 block scales. Generation has a BMM for each that writes # into the same fused Q slice, and the context path expands KV # through kv_b_proj itself, so both dtypes are accepted here. # Context parallelism is Helix for MLA (checked in __init__). if self.apply_rotary_emb or self.kv_cache_dtype == "fp8_ds_mla": raise ValueError( "FP4 MLA requires fused RoPE on the TRTLLM backend without " "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), dtype=mla_weight_dtype, ), requires_grad=False, ) self.k_b_proj_trans_dequant = None self.v_b_proj_dequant = None if has_fp8_block_scales: self.k_b_proj_trans_scale = nn.Parameter( torch.empty( ( self.num_heads_tp, self.kv_lora_rank // 128, self.qk_nope_head_dim // 128, ), dtype=torch.float32, ), requires_grad=False, ) # This parameter will view into self.kv_b_proj.weight_scale after loading weights. # For dummy weight initialization, this parameter is initialized with empty tensor. self.v_b_proj_scale = nn.Parameter( torch.empty( ( self.num_heads_tp_cp, self.v_head_dim // 128, self.kv_lora_rank // 128, ), dtype=torch.float32, ), requires_grad=False, ) if is_sm_100f() and not self.use_cute_dsl_blockscaling_bmm: assert self.dtype == torch.bfloat16 self.k_b_proj_trans_dequant = nn.Parameter( torch.empty( (self.num_heads_tp, self.kv_lora_rank, self.qk_nope_head_dim), dtype=self.dtype, ), requires_grad=False, ) self.v_b_proj_dequant = nn.Parameter( torch.empty( (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank), dtype=self.dtype, ), requires_grad=False, ) else: self.k_b_proj_trans_scale = None self.v_b_proj_scale = None self._locality_domain_k_b_proj_trans_shards = None self._locality_domain_v_b_proj_shards = None self._weights_transformed = False self._init_fused_q_fp8_context() def apply_rope( self, q: torch.Tensor, k_pe: torch.Tensor, position_ids: torch.Tensor, num_heads: Optional[int] = None, ) -> torch.Tensor: num_heads = self.num_heads_tp if num_heads is None else num_heads q = q.view(-1, num_heads, self.qk_head_dim) q_pe = q[..., self.qk_nope_head_dim :].reshape(-1, num_heads * self.qk_rope_head_dim) q_pe, k_pe = self.rotary_emb(position_ids, [q_pe, k_pe]) q[..., self.qk_nope_head_dim :] = q_pe.view(-1, num_heads, self.qk_rope_head_dim) return k_pe def _attn_forward_gen( self, attn_backend: AttentionBackend, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, position_ids: Optional[torch.Tensor], attn_metadata: AttentionMetadata, helix_v_up: bool = False, **kwargs, ): """Run absorbed attention; under Helix CP also exchange and combine the partials. With ``helix_v_up`` the V up-projection is applied to this rank's partial outputs before the all-to-all, so the result holds ``num_heads_tp_cp * v_head_dim`` values per token instead of latents. """ if self.mapping.has_cp_helix(): # partial_o: [num_tokens, num_heads_tp * kv_lora_rank] # softmax_stats: [num_tokens, num_heads_tp, 2] softmax_stats = torch.empty( (q.shape[0], self.num_heads_tp, 2), device=q.device, dtype=torch.float32 ) kwargs["softmax_stats_tensor"] = softmax_stats partial_o = attn_backend.forward( q, k, v, attn_metadata, forward_args=AttentionForwardArgs(**kwargs), ) kv_lora_rank = partial_o.shape[-1] // self.num_heads_tp assert self.kv_lora_rank == kv_lora_rank helix_kv_bounds = getattr(attn_metadata, "helix_kv_bounds", None) # helix_kv_bounds is non-None only on TrtllmAttentionMetadata, # where _helix_spec_tokens_valid is a declared field. if helix_kv_bounds is not None and attn_metadata._helix_spec_tokens_valid: # Speculative verify groups: KV ownership is per-TOKEN. A rank # owning only the tail page of a group has zero visible KV for # the group's leading tokens while its per-sequence kv_len is # nonzero, so the per-sequence mask above misses those rows. # Their decode rows are fully masked with a finite sentinel, # making partial_o an average over (possibly uninitialized) # pool values; the combine multiplies by corr = 0 and # 0 * NaN would poison the token on every CP rank — sanitize # by the per-token bound instead. zero_kv_mask = helix_kv_bounds[: partial_o.shape[0]] == 0 else: # MLA processes only the generation token slice here, so build # the mask from [num_contexts, num_seqs). Skip this expansion # when the more precise per-token bounds above are valid. zero_kv_mask = _helix_zero_kv_mask( attn_metadata, partial_o.shape[0], seq_start=attn_metadata.num_contexts, num_seqs=attn_metadata.num_generations, ) value_dim = kv_lora_rank if helix_v_up: partial_o = self._helix_project_v_partials(partial_o) value_dim = self.v_head_dim return _helix_post_process( partial_o, softmax_stats, self.mapping, self.num_heads_tp_cp, value_dim, self.aux_stream, self.ln_events, zero_kv_mask=zero_kv_mask, ) else: attn_output = attn_backend.forward( q, k, v, attn_metadata, forward_args=AttentionForwardArgs(**kwargs), ) return attn_output def create_output(self, hidden_states: torch.Tensor, num_contexts: int): if isinstance(hidden_states, Fp4QuantizedTensor): assert hidden_states.unquantized_hidden_states is not None, ( "MLA.create_output received an Fp4QuantizedTensor without a " "unquantized_hidden_states view; the producing fusion must use " "return_norm_out=True" ) hidden_states = hidden_states.unquantized_hidden_states num_tokens = hidden_states.shape[0] return hidden_states.new_empty( [num_tokens, self.attention_output_hidden_size], dtype=hidden_states.dtype ) def _create_outputs( self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata ) -> list[torch.Tensor]: """Create the standard output and any algorithm-specific output buffers.""" if self.sparse_attn_hooks is not None: outputs = self.sparse_attn_hooks.prepare_outputs(self, hidden_states, attn_metadata) if outputs is not None: return outputs return [self.create_output(hidden_states, attn_metadata.num_contexts)] def _attention_scaling(self, q, position_ids): def _get_attn_scale(position_ids: torch.Tensor) -> torch.Tensor: positions = position_ids.view(-1) floor = torch.floor((positions + 1.0) / self.floor_scale) attn_scale = torch.log(floor + 1.0) * self.attn_scale + 1.0 return attn_scale.unsqueeze(-1) attn_scale = _get_attn_scale(position_ids) q = (q * attn_scale).to(q.dtype) return q def _resolve_qa_fused_scale(self): """Resolve whether q_a_layernorm can fuse q_b_proj NVFP4 input quantization.""" if self._qa_fused_scale is not None: return self._qa_fused_scale if self._qa_fused_scale is not False else None eligible = False if ( not self.is_lite and getattr(self, "q_a_layernorm", None) is not None and getattr(self, "q_b_proj", None) is not None ): qb = self.q_b_proj qa = self.q_a_layernorm if is_static_nvfp4_input_eligible(qb) and not qa.use_gemma and not qa.is_nvfp4: eligible = True if eligible: self.q_a_layernorm.is_nvfp4 = True self.q_a_layernorm.nvfp4_scale = qb.input_scale self._qa_fused_scale = qb.input_scale if eligible else False return self._qa_fused_scale if eligible else None def _q_a_layernorm_maybe_fused(self, q: torch.Tensor, return_norm_out: bool = False): """Apply q_a_layernorm and fuse static NVFP4 input quantization when eligible.""" scale = self._resolve_qa_fused_scale() if scale is None: normed = self.q_a_layernorm(q) return (normed, normed) if return_norm_out else normed return self.q_a_layernorm(q, return_norm_out=return_norm_out) def forward_impl( self, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, attn_output: list[torch.Tensor], latent_cache_gen: Optional[torch.Tensor] = None, ) -> None: """Run the dense or sparse implementation of the shared MLA module.""" if self.sparse_attn_hooks is not None: self.sparse_attn_hooks.forward( self, position_ids, hidden_states, attn_metadata, attn_output, ) return if attn_metadata is None: # backend-only unit tests mock the forward staging = nullcontext() else: staging = attn_metadata.mixed_precision_kv_staging(self.layer_idx) with staging: self._forward_impl( position_ids, hidden_states, attn_metadata, attn_output[0], latent_cache_gen=latent_cache_gen, ) def _project_q_by_phase( self, q: torch.Tensor, num_ctx_tokens: int ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: if self.num_context_heads == self.num_heads_tp or num_ctx_tokens == 0: projected = self.q_b_proj(q) return projected[:num_ctx_tokens], projected[num_ctx_tokens:] first_head = self.mapping.cp_rank * self.num_context_heads first = first_head * self.qk_head_dim last = first + self.num_context_heads * self.qk_head_dim context = self.q_b_proj.forward_slice(q[:num_ctx_tokens], first, last) generation = self.q_b_proj(q[num_ctx_tokens:]) if q.shape[0] > num_ctx_tokens else None return context, generation def _project_context_kv(self, compressed_kv: torch.Tensor) -> torch.Tensor: if self.num_context_heads == self.num_heads_tp: return self.kv_b_proj(compressed_kv) first_head = self.mapping.cp_rank * self.num_context_heads last_head = first_head + self.num_context_heads key = self.kv_b_proj.forward_slice( compressed_kv, first_head * self.qk_nope_head_dim, last_head * self.qk_nope_head_dim, ) value_start = self.num_heads_tp * self.qk_nope_head_dim value = self.kv_b_proj.forward_slice( compressed_kv, value_start + first_head * self.v_head_dim, value_start + last_head * self.v_head_dim, ) # Preserve the [all K | all V] row stride consumed by context FP8 preprocessing. return torch.cat((key, value), dim=-1) def _forward_impl( self, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, latent_cache_gen: Optional[torch.Tensor] = None, ) -> None: """ Forward pass for the MLA module. Writes result into output tensor in-place. Args: position_ids (Optional[torch.IntTensor]): The position IDs. hidden_states (torch.Tensor): The hidden states. attn_metadata (AttentionMetadata): The attention metadata. output (torch.Tensor): The output tensor to write results into. latent_cache_gen (Optional[torch.Tensor]): The latent cache used in generation. """ # split q, k, v into context and gen batches num_contexts = attn_metadata.num_contexts num_generations = attn_metadata.num_generations num_ctx_tokens = attn_metadata.num_ctx_tokens num_tokens = attn_metadata.num_tokens hidden_states = _slice_hidden_states_to_num_tokens(hidden_states, num_tokens) if position_ids is not None: position_ids = position_ids[..., :num_tokens] if self.is_lite: compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( [self.kv_lora_rank, self.qk_rope_head_dim], -1 ) compressed_kv = self.kv_a_layernorm(compressed_kv) q = hidden_states else: if self.fuse_qkv_a_proj: q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 ) else: q = self.q_a_proj(hidden_states) compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( [self.kv_lora_rank, self.qk_rope_head_dim], -1 ) q, compressed_kv = maybe_execute_in_parallel( lambda: self._q_a_layernorm_maybe_fused(q), lambda: self.kv_a_layernorm(compressed_kv), self.ln_events[0], self.ln_events[1], self.aux_stream, ) projected_q, latent_cache = maybe_execute_in_parallel( lambda: self._project_q_by_phase(q, num_ctx_tokens), lambda: torch.concat([compressed_kv, k_pe], dim=-1), self.ln_events[0], self.ln_events[1], self.aux_stream, ) q_ctx, q_gen = projected_q assert output is not None, "output must be provided" if num_contexts > 0: compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] k_pe_ctx = k_pe[:num_ctx_tokens, ...] latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] if self.apply_rotary_emb: assert position_ids is not None # position_ids spans [ctx..., gen...] in mixed batches; slice to # match q_ctx/k_pe_ctx so external RoPE uses ctx positions. ctx_position_ids = position_ids[..., :num_ctx_tokens] k_pe_ctx = self.apply_rope( q_ctx, k_pe_ctx, ctx_position_ids, self.num_context_heads ) # External RoPE is only used by backends that do not handle # fused RoPE internally, so keep latent_cache in sync. latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], dim=-1) if self.llama_4_scaling: q_ctx = self._attention_scaling(q_ctx, position_ids[..., :num_ctx_tokens]) self.forward_context( q_ctx, compressed_kv_ctx, k_pe_ctx, position_ids, attn_metadata, output[:num_ctx_tokens, :], latent_cache_ctx, ) if num_generations > 0: assert q_gen is not None compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] k_pe_gen = k_pe[num_ctx_tokens:, ...] if latent_cache_gen is None: latent_cache_gen = latent_cache[num_ctx_tokens:, ...] if self.llama_4_scaling: q_gen = self._attention_scaling(q_gen, position_ids[..., num_ctx_tokens:]) self.forward_absorption_generation( q_gen, compressed_kv_gen, k_pe_gen, attn_metadata, output[num_ctx_tokens:num_tokens, :], position_ids=position_ids, latent_cache=latent_cache_gen, ) def _use_kvbproj_strided_out(self) -> bool: """Whether the context KV expansion can skip the k_nope materialization copy by running per-head batched GEMMs that write directly into the strided K buffer (and the v half of a kv-layout buffer). Requires SM100f and an unquantized bf16 kv_b_proj without bias; other configs fall back to the wide-GEMM + copy path. """ if not is_sm_100f(): return False w = getattr(self.kv_b_proj, "weight", None) expected_shape = ( self.num_heads_tp * (self.qk_nope_head_dim + self.v_head_dim), self.kv_lora_rank, ) expected_numel = math.prod(expected_shape) return ( w is not None and w.dtype == torch.bfloat16 and tuple(w.shape) == expected_shape and w.numel() == expected_numel and self.kv_b_proj.bias is None ) def _get_kvbproj_head_weights(self) -> tuple[torch.Tensor, torch.Tensor]: """Per-head views into kv_b_proj.weight, layout [H, N, K] (contiguous). kv_b_proj.weight rows are [H*qk_nope (k_nope block) | H*v_head (v block)] with K = kv_lora_rank, so both blocks reshape to head-major [H, N, K] without any copy. """ cached = getattr(self, "_kvbproj_head_weights_cache", None) w = self.kv_b_proj.weight if cached is not None and cached[0] is w: w_k, w_v = cached[1], cached[2] else: nope_rows = self.num_heads_tp * self.qk_nope_head_dim w_k = w[:nope_rows].view(self.num_heads_tp, self.qk_nope_head_dim, self.kv_lora_rank) w_v = w[nope_rows:].view(self.num_heads_tp, self.v_head_dim, self.kv_lora_rank) self._kvbproj_head_weights_cache = (w, w_k, w_v) if self.num_context_heads != self.num_heads_tp: first = self.mapping.cp_rank * self.num_context_heads return w_k[first : first + self.num_context_heads], w_v[ first : first + self.num_context_heads ] return w_k, w_v def forward_context_default( self, q: torch.Tensor, compressed_kv: torch.Tensor, k_pe: torch.Tensor, position_ids: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Dense MHA context path: expand KV via kv_b_proj and run attention.""" append_owned_kv = ( isinstance(attn_metadata, TrtllmAttentionMetadata) and attn_metadata.helix_context_local_slots is not None ) if append_owned_kv: assert isinstance(self.mha, TrtllmAttention) if self.mha.uses_fp4_mla_attention: raise NotImplementedError( "Owner-only Helix prefill does not support dense FP4 MLA KV cache" ) assert latent_cache is not None self.mha.mla_rope_append_paged_kv_assign_q(q, latent_cache, attn_metadata) k_pe = latent_cache[:, self.kv_lora_rank :] k = torch.empty_like(q).view(-1, self.num_context_heads, self.qk_head_dim) if self._use_kvbproj_strided_out(): # Expand KV via per-head batched GEMMs that write k_nope directly # into the strided K buffer, skipping the [t, H*256] kv # intermediate and the k_nope materialization copy. num_tokens = compressed_kv.shape[0] w_k, w_v = self._get_kvbproj_head_weights() ckv = compressed_kv.unsqueeze(0).expand( self.num_context_heads, num_tokens, self.kv_lora_rank ) # [H, t, qk_nope] view with strides (qk_head_dim, H*qk_head_dim, 1) k_nope_out = k[..., : self.qk_nope_head_dim].transpose(0, 1) # The MLA fp8 quantize kernel (quantizeCopyInputToFp8Kernel, # mlaKernels.cu) hardcodes the V source token stride as # H*(qk_nope+v) -- the layout of the kv.split() view. Keep that # contract: allocate a kv-shaped buffer and only write its v # half; the k_nope half is never written or read. nope_cols = self.num_context_heads * self.qk_nope_head_dim v_cols = self.num_context_heads * self.v_head_dim kv_buf = torch.empty(num_tokens, nope_cols + v_cols, dtype=q.dtype, device=q.device) v = kv_buf[:, nope_cols:] # [H, t, v_head] view with strides (v_head, H*(qk_nope+v), 1) v_out = v.unflatten(-1, (self.num_context_heads, self.v_head_dim)).transpose(0, 1) # cuBLAS consumes the transposed weights and writes strided # outputs directly, without materializing contiguous copies. torch.ops.trtllm.bmm_out(ckv, w_k.transpose(1, 2), k_nope_out) torch.ops.trtllm.bmm_out(ckv, w_v.transpose(1, 2), v_out) else: kv = self._project_context_kv(compressed_kv) k_nope, v = kv.split( [ self.num_context_heads * self.qk_nope_head_dim, self.num_context_heads * self.v_head_dim, ], -1, ) maybe_compiled_copy_( k[..., : self.qk_nope_head_dim], k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim), ) # When rope_fusion=True (apply_rotary_emb=False), the rope portion # of k is left uninitialized here; the fused attention kernel # handles k_pe RoPE via latent_cache instead. if self.apply_rotary_emb or append_owned_kv: k[..., self.qk_nope_head_dim :] = k_pe.view(-1, 1, self.qk_rope_head_dim) k = k.view(-1, self.num_context_heads * self.qk_head_dim) attn_output = self.mha.forward( q, k, v, attn_metadata, forward_args=AttentionForwardArgs( attention_input_type=AttentionInputType.context_only, latent_cache=None if append_owned_kv else latent_cache, out_scale=self.out_scale, output=output, ), ) return attn_output def forward_context_with_cached_kv( self, q: torch.Tensor, latent_cache: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: assert latent_cache is not None trtllm_attention = cast(TrtllmAttention, self.mha) # apply RoPE, append compressed_kv + k_pe to paged kv cache and assign q_pe to q trtllm_attention.mla_rope_append_paged_kv_assign_q(q, latent_cache, attn_metadata) # copy full_compressed_kv and full_k_pe from paged kv cache full_compressed_kv, full_k_pe = trtllm_attention.load_paged_kv_cache_for_mla( attn_metadata, q.dtype ) assert ( full_compressed_kv.shape[0] == attn_metadata.num_ctx_cached_tokens + attn_metadata.num_ctx_tokens ) assert full_compressed_kv.shape[1] == self.kv_lora_rank assert ( full_k_pe.shape[0] == attn_metadata.num_ctx_cached_tokens + attn_metadata.num_ctx_tokens ) assert full_k_pe.shape[1] == self.qk_rope_head_dim assert full_compressed_kv.is_contiguous() assert full_k_pe.is_contiguous() # compute full_k_nope and full_v from full_compressed_kv full_kv = self._project_context_kv(full_compressed_kv) full_k_nope, full_v = full_kv.split( [ self.num_context_heads * self.qk_nope_head_dim, self.num_context_heads * self.v_head_dim, ], -1, ) full_k_nope = full_k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim) full_k_pe = full_k_pe.view(-1, 1, self.qk_rope_head_dim) full_k = maybe_compiled_cat( (full_k_nope, full_k_pe.expand(-1, self.num_context_heads, -1)), dim=-1 ) full_k = full_k.view(-1, self.num_context_heads * self.qk_head_dim) # release pytorch activation memory full_compressed_kv = None full_k_pe = None full_kv = None full_k_nope = None # latent_cache must be None to differentiate from normal context phase, # so that we can skip applying RoPE and appending KV cache inside attention op attn_output = self.mha.forward( q, full_k, full_v, attn_metadata, forward_args=AttentionForwardArgs( attention_input_type=AttentionInputType.context_only, latent_cache=None, out_scale=self.out_scale, output=output, ), ) return attn_output def forward_context_with_chunked_prefill_fp4( self, q: torch.Tensor, compressed_kv: torch.Tensor, latent_cache: torch.Tensor, attn_metadata: TrtllmAttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: """Run dense FP4 MLA chunked prefill through FP8 context partitions. The current chunk is appended to the FP4 cache first. Each cached prefix partition is then gathered from the FP4 cache, expanded, and attended with a full mask; the current chunk is attended causally last, and the partial outputs are merged with their softmax stats. """ trtllm_attention = cast(TrtllmAttention, self.mha) if latent_cache is None: raise RuntimeError("FP4 MLA chunked prefill requires latent_cache.") # Match the FP8 KV path: fuse Q/K RoPE with the cache append, then # consume the rotated current K directly from latent_cache. trtllm_attention.append_fp4_mla_context_cache(q, latent_cache, attn_metadata) current_k_pe = latent_cache.reshape( -1, self.kv_lora_rank + self.qk_rope_head_dim, )[:, self.kv_lora_rank :] self.softmax_stats_tensor = torch.empty( (attn_metadata.num_ctx_tokens, self.num_context_heads, 2), dtype=torch.float, device=q.device, ) self.temp_softmax_stats_tensor = torch.empty_like(self.softmax_stats_tensor) attn_output = output temp_attn_output = q.new_empty( (q.size(0), self.num_context_heads * self.v_head_dim), ) def expand_kv( partition_compressed_kv: torch.Tensor, partition_k_pe: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: partition_kv = self._project_context_kv(partition_compressed_kv) partition_k_nope, partition_v = partition_kv.split( [ self.num_context_heads * self.qk_nope_head_dim, self.num_context_heads * self.v_head_dim, ], -1, ) partition_k_nope = partition_k_nope.view( -1, self.num_context_heads, self.qk_nope_head_dim, ) partition_k_pe = partition_k_pe.view(-1, 1, self.qk_rope_head_dim) partition_k = maybe_compiled_cat( ( partition_k_nope, partition_k_pe.expand(-1, self.num_context_heads, -1), ), dim=-1, ) return ( partition_k.view(-1, self.num_context_heads * self.qk_head_dim), partition_v, ) def partition_forward_args(attention_mask) -> AttentionForwardArgs: return AttentionForwardArgs( attention_input_type=AttentionInputType.context_only, latent_cache=None, out_scale=self.out_scale, attention_mask=attention_mask, softmax_stats_tensor=self.temp_softmax_stats_tensor, chunked_prefill_buffer_batch_size=( attn_metadata.runtime_features.chunked_prefill_buffer_batch_size ), output=temp_attn_output, ) for loop_idx in range(attn_metadata.chunked_loop_num): total_chunked_tokens = int( attn_metadata.host_cu_chunked_seq_len[ loop_idx, attn_metadata.num_contexts, ] ) chunked_compressed_kv, chunked_k_pe = trtllm_attention.load_chunked_kv_cache_for_mla( metadata=attn_metadata, num_ctx_cached_tokens=total_chunked_tokens, cu_chunked_seq_len=attn_metadata.cu_chunked_seq_len[loop_idx], chunked_global_offset=attn_metadata.chunked_global_offset[loop_idx], chunked_max_seq_len=int(attn_metadata.max_chunk_len_per_loop[loop_idx]), out_dtype=q.dtype, ) chunked_k, chunked_v = expand_kv(chunked_compressed_kv, chunked_k_pe) temp_attn_output = trtllm_attention.forward_fp4_mla_context_partition( q, chunked_k, chunked_v, attn_metadata, partition_forward_args(PredefinedAttentionMask.FULL), kv_lens_cuda=attn_metadata.chunked_seq_len[loop_idx], kv_lens_cpu=attn_metadata.host_chunked_seq_len[loop_idx], ) trtllm_attention.merge_attention_for_mla( attn_output, temp_attn_output, self.softmax_stats_tensor, self.temp_softmax_stats_tensor, attn_metadata.merge_op_tensor[loop_idx], attn_metadata, ) current_k, current_v = expand_kv(compressed_kv, current_k_pe) temp_attn_output = trtllm_attention.forward_fp4_mla_context_partition( q, current_k, current_v, attn_metadata, partition_forward_args(PredefinedAttentionMask.CAUSAL), kv_lens_cuda=attn_metadata.prompt_lens_cuda_runtime, kv_lens_cpu=attn_metadata.prompt_lens_cpu_runtime, ) trtllm_attention.merge_attention_for_mla( attn_output, temp_attn_output, self.softmax_stats_tensor, self.temp_softmax_stats_tensor, attn_metadata.merge_op_tensor[attn_metadata.chunked_loop_num], attn_metadata, ) return attn_output def forward_context_with_chunked_prefill( self, q: torch.Tensor, compressed_kv: torch.Tensor, latent_cache: torch.Tensor, # compressed_kv + k_pe [context_tokens, 1, lora_size + rope_size] attn_metadata: TrtllmAttentionMetadata, output: torch.Tensor, ) -> torch.Tensor: trtllm_attention = cast(TrtllmAttention, self.mha) # apply RoPE, append compressed_kv + k_pe to paged kv cache and assign q_pe to q trtllm_attention.mla_rope_append_paged_kv_assign_q(q, latent_cache, attn_metadata) # determine the number of loop # currently we assume that the chunk size is the same as the max_num_tokens chunked_loop_num = attn_metadata.chunked_loop_num # [total_token_q, num_heads, 2] -> [total_token_q, num_heads] float2 self.softmax_stats_tensor = torch.empty( (attn_metadata.num_ctx_tokens, self.num_context_heads, 2), dtype=torch.float, device="cuda", ) self.temp_softmax_stats_tensor = torch.empty( (attn_metadata.num_ctx_tokens, self.num_context_heads, 2), dtype=torch.float, device="cuda", ) attn_output = output temp_attn_output = q.new_empty( (q.size(0), self.num_context_heads * self.v_head_dim), dtype=q.dtype ) # use fake cached_cu_seq_len for chunked loop origin_kv_lens_cuda_runtime = attn_metadata.kv_lens_cuda_runtime origin_kv_lens_runtime = attn_metadata.kv_lens_runtime # host_total_kv_lens is updated in place below, so copy the value out; # a 0-d view would alias it and make the restore a no-op. origin_ctx_total_kv_len = attn_metadata.host_total_kv_lens[0].item() for loop_idx in range(chunked_loop_num): # {b, chunked_unit_size, h, kv_lora_rank + qk_rope_head_dim} zero padded # fetch `loop_idx` chunk from kv cache temp_cu_chunked_seq_len = attn_metadata.cu_chunked_seq_len[loop_idx] total_ctx_chunked_tokens = attn_metadata.host_cu_chunked_seq_len[ loop_idx, attn_metadata.num_contexts ] chunked_global_offset = attn_metadata.chunked_global_offset[loop_idx] chunked_max_seq_len = attn_metadata.max_chunk_len_per_loop[loop_idx] chunked_compressed_kv, chunked_k_pe = trtllm_attention.load_chunked_kv_cache_for_mla( metadata=attn_metadata, num_ctx_cached_tokens=total_ctx_chunked_tokens, cu_chunked_seq_len=temp_cu_chunked_seq_len, chunked_global_offset=chunked_global_offset, chunked_max_seq_len=chunked_max_seq_len, out_dtype=q.dtype, ) # up proj to uncompressed kv # [tokens, 2, h, kv_dim], without rope_dim chunked_kv = self._project_context_kv(chunked_compressed_kv) chunked_k_nope, chunked_v = chunked_kv.split( [ self.num_context_heads * self.qk_nope_head_dim, self.num_context_heads * self.v_head_dim, ], -1, ) chunked_k_nope = chunked_k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim) chunked_k_pe = chunked_k_pe.view(-1, 1, self.qk_rope_head_dim) chunked_k = maybe_compiled_cat( (chunked_k_nope, chunked_k_pe.expand(-1, self.num_context_heads, -1)), dim=-1 ) chunked_k = chunked_k.view(-1, self.num_context_heads * self.qk_head_dim) # release pytorch activation memory chunked_compressed_kv = None chunked_k_pe = None chunked_kv = None chunked_k_nope = None # copy chunked_seq_len to replace kv_lens_runtime attn_metadata.kv_lens_runtime = attn_metadata.host_chunked_seq_len[loop_idx] attn_metadata.kv_lens_cuda_runtime = attn_metadata.chunked_seq_len[loop_idx] attn_metadata.host_total_kv_lens[0] = total_ctx_chunked_tokens # do not apply mask for attention within loop # latent_cache must be None to differentiate from normal context phase, # so that we can skip applying RoPE and appending KV cache inside attention op temp_attn_output = self.mha.forward( q, chunked_k, chunked_v, attn_metadata, forward_args=AttentionForwardArgs( attention_input_type=AttentionInputType.context_only, latent_cache=None, out_scale=self.out_scale, attention_mask=PredefinedAttentionMask.FULL, softmax_stats_tensor=self.temp_softmax_stats_tensor, chunked_prefill_buffer_batch_size=attn_metadata.runtime_features.chunked_prefill_buffer_batch_size, output=temp_attn_output, ), ) # merge attn result temp_merge_op = attn_metadata.merge_op_tensor[loop_idx] trtllm_attention.merge_attention_for_mla( attn_output, temp_attn_output, self.softmax_stats_tensor, self.temp_softmax_stats_tensor, temp_merge_op, attn_metadata, ) # deal with the uncached kv kv = self._project_context_kv(compressed_kv) _, k_pe = latent_cache.view([-1, self.kv_lora_rank + self.qk_rope_head_dim]).split( [self.kv_lora_rank, self.qk_rope_head_dim], -1 ) # final round of attention k_nope, v = kv.split( [ self.num_context_heads * self.qk_nope_head_dim, self.num_context_heads * self.v_head_dim, ], -1, ) k_nope = k_nope.view(-1, self.num_context_heads, self.qk_nope_head_dim) k_pe = k_pe.view(-1, 1, self.qk_rope_head_dim) k = maybe_compiled_cat((k_nope, k_pe.expand(-1, self.num_context_heads, -1)), dim=-1) k = k.view(-1, self.num_context_heads * self.qk_head_dim) # copy q_lens to replace kv_lens_runtime attn_metadata.kv_lens_runtime = attn_metadata.prompt_lens_cpu_runtime attn_metadata.kv_lens_cuda_runtime = attn_metadata.prompt_lens_cuda_runtime attn_metadata.host_total_kv_lens[0] = ( attn_metadata.prompt_lens_cpu_runtime[: attn_metadata.num_contexts].sum().item() ) # latent_cache must be None to differentiate from normal context phase, # so that we can skip applying RoPE and appending KV cache inside attention op temp_attn_output = self.mha.forward( q, k, v, attn_metadata, forward_args=AttentionForwardArgs( attention_input_type=AttentionInputType.context_only, latent_cache=None, out_scale=self.out_scale, softmax_stats_tensor=self.temp_softmax_stats_tensor, chunked_prefill_buffer_batch_size=attn_metadata.runtime_features.chunked_prefill_buffer_batch_size, output=temp_attn_output, ), ) temp_merge_op = attn_metadata.merge_op_tensor[chunked_loop_num] trtllm_attention.merge_attention_for_mla( attn_output, temp_attn_output, self.softmax_stats_tensor, self.temp_softmax_stats_tensor, temp_merge_op, attn_metadata, ) # copy back kv_lens_runtime and kv_lens_cuda_runtime attn_metadata.kv_lens_runtime = origin_kv_lens_runtime attn_metadata.kv_lens_cuda_runtime = origin_kv_lens_cuda_runtime attn_metadata.host_total_kv_lens[0] = origin_ctx_total_kv_len return attn_output @staticmethod @functools.cache def cached_warmup_forward_context_with_cached_kv( num_heads_tp, qk_nope_head_dim, qk_rope_head_dim, kv_lora_rank, v_head_dim, dtype, device, ): """Warmup torch.compile for cat operations with different tensor layouts. Tensors are marked with torch._dynamo.maybe_mark_dynamic(..., 0) on the num_tokens dimension, so for num_tokens != 1 a single warmup run is enough and the compiled kernel generalizes across varying num_tokens at runtime. num_tokens=1 still triggers recompile (torch.compile specializes for it), so it is warmed up separately. Do not use torch.compile with dynamic=True here because it completely ignores tensor layout/stride information, resulting in significantly degraded performance. """ def warmup(num_tokens): chunked_k_nope = k_nope = torch.empty( num_tokens, num_heads_tp * (qk_nope_head_dim + v_head_dim), dtype=dtype, device=device, )[:, : num_heads_tp * qk_nope_head_dim].view(num_tokens, num_heads_tp, qk_nope_head_dim) chunked_k_pe = torch.empty( num_tokens, 1, qk_rope_head_dim, dtype=dtype, device=device ).expand(-1, num_heads_tp, -1) k_pe = torch.empty( num_tokens, 1, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device, )[:, :, -qk_rope_head_dim:].expand(-1, num_heads_tp, -1) torch._dynamo.maybe_mark_dynamic(chunked_k_nope, 0) torch._dynamo.maybe_mark_dynamic(chunked_k_pe, 0) torch._dynamo.maybe_mark_dynamic(k_pe, 0) maybe_compiled_cat((chunked_k_nope, chunked_k_pe), dim=-1) maybe_compiled_cat((k_nope, k_pe), dim=-1) # With dim 0 (num_tokens) marked dynamic, one warmup suffices for all # num_tokens != 1 at runtime. warmup(2) # num_tokens=1 still triggers recompile; warm it separately. warmup(1) def forward_context( self, q: torch.Tensor, compressed_kv: torch.Tensor, k_pe: torch.Tensor, position_ids: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, latent_cache: Optional[torch.Tensor] = None, ) -> torch.Tensor: if isinstance(attn_metadata, FlashInferAttentionMetadata): if ( attn_metadata.enable_context_mla_with_cached_kv and attn_metadata.num_ctx_cached_tokens > 0 ): return self.forward_absorption_context( q, compressed_kv, k_pe, attn_metadata, output, position_ids=position_ids, latent_cache=latent_cache, ) if isinstance(self.mha, TrtllmAttention): assert isinstance(attn_metadata, TrtllmAttentionMetadata) if ( attn_metadata.helix_context_local_slots is not None and attn_metadata.num_ctx_cached_tokens > 0 ): raise NotImplementedError("Helix cached context requires distributed KV gathering") trtllm_attention = cast(TrtllmAttention, self.mha) # Warm up maybe_compiled_cat for both the chunked-prefill path and # the cached-kv path; without this, the cached-kv prefill # (block-reuse without chunked_prefill) recompiles per shape and # may stall inside inductor's compile worker. if trtllm_attention.has_cached_kv_for_mla_context_warmup(attn_metadata): self.cached_warmup_forward_context_with_cached_kv( self.num_context_heads, self.qk_nope_head_dim, self.qk_rope_head_dim, self.kv_lora_rank, self.v_head_dim, q.dtype, q.device, ) if ( trtllm_attention.is_chunked_prefill_for_mla_context(attn_metadata) and get_sm_version() >= 100 ): if trtllm_attention.uses_fp4_mla_attention: return self.forward_context_with_chunked_prefill_fp4( q, compressed_kv, latent_cache, attn_metadata, output ) return self.forward_context_with_chunked_prefill( q, compressed_kv, latent_cache, attn_metadata, output ) elif trtllm_attention.has_cached_kv_for_mla_context( attn_metadata ) or trtllm_attention.is_chunked_prefill_for_mla_context(attn_metadata): if trtllm_attention.uses_fp4_mla_attention: # The cached-KV path writes and reads a dense paged cache. raise RuntimeError( "FP4 MLA attends a cached prefix only through chunked-prefill " "partitions on SM100+; enable chunked prefill." ) return self.forward_context_with_cached_kv(q, latent_cache, attn_metadata, output) return self.forward_context_default( q, compressed_kv, k_pe, position_ids, attn_metadata, output, latent_cache ) def _bmm_bf16_out( self, a, b_no_transpose, b_transposed, output, locality_domain_weight_shards=None ): """BMM with optional CuTe DSL bf16 acceleration on Blackwell/SM107.""" if ( self.use_cute_dsl_bf16_bmm and get_sm_version() == 107 and locality_domain_weight_shards is not None ): torch.ops.trtllm.cute_dsl_bf16_bmm_locality_domain_inplace_rubin( a, locality_domain_weight_shards[0], locality_domain_weight_shards[1], output ) elif self.use_cute_dsl_bf16_bmm and is_sm_100f(): bf16_bmm_op = ( torch.ops.trtllm.cute_dsl_bf16_bmm_rubin if get_sm_version() == 107 else torch.ops.trtllm.cute_dsl_bf16_bmm_blackwell ) bf16_bmm_op(a, b_no_transpose, output) else: if is_sm_120f(): # `a` is a head-major transpose of a [tokens, heads, dim] buffer, so its # batch stride is the token count rather than the tile extent. cuBLAS picks # a TMA-based nvjet kernel for that layout on SM120/121, and # cuTensorMapEncodeTiled cannot describe it -- the kernel then faults with # an MMU page fault (surfaced as CUBLAS_STATUS_INTERNAL_ERROR or a later # illegal memory access). Densify so a non-TMA kernel is selected. a = a.contiguous() torch.ops.trtllm.bmm_out(a, b_transposed, output) def _bind_helix_v_up_weight(self, weight: torch.Tensor) -> None: """Apply the V up-projection with ``weight`` before the Helix CP all-to-all. ``weight`` is the BF16 V weight of every tp-local head, ``[num_heads_tp, v_head_dim, kv_lora_rank]``. """ assert self.mapping.has_cp_helix() assert weight.dtype == torch.bfloat16 assert tuple(weight.shape) == (self.num_heads_tp, self.v_head_dim, self.kv_lora_rank) self.helix_v_b_proj_full = weight self._helix_v_up_before_exchange = True def _build_helix_v_up_weight(self) -> None: """Under Helix CP, build every tp-local head's BF16 V weight from ``kv_b_proj``. ``kv_b_proj.weight`` rows are ``[H * qk_nope_head_dim | H * v_head_dim]``. Runs before ``kv_b_proj`` transforms its own weights (the transform walk visits this module first), so the weight is the loaded BF16 tensor or the loaded FP8 codes with their 128x128 FP32 block scales. Any other layout keeps the V up-projection after the exchange. """ if not self.mapping.has_cp_helix() or self.dtype != torch.bfloat16: return weight = self.kv_b_proj.weight nope_rows = self.num_heads_tp * self.qk_nope_head_dim v_shape = (self.num_heads_tp, self.v_head_dim, self.kv_lora_rank) num_rows = nope_rows + self.num_heads_tp * self.v_head_dim if weight.is_meta or tuple(weight.shape) != (num_rows, self.kv_lora_rank): return if weight.dtype == torch.bfloat16: v_weight = weight[nope_rows:].view(v_shape) plan = self.kv_b_proj.partition_plan if plan.enabled and plan.op_kind == "bf16_linear": # kv_b_proj releases its full weight once it is localized. v_weight = v_weight.clone(memory_format=torch.contiguous_format) elif weight.dtype == torch.float8_e4m3fn: block = 128 scale = self.kv_b_proj.weight_scale if scale.dtype != torch.float32 or tuple(scale.shape) != ( math.ceil(num_rows / block), math.ceil(self.kv_lora_rank / block), ): return row_blocks = torch.arange(nope_rows, num_rows, device=weight.device) // block v_scale = scale[row_blocks].repeat_interleave(block, dim=1)[:, : self.kv_lora_rank] v_weight = (weight[nope_rows:].float() * v_scale).to(self.dtype).view(v_shape) if self.v_b_proj_dequant is not None: # This rank's dequantized V chunk holds the same values. start = self.mapping.cp_rank * self.num_heads_tp_cp self.v_b_proj_dequant = nn.Parameter( v_weight[start : start + self.num_heads_tp_cp], requires_grad=False ) else: return self._bind_helix_v_up_weight(v_weight) def _helix_project_v_partials(self, partial_o: torch.Tensor) -> torch.Tensor: """Apply every tp-local head's V up-projection to this CP rank's partial outputs. The CP combine is a per-head weighted sum of the ranks' partial outputs, so it commutes with the per-head linear V projection. ``partial_o`` is ``[num_tokens, num_heads_tp * kv_lora_rank]``; the result is ``[num_tokens, num_heads_tp * v_head_dim]``. """ num_tokens = partial_o.shape[0] partial_v = partial_o.new_empty((num_tokens, self.num_heads_tp, self.v_head_dim)) self._bmm_bf16_out( partial_o.view(num_tokens, self.num_heads_tp, self.kv_lora_rank).transpose(0, 1), self.helix_v_b_proj_full, self.helix_v_b_proj_full.transpose(1, 2), partial_v.transpose(0, 1), ) return partial_v.view(num_tokens, self.num_heads_tp * self.v_head_dim) def forward_absorption_generation( self, q: torch.Tensor, compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: num_tokens = q.shape[0] q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 ) # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp num_seqs = attn_metadata.num_seqs # The rope kernel writes seqQOffset[batch_idx + 1] for every # batch_idx < batch_size, and CUB-scans total_s_len / seq_len entries # into seqKVOffsets (mlaKernels.cu:579, :650). `batch_size` there is the # count of presented rows, which equals num_seqs only while every # request contributes the same number of query tokens. Under the # token-major presentation a row is a token, so size these by the token # count. Uniform verification keeps the established request-major # allocation rather than scaling these temporary buffers by K. cu_seqlens_rows = ( max(num_seqs, num_tokens) + 1 if getattr(attn_metadata, "is_ragged_verify", False) else num_seqs + 1 ) cu_q_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) cu_kv_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=q.device) use_fp8_mla = getattr(self.mqa, "has_fp8_kv_cache", False) or getattr( self.mqa, "has_fp4_kv_cache", False ) mla_bmm1_scale = None mla_bmm2_scale = None quant_q_buffer = None if use_fp8_mla: mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=q.device) mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=q.device) quant_q_buffer = torch.empty( num_tokens, self.num_heads_tp, (self.kv_lora_rank + self.qk_rope_head_dim), dtype=torch.uint8, device=q.device, ) if hasattr(self, "k_b_proj_trans"): fused_q = torch.empty( [ num_tokens, self.num_heads_tp, (self.kv_lora_rank + self.qk_rope_head_dim), ], dtype=q.dtype, 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 # in the attention kernel. Reuse apply_rope, which rotates # q's q_pe slice in place, then copy rotated k_pe into the # latent cache that the backend appends. assert position_ids is not None assert latent_cache is not None gen_position_ids = position_ids[..., attn_metadata.num_ctx_tokens :] k_pe_rope = self.apply_rope(q, k_pe, gen_position_ids) fused_q[..., self.kv_lora_rank :] = q_pe latent_cache[..., self.kv_lora_rank :] = k_pe_rope else: # Fused backend (TRTLLM): RoPE, latent-cache append and the # trtllm-gen scheduler buffers are all produced in-kernel. if self.kv_cache_dtype == "fp8_ds_mla": fused_q[..., self.kv_lora_rank :] = q_pe else: self.mqa.mla_rope_generation( fused_q, q_pe, latent_cache, attn_metadata, cu_q_seqlens, cu_kv_seqlens, fmha_scheduler_counter, mla_bmm1_scale, mla_bmm2_scale, quant_q_buffer, **fp4_rope_kwargs, ) 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) # [num_heads, num_tokens, self.kv_lora_rank] q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1) # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] # The output of bmm is written directly into fused_q maybe_execute_in_parallel( lambda: self._bmm_bf16_out( q_nope_t, self.k_b_proj_trans, self.k_b_proj_trans.transpose(1, 2), q_nope_out, self._locality_domain_k_b_proj_trans_shards, ), _mla_gen_rope, self.ln_events[0], self.ln_events[1], rope_stream, ) elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: # [num_heads, num_tokens, self.kv_lora_rank] q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1) maybe_execute_in_parallel( lambda: fp8_block_scaling_bmm_out( q_nope, self.k_b_proj_trans, self.k_b_proj_trans_scale, q_nope_out, self.k_b_proj_trans_dequant, self.use_cute_dsl_blockscaling_bmm, ), _mla_gen_rope, self.ln_events[0], self.ln_events[1], rope_stream, ) else: raise NotImplementedError( f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}." ) fused_q = fused_q.view( [ num_tokens, self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim), ] ) else: raise RuntimeError("MLA absorption requires k_b_proj_trans") helix_v_up = self._helix_v_up_before_exchange attn_out_latent = self._attn_forward_gen( self.mqa, fused_q, None, None, position_ids, attn_metadata, attention_input_type=AttentionInputType.generation_only, out_scale=self.out_scale, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by `invokeMLARopeGeneration` cu_q_seqlens=cu_q_seqlens, # used by `mlaGeneration` cu_kv_seqlens=cu_kv_seqlens, # used by `mlaGeneration` fmha_scheduler_counter=fmha_scheduler_counter, # used by `mlaGeneration` mla_bmm1_scale=mla_bmm1_scale, # used by `mlaGeneration` mla_bmm2_scale=mla_bmm2_scale, # used by `mlaGeneration` quant_q_buffer=quant_q_buffer, # used by `mlaGeneration` helix_v_up=helix_v_up, **kwargs, ) fused_q = None if helix_v_up: # The V up-projection ran before the all-to-all; the combine already # produced the [num_tokens, num_heads_tp_cp * v_head_dim] output. output.view(attn_out_latent.shape).copy_(attn_out_latent) return output # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert ( attn_out_latent.shape[0] == q.shape[0] and attn_out_latent.shape[1] == self.num_heads_tp_cp * self.kv_lora_rank ) # [seq, num_heads, kv_lora_rank] attn_out_latent = attn_out_latent.view([-1, self.num_heads_tp_cp, self.kv_lora_rank]) attn_output = output.view([num_tokens, self.num_heads_tp_cp, self.v_head_dim]) if self.v_b_proj.dtype == torch.bfloat16: # [num_heads, seq, kv_lora_rank] x [num_heads, kv_lora_rank, v_head_dim] # -> [num_heads, seq, v_head_dim] self._bmm_bf16_out( attn_out_latent.transpose(0, 1), self.v_b_proj, self.v_b_proj.transpose(1, 2), attn_output.transpose(0, 1), self._locality_domain_v_b_proj_shards, ) elif self.v_b_proj.dtype == torch.float8_e4m3fn: fp8_block_scaling_bmm_out( attn_out_latent, self.v_b_proj, self.v_b_proj_scale, attn_output.transpose(0, 1), self.v_b_proj_dequant, self.use_cute_dsl_blockscaling_bmm, ) else: raise NotImplementedError(f"Missing bmm impl for dtype: {self.v_b_proj.dtype}.") return output def _init_fused_q_fp8_context(self) -> None: """Resolve the static half of the context-phase fused FP8-Q gate (see `_use_fused_q_fp8_context`). Called once from `create_weights()`.""" ok = ( os.environ.get("TRTLLM_DISABLE_FUSED_Q_FP8_QUANT", "0") != "1" # Independent of use_cute_dsl_bf16_bmm (off by default outside PP): # this path is a memory-traffic win, not a GEMM-kernel preference. and IS_CUTLASS_DSL_AVAILABLE and is_sm_100f() # Layout the C++ RoPE kernel instantiates kOutputFp8Q for on this path # (mlaKernels.cu); the CuTe bmm itself has no shape requirement. and self.kv_lora_rank == 512 and self.qk_rope_head_dim == 64 and self.k_b_proj_trans.dtype == torch.bfloat16 and bool(getattr(self.mqa, "has_fp8_kv_cache", False)) ) if ok: # The CuTe epilogue applies no scale, so both writers of the FP8 Q # buffer must agree on unit scale (the PyTorch backend never sets # another value; checked once here, at construction time). scale = getattr(self.mqa, "kv_scale_orig_quant", None) ok = isinstance(scale, torch.Tensor) and bool(torch.all(scale == 1.0).item()) if ok: # The attention op takes Q only for num_tokens / dtype on the fused # path; a stride-0 expand of one row satisfies that without a # [num_tokens, heads * 576] bf16 allocation per layer. Registered # as a (non-persistent) buffer so meta-init / `.to(device)` place # it alongside the weights. self._fused_q_placeholder = torch.empty( 1, self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim), dtype=self.k_b_proj_trans.dtype, device=self.k_b_proj_trans.device, ) self._fused_q_fp8_ctx_ok = ok logger.info(f"MLA fused FP8-Q (absorbed context) {'enabled' if ok else 'disabled'}") def _use_fused_q_fp8_context( self, attn_metadata: AttentionMetadata, rope_applied_in_python: bool ) -> bool: """Context-phase fused FP8-Q for the absorbed MLA path. The FP8 FMHA reads Q as FP8 [tokens, heads, kv_lora + rope]. Instead of a bf16 Q plus a standalone quantize pass, the CuTe-DSL bmm stores FP8 nope columns and the RoPE kernel (kOutputFp8Q) stores FP8 rope columns into one buffer; the C++ op then skips invokeMLAContextFp8Quantize. Requires RoPE fused in the kernel and a context-only batch (generation still reads the bf16 Q). TRTLLM_DISABLE_FUSED_Q_FP8_QUANT=1 disables it. """ return ( not rope_applied_in_python and self._fused_q_fp8_ctx_ok and attn_metadata.num_generations == 0 ) def forward_absorption_context( self, q: torch.Tensor, compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: num_tokens = q.shape[0] q_rope_applied = kwargs.pop("q_rope_applied", False) q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 ) if not hasattr(self, "k_b_proj_trans"): raise RuntimeError("MLA absorption requires k_b_proj_trans") # RoPE on q_pe happened (or will happen) outside the attention kernel. rope_applied_in_python = ( q_rope_applied or self.apply_rotary_emb or self.kv_cache_dtype == "fp8_ds_mla" ) fused_q_fp8 = None if self._use_fused_q_fp8_context(attn_metadata, rope_applied_in_python): # FP8 Q written once: nope columns by the bmm epilogue, rope columns # by the RoPE kernel (kOutputFp8Q). The attention op still takes a # bf16 Q for num_tokens / dtype; nothing reads it, so pass a view. fused_q_fp8 = q.new_empty( [num_tokens, self.num_heads_tp, self.kv_lora_rank + self.qk_rope_head_dim], dtype=torch.float8_e4m3fn, ) torch.ops.trtllm.cute_dsl_bf16_bmm_fp8out_blackwell( q_nope.transpose(0, 1), # [heads, tokens, qk_nope_head_dim] self.k_b_proj_trans, # [heads, kv_lora_rank, qk_nope_head_dim] fused_q_fp8[..., : self.kv_lora_rank].transpose( 0, 1 ), # [heads, tokens, kv_lora_rank] ) fused_q = self._fused_q_placeholder.expand(num_tokens, -1) else: # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp fused_q = torch.empty( [ num_tokens, self.num_heads_tp, (self.kv_lora_rank + self.qk_rope_head_dim), ], dtype=q.dtype, device=q.device, ) 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) # [num_heads, num_tokens, self.kv_lora_rank] q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1) # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] # The output of bmm is written directly into fused_q self._bmm_bf16_out( q_nope_t, self.k_b_proj_trans, self.k_b_proj_trans.transpose(1, 2), q_nope_out, self._locality_domain_k_b_proj_trans_shards, ) elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: # [num_heads, num_tokens, self.kv_lora_rank] q_nope_out = fused_q[..., : self.kv_lora_rank].transpose(0, 1) fp8_block_scaling_bmm_out( q_nope, self.k_b_proj_trans, self.k_b_proj_trans_scale, q_nope_out, self.k_b_proj_trans_dequant, self.use_cute_dsl_blockscaling_bmm, ) else: raise NotImplementedError( f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}." ) if rope_applied_in_python: fused_q[..., self.kv_lora_rank :] = q_pe fused_q = fused_q.view( [ num_tokens, self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim), ] ) attn_out_latent = self._attn_forward_gen( self.mqa, fused_q, None, None, position_ids, attn_metadata, attention_input_type=AttentionInputType.context_only, out_scale=self.out_scale, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by applyMLARopeAndAssignQKVKernelOptContext # fused FP8-Q path only (both None otherwise): the C++ op skips # invokeMLAContextFp8Quantize when both are set. quant_q_buffer=fused_q_fp8, quant_scale_qkv=self.mqa.kv_scale_orig_quant if fused_q_fp8 is not None else None, **kwargs, ) fused_q = None # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert ( attn_out_latent.shape[0] == q.shape[0] and attn_out_latent.shape[1] == self.num_heads_tp_cp * self.kv_lora_rank ) # [seq, num_heads, kv_lora_rank] attn_out_latent = attn_out_latent.view([-1, self.num_heads_tp_cp, self.kv_lora_rank]) attn_output = output.view([num_tokens, self.num_heads_tp_cp, self.v_head_dim]) if self.v_b_proj.dtype == torch.bfloat16: # [num_heads, seq, kv_lora_rank] x [num_heads, kv_lora_rank, v_head_dim] # -> [num_heads, seq, v_head_dim] self._bmm_bf16_out( attn_out_latent.transpose(0, 1), self.v_b_proj, self.v_b_proj.transpose(1, 2), attn_output.transpose(0, 1), self._locality_domain_v_b_proj_shards, ) elif self.v_b_proj.dtype == torch.float8_e4m3fn: fp8_block_scaling_bmm_out( attn_out_latent, self.v_b_proj, self.v_b_proj_scale, attn_output.transpose(0, 1), self.v_b_proj_dequant, self.use_cute_dsl_blockscaling_bmm, ) else: raise NotImplementedError(f"Missing bmm impl for dtype: {self.v_b_proj.dtype}.") return output def _forward_custom_op( self, hidden_states: torch.Tensor, position_ids: Optional[torch.Tensor], attn_output: list[torch.Tensor], latent_cache_gen: Optional[torch.Tensor], ) -> None: """Run the dense or sparse registered custom-op implementation.""" if self.sparse_attn_hooks is not None and self.sparse_attn_hooks.forward_custom_op( self, hidden_states, position_ids, attn_output, latent_cache_gen, ): return output = attn_output[0] if len(attn_output) != 1: raise RuntimeError("MLA custom ops require exactly one output tensor.") if isinstance(hidden_states, Fp4QuantizedTensor): maybe_bcg_mla_custom_op_inplace( hidden_states.unquantized_hidden_states, position_ids, self.layer_idx_str, output, latent_cache_gen, hidden_states.fp4_tensor, hidden_states.scaling_factor, ) else: maybe_bcg_mla_custom_op_inplace( hidden_states, position_ids, self.layer_idx_str, output, latent_cache_gen, ) def _apply_output_gate( self, hidden_states: torch.Tensor, attn_output: torch.Tensor, ) -> torch.Tensor: """Transform the attention output before the output projection. ``attn_output`` is the row-parallel ``o_proj``'s input and ``hidden_states`` is the unquantized module input, so an override that derives a gate from ``hidden_states`` must match ``o_proj``'s input sharding. Runs on both the ``register_to_config`` and eager branches, and ahead of the helix-CP reduce-scatter. Dense output path only: ``__init__`` rejects an override combined with sparse MLA hooks, which replace ``o_proj`` and may hand ``attn_output[0]`` on in a different rank and layout. The base is the identity. """ return attn_output def _project_output( self, attn_output: list[torch.Tensor], position_ids: Optional[torch.Tensor], attn_metadata: AttentionMetadata, all_reduce_params: Optional[AllReduceParams], ) -> torch.Tensor: """Apply the MLA output projection.""" if self.sparse_attn_hooks is not None: output = self.sparse_attn_hooks.project_output( self, attn_output, position_ids, attn_metadata, all_reduce_params, ) if output is not None: return output return self._project_output_impl( attn_output[0], position_ids, attn_metadata, all_reduce_params ) def _project_output_impl( self, attn_output: torch.Tensor, position_ids: Optional[torch.Tensor], attn_metadata: AttentionMetadata, all_reduce_params: Optional[AllReduceParams], ) -> torch.Tensor: """Apply the default MLA output projection implementation.""" return _helix_cp_output_projection( self.o_proj, attn_output, attn_metadata, all_reduce_params, self.mapping, self.mapping_o, self.layer_idx, ) def forward( self, position_ids: Optional[torch.Tensor], hidden_states: torch.Tensor, attn_metadata: AttentionMetadata, all_reduce_params: Optional[AllReduceParams] = None, latent_cache_gen: Optional[torch.Tensor] = None, ) -> torch.Tensor: hidden_states = _helix_cp_allgather_input( hidden_states, attn_metadata, self.mapping, self.layer_idx ) # Unquantized view of the module input, used by create_mla_outputs and # by _apply_output_gate. output_hidden_states = hidden_states if isinstance(hidden_states, Fp4QuantizedTensor): assert hidden_states.unquantized_hidden_states is not None, ( "MLA.forward received an Fp4QuantizedTensor without a " "unquantized_hidden_states view" ) output_hidden_states = hidden_states.unquantized_hidden_states if self.register_to_config: attn_output = [ torch.ops.trtllm.create_mla_outputs(output_hidden_states, self.layer_idx_str) ] self._forward_custom_op( hidden_states, position_ids, attn_output, latent_cache_gen, ) else: attn_output = self._create_outputs(hidden_states, attn_metadata) self.forward_impl( position_ids, hidden_states, attn_metadata, attn_output=attn_output, latent_cache_gen=latent_cache_gen, ) attn_output[0] = self._apply_output_gate(output_hidden_states, attn_output[0]) return self._project_output(attn_output, position_ids, attn_metadata, all_reduce_params) def resmooth_parameters(self, module_weight, module_weight_scale, recipe=(1, 128, 128)): weight, weight_scale = fp8_utils.resmooth_to_fp8_e8m0(module_weight, module_weight_scale) transfromed_scale = fp8_utils.transform_sf_into_required_layout( weight_scale, mn=weight.shape[1], k=weight.shape[2], recipe=recipe, num_groups=weight.shape[0], is_sfa=False, ) weight_param = torch.nn.Parameter(weight, requires_grad=False) scale_param = torch.nn.Parameter(transfromed_scale, requires_grad=False) return weight_param, scale_param def transform_weights(self) -> None: if self._weights_transformed: return sparse_attn_hooks = getattr(self, "sparse_attn_hooks", None) if sparse_attn_hooks is not None: sparse_attn_hooks.transform_weights(self) if sparse_attn_hooks is not None and not sparse_attn_hooks.need_absorption: self._weights_transformed = True return has_fp8_block_scales = bool( self.kv_b_proj.quant_config and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales() ) if is_sm_120f() and has_fp8_block_scales: self.k_b_proj_trans, self.k_b_proj_trans_scale = self.resmooth_parameters( self.k_b_proj_trans, self.k_b_proj_trans_scale, recipe=(1, 128, 128) ) self.v_b_proj, self.v_b_proj_scale = self.resmooth_parameters( self.v_b_proj, self.v_b_proj_scale, recipe=(1, 128, 128) ) self._build_helix_v_up_weight() self._rebuild_locality_domain_bf16_bmm_weight_shards() self._weights_transformed = True def cache_derived_state(self) -> None: if self.num_context_heads != self.num_heads_tp: first_head = self.mapping.cp_rank * self.num_context_heads last_head = first_head + self.num_context_heads self.q_b_proj.prepare_output_slice( first_head * self.qk_head_dim, last_head * self.qk_head_dim ) self.kv_b_proj.prepare_output_slice( first_head * self.qk_nope_head_dim, last_head * self.qk_nope_head_dim ) value_start = self.num_heads_tp * self.qk_nope_head_dim self.kv_b_proj.prepare_output_slice( value_start + first_head * self.v_head_dim, value_start + last_head * self.v_head_dim, ) self._weights_transformed = True def post_load_weights(self) -> None: self.transform_weights() self.cache_derived_state() def _bind_v_b_proj_weight(self, weight: torch.Tensor) -> None: """Bind the V absorption weight without retaining a released full weight.""" plan = self.kv_b_proj.partition_plan if plan.enabled and plan.op_kind == "bf16_linear": weight = weight.clone(memory_format=torch.contiguous_format) self.v_b_proj = nn.Parameter(weight, requires_grad=False) self._locality_domain_k_b_proj_trans_shards = None self._locality_domain_v_b_proj_shards = None self._weights_transformed = False def _has_full_absorption_weights(self) -> bool: if getattr(self, "kv_lora_rank", None) != 512 or getattr(self, "v_head_dim", None) != 128: return False # The Rubin BF16 BMM requires its K dimension to be 16-byte aligned. # qk_nope_head_dim is K for the absorbed K projection. if getattr(self, "qk_nope_head_dim", 0) % 8 != 0: return False k_weight = getattr(self, "k_b_proj_trans", None) v_weight = getattr(self, "v_b_proj", None) if not isinstance(k_weight, torch.Tensor) or not isinstance(v_weight, torch.Tensor): return False if k_weight.is_meta or v_weight.is_meta: return False expected_k_shape = ( self.num_heads_tp, self.kv_lora_rank, self.qk_nope_head_dim, ) expected_v_shape = ( self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank, ) return ( k_weight.dtype == v_weight.dtype and tuple(k_weight.shape) == expected_k_shape and tuple(v_weight.shape) == expected_v_shape and k_weight.numel() == math.prod(expected_k_shape) and v_weight.numel() == math.prod(expected_v_shape) ) def _copy_absorption_weight_shards( self, weight: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: assert self._locality_domain_runtime is not None num_partitions = self._locality_domain_runtime.num_partitions assert weight.shape[1] % num_partitions == 0 shard_n = weight.shape[1] // num_partitions shards = [] for partition_id in range(num_partitions): weight_slice = weight[ :, partition_id * shard_n : (partition_id + 1) * shard_n, :, ] with self._locality_domain_runtime.partition_weight_context(partition_id): shard = torch.empty_like( weight_slice, memory_format=torch.contiguous_format, ) shard.copy_(weight_slice) shards.append(shard) assert len(shards) == 2 return shards[0], shards[1] def _rebuild_locality_domain_bf16_bmm_weight_shards(self) -> None: """Rebuild localized BF16 absorption weights from the full parameters.""" self._locality_domain_k_b_proj_trans_shards = None self._locality_domain_v_b_proj_shards = None if not self._has_full_absorption_weights(): return locality_domain_policy = getattr(self, "_locality_domain_policy", None) if locality_domain_policy is None: return plan = LocalityDomainExecutionPlanner(locality_domain_policy).plan_bf16_bmm( dtype=self.k_b_proj_trans.dtype, use_cute_dsl_bf16_bmm=getattr(self, "use_cute_dsl_bf16_bmm", False), ) if not plan.enabled: return if getattr(self, "_locality_domain_runtime", None) is None: self._locality_domain_runtime = LocalityDomainRuntime(plan.num_partitions) k_shards = self._copy_absorption_weight_shards(self.k_b_proj_trans) v_shards = self._copy_absorption_weight_shards(self.v_b_proj) self._locality_domain_runtime.prepare_for_capture(plan) self._locality_domain_k_b_proj_trans_shards = k_shards self._locality_domain_v_b_proj_shards = v_shards