From e662db2dde80eb197ca0f8cca2eb6bc3964ccf48 Mon Sep 17 00:00:00 2001 From: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> Date: Wed, 2 Sep 2026 08:24:53 +0000 Subject: [PATCH] [TRTLLM-15033][fix] Revert FlashInfer CuTeDSL MLA dispatch Remove the FlashInfer-to-CuTeDSL MLA backend path introduced by #17800 while preserving the standalone CuTeDSL FMHA backend and later DSA, Helix, combined-FMHA, and sysinfo changes. Signed-off-by: Pengbo Wang <221450789+pengbowang-nv@users.noreply.github.com> --- .../deployment-guide-for-kimi-k3-on-trtllm.md | 14 +- examples/kimi_k3/README.md | 17 -- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 28 -- .../attention/backends/fmha/cute_dsl_mla.py | 8 - .../backends/fmha/flashinfer_trtllm_gen.py | 242 +----------------- .../attention/backends/fmha/interface.py | 48 +--- .../_torch/attention/backends/trtllm.py | 34 --- .../_torch/attention/backends/utils.py | 11 - tensorrt_llm/_torch/attention/mla.py | 5 - .../kimi_k3_mla/kimi_k3_mla_attention.py | 118 +-------- .../test_lists/test-db/l0_b200.yml | 4 - ..._dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 8 +- .../_torch/attention/fmha_test_utils.py | 1 - .../sparse/dsa/test_req_idx_per_token.py | 2 - .../_torch/attention/test_fmha_page_index.py | 221 +--------------- .../modules/test_kimi_k3_mla_backend.py | 178 ------------- 16 files changed, 21 insertions(+), 918 deletions(-) delete mode 100644 tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index 4d543f9deee2..d8bc7674df76 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -38,7 +38,7 @@ The checkpoint and the configuration file must live on a shared filesystem visib * **High-throughput and low-latency deployments are provided.** DEP16 (`enable_attention_dp: true`, `moe_expert_parallel_size: 16`) is the high-throughput deployment. TEP16 (`enable_attention_dp: false`, `moe_expert_parallel_size: 16`) is the low-latency deployment. An 8-GPU deployment, TEP8 (`enable_attention_dp: false`, `moe_expert_parallel_size: 8`), is also provided. Select the deployment and concurrency appropriate for your workload. * **CUDA graphs and the overlap scheduler are enabled.** The performance-sweep recipes set `disable_overlap_scheduler: false` and enable CUDA graphs. DEP16 additionally sets `cuda_graph_config.enable_padding: true`. * **Chunked prefill is supported and enabled** (`enable_chunked_prefill: true`), so prompts longer than `max_num_tokens` are scheduled across multiple steps. -* **`kv_cache_config.tokens_per_block` must be `64`** — required by the MLA (576, 512) trtllm-gen generation kernel. +* **`kv_cache_config.tokens_per_block` must be `64`** — required by the MLA (576, 512) generation kernels. * **Speculative decoding and disaggregated serving are not yet available** for Kimi K3; support is under development. See the "Current limitations" section of `examples/kimi_k3/README.md`. ## Deployment Steps @@ -60,16 +60,6 @@ Kimi K3 additionally depends on `fla` and `einops`, installed into the same in-p .venv-3.12/bin/python -m pip install fla-core einops ``` -To use the optimized CuTeDSL MLA kernel, install the FlashInfer revision used by the Kimi K3 example into the same in-place environment: - -```bash -.venv-3.12/bin/python -u -m pip install --force-reinstall --no-deps \ - --no-build-isolation \ - "flashinfer-python[cu13] @ git+https://github.com/PerkzZheng/flashinfer-k3.git@b6cc594918baf76c40c3a6236fd53f0f8fb9d2dc" -``` - -The `packaging>=24.2` requirement of this source build is already satisfied by `requirements.txt`. The TensorRT LLM environment already provides FlashInfer's runtime dependencies. The `--no-deps` option prevents `pip` from replacing the pinned PyTorch, Triton, CUDA, and CuTeDSL packages. Install FlashInfer after TensorRT LLM because a later dependency-resolving TensorRT LLM installation can replace this source revision with the currently pinned `flashinfer-python==0.6.16`. - For general build-from-source instructions see [https://nvidia.github.io/TensorRT-LLM/latest/installation/build-from-source.html](https://nvidia.github.io/TensorRT-LLM/latest/installation/build-from-source.html). ### Recommended Performance Settings @@ -211,7 +201,7 @@ These options are set within the YAML file passed to `trtllm-serve` via the `--c * **Options:** * `enable_block_reuse`: Off by default; set to `true` to enable prefix-cache reuse across requests. * `mamba_state_config.periodic_snapshot_interval`: With block reuse on, the KDA recurrent state is snapshotted every this many tokens so prefix hits can restore it (default `0` = snapshots off; hybrid models only expose reusable prefixes at snapshot boundaries, so set e.g. `256` for block reuse to engage; see `examples/kimi_k3/eval_extra_llm_options_reuse.yaml`). - * `tokens_per_block`: Must be `64`, required by the MLA (576, 512) trtllm-gen generation kernel. + * `tokens_per_block`: Must be `64`, required by the MLA (576, 512) generation kernels. * `free_gpu_memory_fraction`: Fraction of free GPU memory reserved for the paged KV cache after model load. The configuration above uses `0.25` to leave runtime headroom. Lower it if you hit out-of-memory errors. #### `trust_remote_code` diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 27f75b8d2b70..6a07083b965e 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -66,23 +66,6 @@ architectures may be added in a future release. ```bash .venv-3.12/bin/python -m pip install fla-core einops ``` -- To use the optimized CuTeDSL MLA kernel, install the following FlashInfer - revision into the same in-place environment after installing TensorRT-LLM: - - ```bash - .venv-3.12/bin/python -u -m pip install --force-reinstall --no-deps \ - --no-build-isolation \ - "flashinfer-python[cu13] @ git+https://github.com/PerkzZheng/flashinfer-k3.git@b6cc594918baf76c40c3a6236fd53f0f8fb9d2dc" - ``` - - The `packaging>=24.2` requirement of this source build is already - satisfied by `requirements.txt`. The TensorRT-LLM environment already - provides FlashInfer's runtime - dependencies; `--no-deps` prevents pip from replacing its pinned PyTorch, - Triton, CUDA, and CuTeDSL packages. Install FlashInfer last: TensorRT-LLM - currently pins `flashinfer-python==0.6.14`, so a later - dependency-resolving TensorRT-LLM install can replace this source revision. - ## Run the model Kimi K3 requires a multi-node launch. From the repository root, submit the diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 1db7ff0fbc4d..008a575114cb 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -329,34 +329,6 @@ request checks. For mixed non-MLA batches, the manager checks each active phase independently with `is_supported(..., phase=...)`; a phased library accepts only phases backed by its corresponding `run_*()` entry point. -The `TrtllmAttention` constructor's optional `flashinfer_mla_backend` argument -explicitly selects the MLA generation kernel inside -`FlashInferTrtllmGenFmha` for that attention instance. It accepts -`trtllm-gen` or `cute-dsl`; the latter uses the monolithic CuTeDSL decode -implementation. When the argument is `None`, the ordered FMHA-library -dispatch is preserved and FlashInfer uses `trtllm-gen` if reached. When it is -set, the standalone `CuteDslMlaFmha` defers to the explicit FlashInfer -selection. Selecting `cute-dsl` for an MLA layer using FP8 KV cache raises an -exception because the current CuTeDSL kernel does not accept the -device-resident BMM scale tensors produced for FP8 KV. - -`TrtllmAttention.mla_backend_policy` is an optional per-batch override hook: -model code may install a callable -`(static_backend, metadata, num_gen_tokens) -> backend` on an attention -instance to adjust the selection to the batch composition. - -Kimi K3 defaults its absorbed-generation MLA backend to `cute-dsl` for BF16 KV -cache (override with `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen`; other values are -rejected at model build). FP8 KV cache forces `trtllm-gen`. K3 also installs a -per-batch policy that falls back to `trtllm-gen` for mixed -context/generation batches and multi-token generation (speculative -verification), keeping `cute-dsl` for plain one-token-per-request decode. -Any H=96 batch (K3's attention-DP shape) remains on `cute-dsl` regardless of -batch composition: TRTLLM-Gen may select a 64-head Q tile, which does not -divide 96 after K3's head padding removal, and its decode gate rejects -`64 < num_heads_q < 128` — so falling back there would fail engine -initialization (this covers attention-DP speculative verification). - The FMHA package is split by role: - `fmha/interface.py` defines the `Fmha` runtime contract. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index d9a093a958a7..9584d131e245 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -54,14 +54,6 @@ def is_available(cls, attn: "TrtllmAttention") -> bool: ) return False - if attn.flashinfer_mla_backend is not None: - logger.debug( - "Standalone CuTe DSL MLA FMHA is unavailable: an explicit " - "flashinfer_mla_backend delegates MLA generation to " - "FlashInferTrtllmGenFmha." - ) - return False - if not IS_CUTLASS_DSL_AVAILABLE: logger.debug("CuTe DSL MLA FMHA is unavailable: nvidia-cutlass-dsl is not installed.") return False diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index 2dcdb980a915..dfc5d2e4af45 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -16,9 +16,9 @@ """ FlashInfer TRTLLM-Gen FMHA -This module implements attention computation using flashinfer's trtllm-gen kernels, -with an optional CuTeDSL kernel for MLA generation. It provides a TRT-LLM attention -FMHA library for Blackwell architecture (SM100/SM103). Enable or disable it through +This module implements attention computation using flashinfer's trtllm-gen kernels. +It provides a TRT-LLM attention FMHA library for trtllm-gen kernels +(Blackwell architecture: SM100/SM103). Enable or disable it through ``TLLM_FMHA_LIBS``. Architecture: @@ -61,7 +61,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantMode -from .interface import FmhaPhase, _CuteDslMlaStagingKey +from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha from .utils import ( get_attention_chunk_size, @@ -217,18 +217,6 @@ def _cached_build( _install_flashinfer_mla_decode_tuning_config_cache() -_SUPPORTED_MLA_BACKENDS = {"cute-dsl", "trtllm-gen"} - - -def _get_mla_backend(backend: str) -> str: - backend = backend.strip().lower() - if backend not in _SUPPORTED_MLA_BACKENDS: - raise ValueError( - f"flashinfer_mla_backend must be one of {_SUPPORTED_MLA_BACKENDS}, got {backend!r}." - ) - return backend - - _MULTI_CTAS_KV_COUNTER_ALIGNMENT = 8 @@ -252,89 +240,6 @@ def _get_bmm1_scale_log2(bmm1_scale: torch.Tensor) -> torch.Tensor: return bmm1_scale.narrow(0, 1, 1) -@lru_cache(maxsize=128) -def _get_cute_dsl_mla_workspace_size( - max_batch_size: int, - q_len: int, - num_heads: int, - kv_lora_rank: int, - multi_processor_count: int, -) -> int: - from flashinfer.cute_dsl.attention.monolithic.mla_decode import _get_split_kv_and_workspace_size - - return max( - _get_split_kv_and_workspace_size( - batch_size, - q_len, - num_heads, - kv_lora_rank, - multi_processor_count, - )[1] - for batch_size in range(1, max_batch_size + 1) - ) - - -def _get_cute_dsl_mla_buffer_layout( - batch_size: int, - padded_num_pages: int, -) -> Tuple[int, int, int]: - buffer_alignment_bytes = 32 - page_table_bytes = batch_size * padded_num_pages * torch.int32.itemsize - sequence_lengths_offset = ( - math.ceil(page_table_bytes / buffer_alignment_bytes) * buffer_alignment_bytes - ) - kernel_workspace_offset = ( - math.ceil( - (sequence_lengths_offset + batch_size * torch.int32.itemsize) / buffer_alignment_bytes - ) - * buffer_alignment_bytes - ) - return page_table_bytes, sequence_lengths_offset, kernel_workspace_offset - - -def _prepare_cute_dsl_mla_buffers( - workspace: torch.Tensor, - block_tables: torch.Tensor, - sequence_lengths: torch.Tensor, - padded_num_pages: int, - skip_copy: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if block_tables.size(-1) > padded_num_pages: - raise RuntimeError("CuTeDSL MLA page table exceeds its pre-allocated shape.") - batch_size = block_tables.size(0) - if sequence_lengths.numel() < batch_size: - raise RuntimeError("CuTeDSL MLA sequence lengths are smaller than the batch size.") - - page_table_bytes, sequence_lengths_offset, kernel_workspace_offset = ( - _get_cute_dsl_mla_buffer_layout(batch_size, padded_num_pages) - ) - workspace_bytes = workspace.view(torch.uint8).flatten() - if workspace_bytes.numel() <= kernel_workspace_offset: - raise RuntimeError( - f"CuTeDSL MLA workspace has {workspace_bytes.numel()} bytes; " - f"staging metadata alone needs {kernel_workspace_offset} bytes." - ) - page_table_storage = ( - workspace_bytes[:page_table_bytes].view(torch.int32).view(batch_size, padded_num_pages) - ) - if not skip_copy: - page_table_storage.zero_() - page_table_storage[:, : block_tables.size(-1)].copy_(block_tables[:, 0, :]) - - sequence_lengths_storage = workspace_bytes[ - sequence_lengths_offset:kernel_workspace_offset - ].view(torch.int32) - aligned_sequence_lengths = sequence_lengths_storage[:batch_size] - if not skip_copy: - aligned_sequence_lengths.copy_(sequence_lengths.flatten()[:batch_size]) - - kernel_workspace_bytes = (workspace_bytes.numel() - kernel_workspace_offset) // 4 * 4 - kernel_workspace = workspace_bytes[ - kernel_workspace_offset : kernel_workspace_offset + kernel_workspace_bytes - ].view(-1, 4) - return kernel_workspace, page_table_storage, aligned_sequence_lengths - - @lru_cache(maxsize=128) def _get_generation_workspace_layout( dtype: torch.dtype, @@ -417,7 +322,7 @@ def _get_workspace_size( class FlashInferTrtllmGenFmha(PhasedFmha): """ - An attention backend using FlashInfer trtllm-gen and optional MLA CuTeDSL kernels. + An attention backend using pure trtllm-gen kernels from flashinfer. """ # Default KV layout for flashinfer @@ -479,13 +384,6 @@ class FlashInferTrtllmGenFmha(PhasedFmha): def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) self._layout = self.DEFAULT_KV_LAYOUT - requested_mla_backend = _get_mla_backend(attn.flashinfer_mla_backend or "trtllm-gen") - if requested_mla_backend == "cute-dsl" and attn.is_mla_enable and attn.has_fp8_kv_cache: - raise ValueError( - "flashinfer_mla_backend='cute-dsl' does not support FP8 KV cache device scales; " - "use 'trtllm-gen' instead." - ) - self._mla_backend = requested_mla_backend # Read once so the hot path is not sensitive to later environment changes. self._enable_pdl = get_env_enable_pdl() @@ -592,7 +490,6 @@ def _check_mla_generation_support( cls, head_size: int, tokens_per_block: int, - mla_backend: str, kv_lora_rank: Optional[int], qk_rope_head_dim: Optional[int], ) -> Tuple[bool, str]: @@ -630,14 +527,7 @@ def _check_mla_generation_support( f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.", ) - # Scoped to trtllm-gen: SLOWER_MLA_GENERATION_KERNELS was measured on that - # kernel. Callers pass the backend that will actually run this batch, so a - # cute-dsl batch stays selectable while a batch a policy downgraded to - # trtllm-gen is still gated. - if ( - mla_backend == "trtllm-gen" - and (head_dim_qk, head_dim_v, tokens_per_block) in cls.SLOWER_MLA_GENERATION_KERNELS - ): + if (head_dim_qk, head_dim_v, tokens_per_block) in cls.SLOWER_MLA_GENERATION_KERNELS: return ( False, f"[Generation][MLA] slower TRTLLM-GEN decode kernel for " @@ -883,15 +773,9 @@ def _is_supported_with_reason( f"Q={q_dtype}, KV={kv_cache_dtype}, O={o_dtype}." ) if is_mla_enable: - # The effective backend, not the configured one: a policy may downgrade - # cute-dsl to trtllm-gen for this batch, and it is the kernel that - # actually runs that the gate is about. MLA reaches here only as - # generation-only (checked above), so num_gen_tokens == q.size(0), - # matching prepare_workspace's is_gen_only branch. supported, reason = self._check_mla_generation_support( head_size=attn.head_dim, tokens_per_block=tokens_per_block, - mla_backend=self._get_effective_mla_backend(meta, q.size(0)), kv_lora_rank=attn.kv_lora_rank, qk_rope_head_dim=attn.qk_rope_head_dim, ) @@ -994,39 +878,12 @@ def prepare_workspace( fp8_context_fmha=fp8_context_fmha, ) - effective_mla_backend = self._get_effective_mla_backend(metadata, num_gen_tokens) - if is_gen_only and attn.is_mla_enable and effective_mla_backend == "cute-dsl": - if metadata.kv_cache_manager is None: - raise RuntimeError("CuTeDSL MLA requires a paged KV cache manager.") - max_batch_size = metadata.max_num_sequences or metadata.max_num_requests - max_num_pages = metadata.kv_cache_manager.max_blocks_per_seq - tokens_per_block = metadata.tokens_per_block or 64 - pages_per_superblock = 128 // tokens_per_block - padded_num_pages = ( - math.ceil(max_num_pages / pages_per_superblock) * pages_per_superblock - ) - _, _, buffer_metadata_size = _get_cute_dsl_mla_buffer_layout( - max_batch_size, padded_num_pages - ) - assert self._multi_processor_count is not None - cute_dsl_workspace_size = _get_cute_dsl_mla_workspace_size( - max_batch_size, - max(1, attn.predicted_tokens_per_seq), - attn.num_heads, - int(attn.kv_lora_rank or 0), - self._multi_processor_count, - ) - required_workspace_size = max( - required_workspace_size, buffer_metadata_size + cute_dsl_workspace_size - ) - required_workspace_size = math.ceil(required_workspace_size / 4) * 4 - current_workspace_size = workspace.numel() * workspace.element_size() if current_workspace_size < required_workspace_size: if metadata.is_cuda_graph and torch.cuda.is_current_stream_capturing(): raise RuntimeError( "Attention CUDA graph workspace is smaller than the " - f"required size for {effective_mla_backend}." + "required size for trtllm-gen." ) required_workspace_numel = math.ceil(required_workspace_size / workspace.element_size()) workspace.resize_((required_workspace_numel,)) @@ -1037,26 +894,6 @@ def _get_multi_ctas_kv_counter_buffer(self) -> torch.Tensor: raise RuntimeError("The trtllm-gen multi-CTA KV counter buffer is not initialized.") return counter_buffer - def _get_effective_mla_backend( - self, meta: "TrtllmAttentionMetadata", num_gen_tokens: int - ) -> str: - """Resolve the MLA decode backend for the current scheduler batch. - - ``attn.mla_backend_policy`` is a neutral per-batch override hook. - A model whose backend choice depends on the batch composition (e.g. - Kimi K3's MLA module) installs a policy on the attention instance it - owns; the policy receives the statically configured backend, the batch - metadata, and the batch's generation-token count, and returns the - backend to use for this batch. Without a policy (the default for every - model) this returns ``self._mla_backend`` unchanged, so every call - site behaves exactly like the plain ``self._mla_backend`` checks it - replaced. - """ - policy = self.attn.mla_backend_policy - if policy is None: - return self._mla_backend - return policy(self._mla_backend, meta, num_gen_tokens) - @staticmethod def _compute_window_left( cyclic_attention_window_size: int, @@ -1375,9 +1212,6 @@ def run_mla_generation( attn = params.attn meta = params.meta fwd = params.fwd - # ``params.num_tokens`` is the generation-token count here: phased.py - # dispatches MLA generation with the generation slice only. - effective_mla_backend = self._get_effective_mla_backend(meta, params.num_tokens) if 0 < params.cyclic_attention_window_size < params.max_past_kv_length: raise NotImplementedError( "Sliding-window attention is not supported by MLA decode path." @@ -1447,52 +1281,7 @@ def run_mla_generation( ) bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) bmm2_scale = 1.0 - if effective_mla_backend == "cute-dsl": - pages_per_superblock = 128 // params.tokens_per_block - padded_num_pages = ( - math.ceil(block_tables.size(-1) / pages_per_superblock) * pages_per_superblock - ) - # Every MLA layer of one forward step shares this metadata, its - # workspace, and (for a single paged pool) identical block tables - # and sequence lengths — so the staged copies are byte-identical - # across layers. Stage once per step: the first generation-only - # call copies, subsequent layers with a matching key skip the 3 - # copy kernels. ``prepare()`` / ``update_for_spec_dec`` reset - # the key each step so eager forwards always re-stage; under CUDA - # graphs the first layer's captured copies replay once per step. - # The capture flag is part of the key: CUDA-graph capture is - # preceded by warmup forwards on the same metadata without an - # intervening prepare(), and the capture pass MUST re-record the - # staging copies (a skip would freeze stale page tables into the - # graph). - staging_key = _CuteDslMlaStagingKey( - is_capturing=torch.cuda.is_current_stream_capturing(), - workspace_ptr=params.workspace.data_ptr(), - block_tables_ptr=block_tables.data_ptr(), - block_tables_shape=tuple(block_tables.shape), - sequence_lengths_ptr=params.sequence_lengths.data_ptr(), - sequence_lengths_offset=params.seq_offset, - batch_beam=batch_beam, - padded_num_pages=padded_num_pages, - ) - skip_staging_copy = ( - meta.num_contexts == 0 - and getattr(meta, "_cute_dsl_mla_staging_key", None) == staging_key - ) - workspace_buffer, block_tables, sequence_lengths = _prepare_cute_dsl_mla_buffers( - params.workspace, - block_tables, - params.sequence_lengths, - padded_num_pages, - skip_copy=skip_staging_copy, - ) - if meta.num_contexts == 0: - meta._cute_dsl_mla_staging_key = staging_key - uses_shared_paged_kv_idx = True - else: - sequence_lengths = params.sequence_lengths - workspace_buffer = params.workspace.view(-1, 4) - uses_shared_paged_kv_idx = self.USE_SHARED_PAGED_KV_IDX + workspace_buffer = params.workspace.view(-1, 4) flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla( query, # query @@ -1502,7 +1291,7 @@ def run_mla_generation( kv_lora_rank, # kv_lora_rank qk_rope_head_dim, # qk_rope_head_dim block_tables, # block_tables - sequence_lengths, # seq_lens + params.sequence_lengths, # seq_lens params.max_past_kv_length, # max_seq_len 0, # sparse_mla_top_k params.context_buf.view(batch_beam, q_len_per_req, attn.num_heads, kv_lora_rank), # out @@ -1511,15 +1300,8 @@ def run_mla_generation( fwd.attention_sinks, # sinks None, # skip_softmax_threshold_scale_factor self._enable_pdl, # enable_pdl - backend=effective_mla_backend, + backend="trtllm-gen", is_var_seq=True, - uses_shared_paged_kv_idx=uses_shared_paged_kv_idx, - cute_dsl_impl="monolithic", - # flashinfer rejects the counter buffer unless the trtllm-gen - # runner is selected; the cute-dsl MLA path must pass None. - multi_ctas_kv_counter_buffer=( - self._get_multi_ctas_kv_counter_buffer() - if effective_mla_backend != "cute-dsl" - else None - ), + uses_shared_paged_kv_idx=self.USE_SHARED_PAGED_KV_IDX, + multi_ctas_kv_counter_buffer=self._get_multi_ctas_kv_counter_buffer(), ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index dcefd3ee1c8b..1e5bce7161aa 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/interface.py @@ -16,7 +16,7 @@ import weakref from abc import ABC, abstractmethod from enum import Enum -from typing import TYPE_CHECKING, NamedTuple, Optional, Protocol +from typing import TYPE_CHECKING, Optional import torch @@ -29,52 +29,6 @@ ) -class _CuteDslMlaStagingKey(NamedTuple): - """Identifies CuTe-DSL MLA inputs staged into a shared workspace. - - Attributes: - is_capturing: Whether the staging occurred during CUDA graph capture. - workspace_ptr: Address of the shared staging workspace. - block_tables_ptr: Address of the source block tables. - block_tables_shape: Shape of the source block tables. - sequence_lengths_ptr: Address of the source sequence lengths. - sequence_lengths_offset: Offset applied to the source sequence lengths. - batch_beam: Number of generation sequences, including beam expansion. - padded_num_pages: Page-table width after CuTe-DSL alignment padding. - """ - - is_capturing: bool - workspace_ptr: int - block_tables_ptr: int - block_tables_shape: tuple[int, ...] - sequence_lengths_ptr: int - sequence_lengths_offset: int - batch_beam: int - padded_num_pages: int - - -class MlaBackendPolicy(Protocol): - """Selects the MLA generation backend for one scheduler batch.""" - - def __call__( - self, - requested_backend: str, - metadata: "TrtllmAttentionMetadata", - num_gen_tokens: int, - ) -> str: - """Return the backend to use for the supplied batch composition. - - Args: - requested_backend: Backend selected by the attention instance. - metadata: Runtime metadata for the current scheduler batch. - num_gen_tokens: Number of generation tokens in the batch. - - Returns: - Backend name to use for MLA generation in this batch. - """ - ... - - class FmhaPhase(str, Enum): """Attention phase checked by a phased FMHA library.""" diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 5a5715b1d273..914f6f3a654e 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -29,8 +29,6 @@ from ...speculative.interface import SpecMetadata from ...speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm._torch.attention.backends.fmha.interface import ( - MlaBackendPolicy, _CuteDslMlaStagingKey) from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType @@ -200,18 +198,6 @@ def effective_beam_width(self) -> int: init=False, repr=False) - # Per-forward-pass staging key for the CuTeDSL MLA generation workspace - # (page table + sequence lengths). All MLA layers of one step stage - # byte-identical data into the shared workspace, so the first layer - # copies and later layers skip. Reset whenever kv lens can change so - # eager forwards always re-stage (under CUDA graphs the first layer's - # captured copies replay once per step). - _cute_dsl_mla_staging_key: Optional[_CuteDslMlaStagingKey] = field( - default=None, - init=False, - repr=False, - ) - use_paged_context_fmha: bool = field(init=False, default=False, repr=False) # FMHA prologue buffers for the MLA generation path, hoisted out of the per-layer @@ -586,9 +572,6 @@ def _invalidate_mla_scheduler_buffers(self) -> None: # buffers below must be rebuilt before the next MLA layer reads them. self._mla_scheduler_buffers_valid = False self._mla_ctx_cu_seqlens_valid = False - # The staged CuTe DSL page table and sequence lengths are derived from - # the same per-iteration scheduler state. - self._cute_dsl_mla_staging_key = None def update_helix_param( self, @@ -1416,7 +1399,6 @@ def __init__( attention_chunk_size: Optional[int] = None, sparse_params: Optional[SparseParams] = None, kv_cache_dtype: str = "auto", - flashinfer_mla_backend: Optional[str] = None, skip_correction_threshold: float = 0.0, **kwargs, ) -> None: @@ -1437,10 +1419,6 @@ def __init__( values are ``auto``, ``fp8``, ``fp8_ds_mla``, ``nvfp4``, and supported torch dtype strings. ``fp8_ds_mla`` selects the packed sparse-MLA cache used by DeepSeek-V4 and DSA on SM120/SM121. - flashinfer_mla_backend (Optional[str]): FlashInfer MLA generation backend - selected for this attention instance. - None preserves the ordered FMHA-library - dispatch. skip_correction_threshold (float): Runtime MLA threshold. Zero disables skip-correction. """ @@ -1449,18 +1427,6 @@ def __init__( self.sparse_params = sparse_params self.kv_cache_dtype = kv_cache_dtype self.use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla" - self.flashinfer_mla_backend = flashinfer_mla_backend - # Per-batch MLA decode backend override hook. Maps (statically - # configured backend, batch metadata, generation-token count) to the - # backend name for the current batch. None (the default) keeps the - # static ``flashinfer_mla_backend`` selection unchanged. Model code - # that needs batch-dependent selection (e.g. Kimi K3's MLA module) - # installs a policy on the attention instances it owns; it lives on - # the backend object rather than the FMHA lib instances because - # ``update_quant_config`` may recreate the manager after model - # construction. - self.mla_backend_policy: Optional[MlaBackendPolicy] = None - self.is_mla_enable = mla_params is not None sparse_algorithm = getattr(self.sparse_params, "algorithm", None) if (self.is_mla_enable and sparse_algorithm in ("deepseek_v4", "dsa") diff --git a/tensorrt_llm/_torch/attention/backends/utils.py b/tensorrt_llm/_torch/attention/backends/utils.py index 67dbfe38eb6b..07bbbe542c59 100644 --- a/tensorrt_llm/_torch/attention/backends/utils.py +++ b/tensorrt_llm/_torch/attention/backends/utils.py @@ -76,7 +76,6 @@ def create_attention( dtype: Optional[torch.dtype] = None, aux_stream: Optional[torch.cuda.Stream] = None, kv_cache_dtype: str = "auto", - flashinfer_mla_backend: Optional[str] = None, skip_correction_threshold: float = 0.0, ) -> AttentionBackend: if attention_chunk_size is not None and backend_name.upper() != "TRTLLM": @@ -115,16 +114,6 @@ def create_attention( kv_cache_dtype=kv_cache_dtype, skip_correction_threshold=skip_correction_threshold, ) - if flashinfer_mla_backend is not None: - # Only TrtllmAttention understands this selector. Raise instead of - # silently dropping it: a model that configured a specific MLA - # generation kernel must not run on another backend's default. - if not issubclass(attn_cls, TrtllmAttention): - raise ValueError( - f"flashinfer_mla_backend={flashinfer_mla_backend!r} is only " - "supported by the TRTLLM attention backend, but backend " - f"{backend_name} resolves to {attn_cls.__name__}.") - kwargs["flashinfer_mla_backend"] = flashinfer_mla_backend return attn_cls( layer_idx, diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index f7d9e2cc509a..327d5b76739f 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -225,7 +225,6 @@ def __init__( o_lora_rank: int = 1024, fuse_qkv_a_proj: bool = True, rms_norm_eps: Optional[float] = None, - flashinfer_mla_backend: Optional[str] = None, ) -> None: """ Initialize the MLA module. @@ -257,9 +256,6 @@ def __init__( 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. - flashinfer_mla_backend (Optional[str]): Generation backend for the - FlashInfer/TRTLLM-Gen MLA dispatcher. ``None`` preserves the - attention backend default. """ super().__init__() self.layer_idx = layer_idx @@ -571,7 +567,6 @@ def yarn_get_mscale(scale=1, mscale=1): 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, - flashinfer_mla_backend=flashinfer_mla_backend, skip_correction_threshold=config.skip_correction_threshold, ) if self.mqa is None: diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index 34fa50755c3c..4ab1d52f1c55 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -10,121 +10,18 @@ from __future__ import annotations -import os -from functools import partial from typing import Optional import torch from ....functional import PositionEmbeddingType -from ....logger import logger from ....mapping import Mapping -from ....models.modeling_utils import QuantConfig -from ...attention.backends import TrtllmAttention, TrtllmAttentionMetadata +from ...attention.backends import TrtllmAttention from ...attention.backends.interface import PositionalEmbeddingParams, RopeParams from ...attention.mla import MLA from ...model_config import ModelConfig from ..linear import Linear, TensorParallelMode -_KIMI_K3_MLA_GEN_BACKEND_ENV = "TLLM_K3_MLA_GEN_BACKEND" -_KIMI_K3_MLA_GEN_BACKENDS = ("cute-dsl", "trtllm-gen") - - -def _select_mla_generation_backend(quant_config: Optional[QuantConfig]) -> str: - """Select K3's absorbed-generation MLA backend. - - K3 was tuned with the FlashInfer CuTe-DSL backend for BF16 KV cache. - FP8 KV cache carries device scales that CuTe-DSL does not support, so - retain the TRTLLM-Gen fallback used by the pre-refactor implementation. - """ - backend = os.environ.get(_KIMI_K3_MLA_GEN_BACKEND_ENV, "cute-dsl") - # Validate here, where the env var is read: an invalid value would - # otherwise surface only deep inside attention-backend construction, - # with an error that never names the knob that caused it. - if backend not in _KIMI_K3_MLA_GEN_BACKENDS: - raise ValueError( - f"{_KIMI_K3_MLA_GEN_BACKEND_ENV}={backend!r} is invalid; " - f"expected one of {list(_KIMI_K3_MLA_GEN_BACKENDS)}." - ) - has_fp8_kv_cache = bool( - quant_config is not None and quant_config.layer_quant_mode.has_fp8_kv_cache() - ) - if has_fp8_kv_cache and backend != "trtllm-gen": - # info_once: this runs once per MLA layer (~60x at startup for - # FP8-KV) and the decision is identical for every layer. - logger.info_once( - "Kimi K3 MLA: FP8 KV cache requires the trtllm-gen MLA " - f"generation backend; overriding '{backend}' -> 'trtllm-gen'.", - key="kimi_k3_mla_gen_backend_fp8_override", - ) - return "trtllm-gen" - return backend - - -def _validate_mla_generation_backend(backend: str, num_heads: int) -> None: - """Fail fast when `backend` can never run at this per-rank head count. - - FlashInfer's `trtllm_batch_decode_with_kv_cache_mla` rejects - `64 < num_heads_q < 128` for every batch shape, and the per-batch policy - never demotes away from an explicit `trtllm-gen` selection (the - FP8-KV-cache override or a `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen` request). - Without this check the conflict only surfaces as a FlashInfer error deep - in attention warmup. - - The bound mirrors FlashInfer's validation gate verbatim: it predicts - FlashInfer's rejection, it is not a verified support claim for the head - counts outside the range. For Kimi K3 the open side above 128 is - unreachable anyway — per-rank Q heads never exceed 96 (all heads - replicated under attention-DP, `96 / tp_size` under TEP head sharding). - """ - if backend == "trtllm-gen" and 64 < num_heads < 128: - raise ValueError( - "Kimi K3 MLA: the trtllm-gen generation backend cannot run with " - f"{num_heads} query heads per rank (trtllm-gen MLA decode rejects " - "64 < num_heads_q < 128; under attention-DP every rank keeps all " - "heads). trtllm-gen was selected explicitly — by the FP8-KV-cache " - f"override or by {_KIMI_K3_MLA_GEN_BACKEND_ENV}=trtllm-gen. Use " - "tensor-parallel head sharding (TEP) so each rank has <= 64 " - "heads, or a BF16 KV cache with the default cute-dsl backend." - ) - - -def _kimi_k3_mla_decode_backend_policy( - requested_backend: str, - metadata: TrtllmAttentionMetadata, - num_gen_tokens: int, - *, - num_heads: int, -) -> str: - """Per-batch MLA decode backend selection for Kimi K3. - - Installed as ``mla_backend_policy`` on K3's generation attention backend - (see :class:`KimiK3MLAAttention`); the general attention code applies no - such policy on its own. - - CuTe-DSL reuses one staged page table across MLA layers for a - generation-only, one-token-per-request batch. Other mixed batches repeat - the staging copies in every MLA layer and regress time to first token, so - they fall back to TRTLLM-Gen. The H=96 path is the correctness exception - and applies to EVERY fallback candidate: TRTLLM-Gen may select a 64-head - Q tile, which does not divide 96 (invalid after K3's head padding was - removed), and its decode gate rejects 64 < num_heads_q < 128 outright — - falling back would fail engine initialization. H=96 per rank is K3's - attention-DP shape, so this keeps attention-DP + speculative - verification (a generation-only multi-token batch) on CuTe-DSL, which - accepts multi-token queries; K3's decode tuning preference for - TRTLLM-Gen only applies where TRTLLM-Gen is valid at all. - """ - is_single_token_generation = num_gen_tokens == metadata.num_generations - requires_cute_dsl = num_heads == 96 - if ( - requested_backend == "cute-dsl" - and not requires_cute_dsl - and (metadata.num_contexts > 0 or not is_single_token_generation) - ): - return "trtllm-gen" - return requested_backend - def _meta_safe_cast_dtype(module, dtype): """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``. @@ -289,7 +186,6 @@ def __init__( reduce_output=False, fuse_qkv_a_proj=False, rms_norm_eps=rms_norm_eps, - flashinfer_mla_backend=_select_mla_generation_backend(model_config.get_quant_config()), ) # Keep the base MLA registration enabled so breakable CUDA graphs use # the shared custom op. The output gate is a base hook and runs on both @@ -322,18 +218,6 @@ def __init__( assert isinstance(self.mqa, TrtllmAttention) _install_identity_rope_table(self.mha) _install_identity_rope_table(self.mqa) - # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so - # only it needs K3's per-batch fallback policy; mha keeps the - # default trtllm-gen selection. - # Validate here rather than in _select_mla_generation_backend: the - # per-rank head count (replicated under attention-DP, sharded under - # TEP) is only authoritative once the base MLA module has built its - # generation backend. - _validate_mla_generation_backend(self.mqa.flashinfer_mla_backend, self.mqa.num_heads) - self.mqa.mla_backend_policy = partial( - _kimi_k3_mla_decode_backend_policy, - num_heads=self.mqa.num_heads, - ) self.rotary_emb = None self.apply_rotary_emb = False diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 407c04da0168..1a53de7db5fb 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -142,10 +142,6 @@ l0_b200: - unittest/disaggregated/test_kda_mamba_transfer.py - unittest/_torch/moe/test_kimi_k3_mlp.py - unittest/_torch/moe/test_kimi_k3_situ_moe.py - # Kimi K3 MLA backend selection + config/checkpoint plumbing (monkeypatched, - # HW-light). Collected nowhere else: the CPU-Generic stage only collects files - # carrying the literal cpu_only marker, and these carry none. - - unittest/_torch/modules/test_kimi_k3_mla_backend.py - unittest/_torch/modeling/test_kimi_k3_config_routing.py - unittest/_torch/modeling/test_kimi_linear_checkpoint.py - unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index dbc2d64a5a0f..0ab917aff209 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -37,14 +37,14 @@ environment: trtllm_repo: '' build_wheel: false work_dir: - # TLLM_K3_MLA_GEN_BACKEND / TLLM_KV_CACHE_MANAGER_V2_BACKEND / - # TRTLLM_ENABLE_PDL / NCCL_GRAPH_MIXING_SUPPORT act on the model runtime, so + # TLLM_KV_CACHE_MANAGER_V2_BACKEND / TRTLLM_ENABLE_PDL / + # NCCL_GRAPH_MIXING_SUPPORT act on the model runtime, so # they belong on the ctx and gen workers only; the disagg proxy server does # not read them. TRTLLM_SERVER_DISABLE_GC / TRTLLM_WORKER_DISABLE_GC are set # on both sides. worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 - TRTLLM_ENABLE_PDL=1 NCCL_GRAPH_MIXING_SUPPORT=0 TLLM_K3_MLA_GEN_BACKEND=trtllm-gen - TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp ENROOT_ALLOW_DEV=yes + TRTLLM_ENABLE_PDL=1 NCCL_GRAPH_MIXING_SUPPORT=0 TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp + ENROOT_ALLOW_DEV=yes server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false diff --git a/tests/unittest/_torch/attention/fmha_test_utils.py b/tests/unittest/_torch/attention/fmha_test_utils.py index 01e5609a686d..f109d523a170 100644 --- a/tests/unittest/_torch/attention/fmha_test_utils.py +++ b/tests/unittest/_torch/attention/fmha_test_utils.py @@ -31,7 +31,6 @@ def __init__(self, local_layer_idx: int = 0) -> None: self.num_heads = 1 self.num_kv_heads = 1 self.predicted_tokens_per_seq = 1 - self.flashinfer_mla_backend = None self.has_fp8_kv_cache = False self.local_layer_idx = local_layer_idx diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py index 9f4da22b2f29..019b790a762f 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py @@ -70,7 +70,6 @@ def test_on_update_kv_lens_invalidates_base_mla_state() -> None: md.enable_flash_mla = False md._mla_scheduler_buffers_valid = True md._mla_ctx_cu_seqlens_valid = True - md._cute_dsl_mla_staging_key = object() md._invalidate_pool_view_cache = Mock() md._num_tokens = 0 md._num_generations = 0 @@ -83,7 +82,6 @@ def test_on_update_kv_lens_invalidates_base_mla_state() -> None: assert not md._mla_scheduler_buffers_valid assert not md._mla_ctx_cu_seqlens_valid - assert md._cute_dsl_mla_staging_key is None def test_on_update_kv_lens_rebuilds_stale_map() -> None: diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index e008f060d9e6..1857d516e4f6 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -1,12 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import ast -import inspect -import textwrap -from collections.abc import Callable from types import SimpleNamespace -from typing import TypeAlias import pytest import torch @@ -19,39 +14,16 @@ FlashInferTrtllmGenFmha, _get_multi_ctas_kv_counter_size, ) -from tensorrt_llm._torch.attention.backends.fmha.interface import _CuteDslMlaStagingKey from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, ) -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.autotuner import AutoTuner from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -class _AttentionStub: - def __init__( - self, - *, - is_mla_enable: bool, - has_fp8_kv_cache: bool, - flashinfer_mla_backend: str | None = None, - local_layer_idx: int = 0, - ) -> None: - self.is_mla_enable = is_mla_enable - self.has_fp8_kv_cache = has_fp8_kv_cache - self.flashinfer_mla_backend = flashinfer_mla_backend - self.kv_lora_rank = 512 if is_mla_enable else None - self.head_dim = 576 - self.v_head_dim = 512 if is_mla_enable else None - self.local_layer_idx = local_layer_idx - - -_MlaBackendPolicy: TypeAlias = Callable[[str, SimpleNamespace, int], str] - - def test_flashinfer_uses_v2_page_index_upper_bound_directly() -> None: calls: list[tuple[int, object]] = [] bounds = iter((97, 101)) @@ -286,30 +258,6 @@ def generation_preprocess(*args: object) -> tuple[object, ...]: assert len(decode_calls) == 1 -def test_flashinfer_cute_dsl_mla_backend_rejects_fp8_kv_cache() -> None: - attn = _AttentionStub( - is_mla_enable=True, - has_fp8_kv_cache=True, - flashinfer_mla_backend="cute-dsl", - ) - - with pytest.raises(ValueError, match="does not support FP8 KV cache"): - FlashInferTrtllmGenFmha(attn) - - -@pytest.mark.parametrize("configured_backend", ["cute-dsl", "trtllm-gen"]) -def test_standalone_cute_dsl_mla_defers_to_explicit_flashinfer_backend( - configured_backend: str, -) -> None: - attn = _AttentionStub( - is_mla_enable=True, - has_fp8_kv_cache=False, - flashinfer_mla_backend=configured_backend, - ) - - assert not CuteDslMlaFmha.is_available(attn) - - def _cute_dsl_mla_helix_support( monkeypatch: pytest.MonkeyPatch, *, @@ -399,112 +347,7 @@ def test_cute_dsl_mla_rejects_invalid_helix_contract( assert reason in actual_reason -def test_flashinfer_mla_backend_defaults_to_trtllm_gen() -> None: - attn = _AttentionStub( - is_mla_enable=True, - has_fp8_kv_cache=False, - ) - - assert FlashInferTrtllmGenFmha(attn)._mla_backend == "trtllm-gen" - - -def test_mla_scheduler_invalidation_resets_cute_dsl_staging_key() -> None: - metadata = object.__new__(TrtllmAttentionMetadata) - metadata._mla_scheduler_buffers_valid = True - metadata._mla_ctx_cu_seqlens_valid = True - metadata._cute_dsl_mla_staging_key = _CuteDslMlaStagingKey( - is_capturing=True, - workspace_ptr=1, - block_tables_ptr=2, - block_tables_shape=(3, 4), - sequence_lengths_ptr=5, - sequence_lengths_offset=6, - batch_beam=7, - padded_num_pages=8, - ) - - metadata._invalidate_mla_scheduler_buffers() - - assert not metadata._mla_scheduler_buffers_valid - assert not metadata._mla_ctx_cu_seqlens_valid - assert metadata._cute_dsl_mla_staging_key is None - - -def test_flashinfer_mla_backend_rejects_unknown_backend() -> None: - attn = _AttentionStub( - is_mla_enable=True, - has_fp8_kv_cache=False, - flashinfer_mla_backend="cutedsl", - ) - - with pytest.raises(ValueError, match="flashinfer_mla_backend must be one of"): - FlashInferTrtllmGenFmha(attn) - - -def _make_fmha( - requested_backend: str, - mla_backend_policy: _MlaBackendPolicy | None, -) -> FlashInferTrtllmGenFmha: - fmha = object.__new__(FlashInferTrtllmGenFmha) - fmha._mla_backend = requested_backend - # ``Fmha.attn`` is a read-only property that dereferences ``_attn_ref`` - # (normally a weakref to the owning TrtllmAttention). SimpleNamespace is - # not weak-referenceable, so stand in with a closure of the same shape. - attn = SimpleNamespace(mla_backend_policy=mla_backend_policy) - fmha._attn_ref = lambda: attn - return fmha - - -@pytest.mark.parametrize("requested_backend", ["cute-dsl", "trtllm-gen"]) -@pytest.mark.parametrize( - ("num_contexts", "num_generations", "num_gen_tokens"), - [ - (0, 4, 4), # generation-only, one token per request - (1, 3, 3), # mixed context/generation batch - (0, 4, 8), # multi-token generation (speculative verification) - ], -) -def test_flashinfer_mla_backend_default_matches_static_selection( - requested_backend: str, - num_contexts: int, - num_generations: int, - num_gen_tokens: int, -) -> None: - """Without an installed policy the static backend is used for every batch - composition, matching the behavior before the policy hook existed.""" - fmha = _make_fmha(requested_backend, mla_backend_policy=None) - - assert ( - fmha._get_effective_mla_backend( - SimpleNamespace( - num_contexts=num_contexts, - num_generations=num_generations, - ), - num_gen_tokens, - ) - == requested_backend - ) - - -def test_flashinfer_mla_backend_policy_hook_is_consulted() -> None: - calls: list[tuple[str, SimpleNamespace, int]] = [] - - def policy( - requested_backend: str, - meta: SimpleNamespace, - num_gen_tokens: int, - ) -> str: - calls.append((requested_backend, meta, num_gen_tokens)) - return "trtllm-gen" - - fmha = _make_fmha("cute-dsl", mla_backend_policy=policy) - meta = SimpleNamespace(num_contexts=0, num_generations=4) - - assert fmha._get_effective_mla_backend(meta, 4) == "trtllm-gen" - assert calls == [("cute-dsl", meta, 4)] - - -# The six tests below guard the MLA generation perf gate that #15300 removed as +# The tests below guard the MLA generation perf gate that #15300 removed as # refactoring collateral, costing ~3% output token throughput on DeepSeek-V3-family # and Kimi-K2 MLA decode at the default tokens_per_block. They deliberately call the # checker instead of asserting on SLOWER_MLA_GENERATION_KERNELS itself: a test that @@ -520,7 +363,6 @@ def test_mla_generation_declines_slower_trtllm_gen_decode_kernel() -> None: supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( head_size=576, tokens_per_block=32, - mla_backend="trtllm-gen", kv_lora_rank=512, qk_rope_head_dim=64, ) @@ -538,21 +380,6 @@ def test_mla_generation_gate_is_scoped_to_one_page_size(tokens_per_block: int) - supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( head_size=576, tokens_per_block=tokens_per_block, - mla_backend="trtllm-gen", - kv_lora_rank=512, - qk_rope_head_dim=64, - ) - assert supported, reason - assert reason == "" - - -def test_mla_generation_gate_is_scoped_to_the_trtllm_gen_backend() -> None: - # The gated kernel is the trtllm-gen one; the cute-dsl MLA decode path shares - # this class and these head dims, and must stay selectable. - supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( - head_size=576, - tokens_per_block=32, - mla_backend="cute-dsl", kv_lora_rank=512, qk_rope_head_dim=64, ) @@ -560,57 +387,11 @@ def test_mla_generation_gate_is_scoped_to_the_trtllm_gen_backend() -> None: assert reason == "" -def test_mla_generation_gate_declines_a_policy_downgrade_to_trtllm_gen() -> None: - # A cute-dsl config whose per-batch policy downgrades to trtllm-gen (K3 does so - # for mixed batches and for speculative verification) still runs the gated - # kernel, so the gate must fire on the *effective* backend. Reading the static - # self._mla_backend here would let the slower kernel through. - fmha = _make_fmha("cute-dsl", mla_backend_policy=lambda *_: "trtllm-gen") - meta = SimpleNamespace(num_contexts=1, num_generations=3) - effective = fmha._get_effective_mla_backend(meta, 3) - assert effective == "trtllm-gen" - - supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( - head_size=576, - tokens_per_block=32, - mla_backend=effective, - kv_lora_rank=512, - qk_rope_head_dim=64, - ) - assert not supported - assert "slower" in reason - - -def test_mla_generation_gate_reads_the_effective_mla_backend() -> None: - # The composition above is only load-bearing if the production call site feeds - # the gate the effective backend. Assert that structurally: reverting to the - # static self._mla_backend is a one-word change, and no behavioural test here - # would catch it because driving _is_supported_with_reason needs a full - # metadata/forward-args stub. - source = textwrap.dedent(inspect.getsource(FlashInferTrtllmGenFmha._is_supported_with_reason)) - calls = [ - node - for node in ast.walk(ast.parse(source)) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Attribute) - and node.func.attr == "_check_mla_generation_support" - ] - assert len(calls) == 1, "expected exactly one MLA generation gate call site" - kwargs = {kw.arg: kw.value for kw in calls[0].keywords} - passed = kwargs.get("mla_backend") - assert passed is not None, "gate call site lost its mla_backend argument" - assert ( - isinstance(passed, ast.Call) - and getattr(passed.func, "attr", None) == "_get_effective_mla_backend" - ), f"gate must receive the effective backend, got {ast.dump(passed)}" - - def test_mla_generation_allows_other_supported_head_dims() -> None: # (320, 256) is unaffected at every page size. supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( head_size=320, tokens_per_block=32, - mla_backend="trtllm-gen", kv_lora_rank=256, qk_rope_head_dim=64, ) diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py deleted file mode 100644 index 33f32cf693c5..000000000000 --- a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py +++ /dev/null @@ -1,178 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from types import SimpleNamespace -from typing import Optional - -import pytest -import torch - -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention -from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import ( - _KIMI_K3_MLA_GEN_BACKEND_ENV, - _kimi_k3_mla_decode_backend_policy, - _select_mla_generation_backend, - _validate_mla_generation_backend, -) -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig - - -@pytest.mark.parametrize( - ("configured_backend", "expected_backend"), - [(None, "cute-dsl"), ("trtllm-gen", "trtllm-gen")], -) -def test_select_kimi_k3_mla_generation_backend( - monkeypatch: pytest.MonkeyPatch, - configured_backend: Optional[str], - expected_backend: str, -) -> None: - if configured_backend is None: - monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) - else: - monkeypatch.setenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, configured_backend) - - assert _select_mla_generation_backend(None) == expected_backend - - -@pytest.mark.parametrize("invalid_backend", ["cutedsl", "", "CUTE-DSL "]) -def test_select_kimi_k3_mla_generation_backend_rejects_invalid_env( - monkeypatch: pytest.MonkeyPatch, - invalid_backend: str, -) -> None: - """An invalid env value must fail at read time with a message naming the - knob, not propagate until attention-backend construction.""" - monkeypatch.setenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, invalid_backend) - - with pytest.raises(ValueError, match=_KIMI_K3_MLA_GEN_BACKEND_ENV): - _select_mla_generation_backend(None) - - -def test_select_kimi_k3_mla_generation_backend_uses_trtllm_gen_for_fp8_kv_cache( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, "cute-dsl") - quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) - - assert _select_mla_generation_backend(quant_config) == "trtllm-gen" - - -@pytest.mark.parametrize( - ("backend", "num_heads"), - [ - ("cute-dsl", 96), - ("trtllm-gen", 6), - ("trtllm-gen", 64), - ("trtllm-gen", 128), - ], -) -def test_validate_mla_generation_backend_accepts_runnable_configs( - backend: str, num_heads: int -) -> None: - _validate_mla_generation_backend(backend, num_heads) - - -@pytest.mark.parametrize("num_heads", [65, 96, 127]) -def test_validate_mla_generation_backend_rejects_trtllm_gen_mid_head_counts( - num_heads: int, -) -> None: - """trtllm-gen with 64 < H < 128 per-rank heads can never run any batch.""" - with pytest.raises(ValueError, match="query heads per rank"): - _validate_mla_generation_backend("trtllm-gen", num_heads) - - -def test_fp8_kv_cache_with_attention_dp_head_count_fails_fast( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """The FP8-KV override forces trtllm-gen; with 96 per-rank heads - (attention-DP replication) construction must fail fast rather than - crash inside FlashInfer at attention warmup.""" - monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) - quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) - - backend = _select_mla_generation_backend(quant_config) - with pytest.raises(ValueError, match="FP8-KV-cache"): - _validate_mla_generation_backend(backend, num_heads=96) - - -@pytest.mark.skipif( - not torch.cuda.is_available(), - reason="KimiK3MLAAttention builds real TRTLLM attention backends", -) -def test_kimi_k3_mla_construction_fails_fast_for_fp8_kv_with_attention_dp( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Constructor-level regression for the FP8-KV + attention-DP conflict: - the override selects trtllm-gen, attention-DP keeps all 96 query heads - on every rank, and `KimiK3MLAAttention.__init__` itself must raise — - before attention warmup ever runs.""" - monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) - model_config = ModelConfig( - skip_create_weights_in_init=True, - quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8), - mapping=Mapping(world_size=4, tp_size=4, rank=0, enable_attention_dp=True), - ) - - with pytest.raises(ValueError, match="query heads per rank"): - # Kimi K3 MLA geometry (96 Q heads); small max positions keep the - # identity-RoPE table allocation negligible. - KimiK3MLAAttention( - hidden_size=7168, - num_heads=96, - q_lora_rank=1536, - kv_lora_rank=512, - qk_nope_head_dim=128, - qk_rope_head_dim=64, - v_head_dim=128, - rms_norm_eps=1e-6, - dtype=torch.bfloat16, - layer_idx=0, - max_position_embeddings=256, - model_config=model_config, - ) - - -@pytest.mark.parametrize( - ( - "requested_backend", - "num_contexts", - "num_generations", - "num_gen_tokens", - "num_heads", - "expected_backend", - ), - [ - ("cute-dsl", 0, 4, 4, 96, "cute-dsl"), - ("cute-dsl", 1, 3, 3, 12, "trtllm-gen"), - ("cute-dsl", 1, 3, 3, 96, "cute-dsl"), - # Generation-only multi-token (speculative verification): H=96 must - # stay on cute-dsl — trtllm-gen's decode gate rejects that head - # count, so falling back fails engine init. Smaller per-rank head - # counts (non-attention-DP shapes) keep the tuning fallback. - ("cute-dsl", 0, 4, 8, 96, "cute-dsl"), - ("cute-dsl", 0, 4, 8, 12, "trtllm-gen"), - ("trtllm-gen", 1, 3, 3, 96, "trtllm-gen"), - ], -) -def test_kimi_k3_mla_decode_backend_policy_by_batch_shape( - requested_backend: str, - num_contexts: int, - num_generations: int, - num_gen_tokens: int, - num_heads: int, - expected_backend: str, -) -> None: - """K3 falls back outside plain decode except when H=96 breaks trtllm-gen.""" - assert ( - _kimi_k3_mla_decode_backend_policy( - requested_backend, - SimpleNamespace( - num_contexts=num_contexts, - num_generations=num_generations, - ), - num_gen_tokens, - num_heads=num_heads, - ) - == expected_backend - )