[TRTLLM-15033][fix] Revert #17800 FlashInfer CuTeDSL MLA dispatch

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`

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

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.

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