[None][fix] Unwaive supported attention backend cases

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -43,6 +43,80 @@
 _FORCE_RAGGED_FA2 = False
 """Used for testing."""
 
+_MAX_CUDA_THREADS_PER_BLOCK = 1024
+
+
+def _slice_paged_kv_cache_heads(
+    paged_kv_cache: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    start: int,
+    end: int,
+    kv_layout: str,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    if kv_layout == "HND":
+        head_axis = 2
+    elif kv_layout == "NHD":
+        head_axis = 3
+    else:
+        raise ValueError(f"Unsupported kv_layout: {kv_layout}")
+
+    if isinstance(paged_kv_cache, tuple):
+        head_axis -= 1
+        index = [slice(None)] * 4
+        index[head_axis] = slice(start, end)
+        return tuple(cache[tuple(index)] for cache in paged_kv_cache)
+
+    index = [slice(None)] * 5
+    index[head_axis] = slice(start, end)
+    return paged_kv_cache[tuple(index)]
+
+
+def _append_paged_kv_cache(
+    append_key: torch.Tensor,
+    append_value: torch.Tensor,
+    batch_indices: torch.Tensor,
+    positions: torch.Tensor,
+    paged_kv_cache: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    kv_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    kv_last_page_len: torch.Tensor,
+    kv_layout: str = "NHD",
+) -> None:
+    """Split FlashInfer paged-KV appends that exceed CUDA's CTA limit."""
+    head_dim = append_key.shape[-1]
+    vec_size = max(16 // append_key.element_size(), head_dim // 32)
+    threads_per_head = head_dim // vec_size
+    max_heads_per_launch = _MAX_CUDA_THREADS_PER_BLOCK // threads_per_head
+
+    num_kv_heads = append_key.shape[1]
+    if num_kv_heads <= max_heads_per_launch:
+        flashinfer.page.append_paged_kv_cache(
+            append_key=append_key,
+            append_value=append_value,
+            batch_indices=batch_indices,
+            positions=positions,
+            paged_kv_cache=paged_kv_cache,
+            kv_indices=kv_indices,
+            kv_indptr=kv_indptr,
+            kv_last_page_len=kv_last_page_len,
+            kv_layout=kv_layout,
+        )
+        return
+
+    for start in range(0, num_kv_heads, max_heads_per_launch):
+        end = min(start + max_heads_per_launch, num_kv_heads)
+        flashinfer.page.append_paged_kv_cache(
+            append_key=append_key[:, start:end],
+            append_value=append_value[:, start:end],
+            batch_indices=batch_indices,
+            positions=positions,
+            paged_kv_cache=_slice_paged_kv_cache_heads(paged_kv_cache, start,
+                                                       end, kv_layout),
+            kv_indices=kv_indices,
+            kv_indptr=kv_indptr,
+            kv_last_page_len=kv_last_page_len,
+            kv_layout=kv_layout,
+        )
+
 
 @dataclass(kw_only=True, frozen=True)
 class FlashInferMultiItemParams:
@@ -1748,7 +1822,7 @@ def forward_impl(
                 f"KV cache dtype {kv_cache.dtype} does not match k/v dtype {k.dtype}/{v.dtype}"
             )
 
-            flashinfer.page.append_paged_kv_cache(
+            _append_paged_kv_cache(
                 append_key=k,
                 append_value=v,
                 batch_indices=metadata.batch_indices,

tensorrt_llm/_torch/attention_backend/fmha/fallback.py

@@ -36,6 +36,7 @@
         "attention_mask_data",  # custom-mask code path
         "out_scale_sf",  # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path
         "skip_mla_rope_generation",  # handled in ``TrtllmAttention.forward`` for the test-only MLA path
+        "timestep",  # used to populate skip-softmax params in ``TrtllmAttention.forward``
     }
 )
 

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py

@@ -408,9 +408,6 @@ class FlashInferTrtllmGenFmha(PhasedFmha):
         (320, 256),
         (576, 512),
     }
-    MISSING_MLA_GENERATION_KERNELS = {
-        (576, 512, 32),
-    }
 
     def __init__(self, attn: "TrtllmAttention"):
         super().__init__(attn)
@@ -529,7 +526,6 @@ def _get_attention_chunk_size(attn: "TrtllmAttention") -> int:
     def _check_mla_generation_support(
         cls,
         head_size: int,
-        tokens_per_block: int,
         kv_lora_rank: Optional[int],
         qk_rope_head_dim: Optional[int],
     ) -> Tuple[bool, str]:
@@ -567,14 +563,6 @@ def _check_mla_generation_support(
                 f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.",
             )
 
-        if (head_dim_qk, head_dim_v, tokens_per_block) in cls.MISSING_MLA_GENERATION_KERNELS:
-            return (
-                False,
-                f"[Generation][MLA] Missing TRTLLM-GEN decode kernel for "
-                f"headDimQk={head_dim_qk}, headDimV={head_dim_v}, "
-                f"tokens_per_block={tokens_per_block}.",
-            )
-
         return True, ""
 
     def is_supported(
@@ -765,7 +753,6 @@ def _is_supported_with_reason(
             if is_mla_enable:
                 supported, reason = self._check_mla_generation_support(
                     head_size=attn.head_dim,
-                    tokens_per_block=tokens_per_block,
                     kv_lora_rank=attn.kv_lora_rank,
                     qk_rope_head_dim=attn.qk_rope_head_dim,
                 )
@@ -1192,7 +1179,7 @@ def run_mla_generation(
         batch_beam = params.num_requests * meta.beam_width
         if params.attention_input is None:
             raise RuntimeError("MLA generation requires attention_input.")
-        kv_cache, block_tables = thop.build_trtllm_gen_kv_cache_metadata(
+        kv_cache, block_tables, _kv_scale_pool = thop.build_trtllm_gen_kv_cache_metadata(
             meta.host_kv_cache_pool_pointers,  # host_kv_cache_pool_pointers
             meta.host_kv_cache_pool_mapping,  # host_kv_cache_pool_mapping
             meta.kv_cache_block_offsets,  # kv_cache_block_offsets
@@ -1222,9 +1209,43 @@ def run_mla_generation(
         mla_head_dim_qk = kv_lora_rank + qk_rope_head_dim
         q_len_per_req = params.num_tokens // batch_beam if batch_beam > 0 else 1
 
-        query = params.qkv_input.view(batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk)
+        if QuantMode(attn.quant_mode).has_fp8_kv_cache():
+            quant_q_buffer = fwd.quant_q_buffer
+            bmm1_scale_buffer = fwd.mla_bmm1_scale
+            bmm2_scale_buffer = fwd.mla_bmm2_scale
+            if quant_q_buffer is None or bmm1_scale_buffer is None or bmm2_scale_buffer is None:
+                raise RuntimeError(
+                    "FP8 MLA generation requires quant_q_buffer, "
+                    "mla_bmm1_scale, and mla_bmm2_scale."
+                )
 
-        bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim))
+            expected_q_elements = params.num_tokens * attn.num_heads * mla_head_dim_qk
+            if quant_q_buffer.numel() < expected_q_elements:
+                raise RuntimeError(
+                    f"FP8 MLA quant_q_buffer has {quant_q_buffer.numel()} elements; "
+                    f"expected at least {expected_q_elements}."
+                )
+            if bmm1_scale_buffer.dtype != torch.float32 or bmm1_scale_buffer.numel() < 1:
+                raise RuntimeError("FP8 MLA bmm1 scale must contain a float32 value.")
+            if bmm2_scale_buffer.dtype != torch.float32 or bmm2_scale_buffer.numel() < 1:
+                raise RuntimeError("FP8 MLA bmm2 scale must contain a float32 value.")
+
+            query = (
+                quant_q_buffer.view(torch.uint8)
+                .flatten()[:expected_q_elements]
+                .view(torch.float8_e4m3fn)
+                .view(batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk)
+            )
+            # FlashInfer converts tensor BMM1 scales to log2 internally. The
+            # producer stores the regular scale at index 0 and log2 at index 1.
+            bmm1_scale = bmm1_scale_buffer.flatten()[:1]
+            bmm2_scale = bmm2_scale_buffer.flatten()[:1]
+        else:
+            query = params.qkv_input.view(
+                batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk
+            )
+            bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim))
+            bmm2_scale = 1.0
         workspace_buffer = params.workspace.view(-1, 4)
         _clear_multi_ctas_kv_counter_workspace(
             workspace_buffer, attn.num_heads, meta.max_num_requests, self._multi_processor_count
@@ -1243,7 +1264,7 @@ def run_mla_generation(
             0,  # sparse_mla_top_k
             params.context_buf.view(batch_beam, q_len_per_req, attn.num_heads, kv_lora_rank),  # out
             bmm1_scale,  # bmm1_scale
-            1.0,  # bmm2_scale
+            bmm2_scale,  # bmm2_scale
             fwd.attention_sinks,  # sinks
             None,  # skip_softmax_threshold_scale_factor
             self._enable_pdl,  # enable_pdl

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1312,11 +1312,6 @@ def use_nvfp4_output(
             or metadata.runtime_features.has_speculative_draft_tokens
         ) if metadata.runtime_features else False
 
-        # This is a workaround for https://nvbugs/5624818
-        # Paged context FMHA is forced on SM90 for correctness
-        if get_sm_version() == 90:
-            use_paged_context_fmha = True
-
         return self._is_nvfp4_output_kernel_available(
             tokens_per_block=metadata.tokens_per_block,
             attention_mask=attention_mask,
@@ -1487,17 +1482,6 @@ def forward(
         # Cross-attention uses the THOP path; the trtllm-gen backend API does
         # not carry encoder K/V tensors yet.
 
-        # cpp/tensorrt_llm/thop/attentionOp.cpp enables mFP8ContextFMHA for an
-        # FP8 KV cache only when use_paged_context_fmha is true. Force paged
-        # context so QKV preprocessing and context FMHA use the FP8 path.
-        if self.has_fp8_kv_cache:
-            metadata.use_paged_context_fmha = True
-
-        # SM90 forces `use_paged_context_fmha` on for correctness
-        # (https://nvbugs/5624818).
-        if get_sm_version() == 90:
-            metadata.use_paged_context_fmha = True
-
         # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer).
         # Force paged context FMHA so QKV preprocessing writes Q to q_buf_2_.
         if (self.sparse_params is not None and getattr(