[None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel

tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py

@@ -16,12 +16,23 @@
 try:
     # A missing build raises ImportError; a CuTe/CUTLASS mismatch raises
     # RuntimeError (mirror FlashInfer's own guard) -> Triton fallback.
+    # gated_delta_rule: T=1 decode entry (dispatches to the wide_vec fast path
+    # when B*HV is large). gated_delta_rule_mtp: T>=1 with batch-scoped
+    # intermediate_states_buffer and disable_state_update support, used by the
+    # speculative-decoding target-verify path.
     from flashinfer.gdn_kernels.gdn_decode_bf16_state import \
         gated_delta_rule as _fi_gdn_decode_bf16_state_t1
+    from flashinfer.gdn_kernels.gdn_decode_bf16_state import \
+        gated_delta_rule_mtp as _fi_gdn_decode_bf16_state_mtp
     _FLASHINFER_GDN_BF16_STATE_AVAILABLE = True
 except (ImportError, RuntimeError):
     _FLASHINFER_GDN_BF16_STATE_AVAILABLE = False
 
+# Max per-sequence token count served by the FlashInfer MTP verify kernel; the
+# parity test (test_flashinfer_gdn_verify.py) covers T=1..8 against the Triton
+# reference. Longer drafts fall back to the Triton recurrent kernel.
+_FI_GDN_MAX_MTP_T = 8
+
 
 @triton.heuristics({
     "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None,
@@ -288,6 +299,103 @@ def _flashinfer_gdn_decode(
     return output.reshape(1, T_total, HV, -1)
 
 
+def _can_use_flashinfer_gdn_verify(
+    initial_state_source: Optional[torch.Tensor],
+    head_k_dim: int,
+    head_v_dim: int,
+    draft_token_num: int,
+) -> bool:
+    """Whether the FlashInfer MTP kernel should serve the speculative verify step.
+
+    Default ON when eligible; set ``TRTLLM_FLA_DISABLE_FLASHINFER_GDN_VERIFY=1``
+    to force the Triton recurrent verify kernel (``TRTLLM_FLA_DISABLE_FLASHINFER_GDN=1``
+    disables all FlashInfer GDN decode paths, including this one). The same
+    constraints as the decode path apply (bf16 state pool, K==V==128, supported
+    arch, FI MTP API available) plus a per-sequence draft length in
+    [1, _FI_GDN_MAX_MTP_T]; longer drafts fall back to Triton.
+    """
+    if os.environ.get("TRTLLM_FLA_DISABLE_FLASHINFER_GDN", "0") == "1":
+        return False
+    if os.environ.get("TRTLLM_FLA_DISABLE_FLASHINFER_GDN_VERIFY", "0") == "1":
+        return False
+    if not _FLASHINFER_GDN_BF16_STATE_AVAILABLE:
+        return False
+    if not is_flashinfer_gdn_supported_arch():
+        return False
+    if initial_state_source is None or initial_state_source.dtype != torch.bfloat16:
+        return False
+    if head_k_dim != 128 or head_v_dim != 128:
+        return False
+    if not (1 <= draft_token_num <= _FI_GDN_MAX_MTP_T):
+        return False
+    return True
+
+
+def _flashinfer_gdn_verify(
+    A_log: torch.Tensor,
+    a: torch.Tensor,
+    dt_bias: torch.Tensor,
+    softplus_beta: float,
+    softplus_threshold: float,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    b: torch.Tensor,
+    initial_state_source: torch.Tensor,
+    initial_state_indices: torch.Tensor,
+    intermediate_states_buffer: torch.Tensor,
+    scale: float,
+    use_qk_l2norm_in_kernel: bool,
+    output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """GDN MTP *verify* via the FlashInfer bf16-state kernel.
+
+    Inputs are batched ``[N, draft_token_num, H, D]``. The kernel gathers the
+    initial state from the pool via ``initial_state_indices`` (no host-side
+    gather copy), writes the SSM state after each draft token into the
+    batch-scoped ``intermediate_states_buffer`` (``[N, draft_token_num, HV, V,
+    K]``, matching the Triton verify kernel) and leaves the live state pool
+    untouched (``disable_state_update``) so the cache manager selects the
+    accepted-position state afterwards. Returns the attention output
+    ``[N, draft_token_num, HV, V]``.
+    """
+    logger.info_once(
+        "Using FlashInfer CuTe-DSL kernel for GDN MTP verify "
+        "(bf16 state, K=V=128)",
+        key="flashinfer_gdn_verify")
+    N, T = q.shape[0], q.shape[1]
+    HV, V = v.shape[2], v.shape[3]
+    output = (output.view(N, T, HV, V) if output is not None else q.new_empty(
+        N, T, HV, V))
+    # The FI CuTe-DSL kernel asserts 32-byte data alignment on every tensor
+    # argument. The int32 index tensor may be a slice of a larger buffer
+    # (e.g. state_indices_d = cache_indices[num_prefills:]) whose 4*offset
+    # storage offset breaks that; .int() is a no-op for int32, so realign
+    # with an explicit copy when needed.
+    initial_state_indices = initial_state_indices.int()
+    if initial_state_indices.data_ptr() % 32 != 0:
+        initial_state_indices = initial_state_indices.clone()
+    _fi_gdn_decode_bf16_state_mtp(
+        A_log=A_log,
+        a=a,
+        dt_bias=dt_bias,
+        softplus_beta=softplus_beta,
+        softplus_threshold=softplus_threshold,
+        q=q,
+        k=k,
+        v=v,
+        b=b,
+        initial_state_source=initial_state_source,
+        initial_state_indices=initial_state_indices,
+        intermediate_states_buffer=intermediate_states_buffer,
+        disable_state_update=True,
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+        scale=scale,
+        output=output,
+    )
+    return output
+
+
 def fused_sigmoid_gating_delta_rule_update(
     A_log: torch.Tensor,
     a: torch.Tensor,

tensorrt_llm/_torch/modules/mamba/gdn_mixer.py

@@ -16,6 +16,8 @@
 
 from tensorrt_llm._torch.modules.fla.fused_recurrent import fused_recurrent_gated_delta_rule_update
 from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import (
+    _can_use_flashinfer_gdn_verify,
+    _flashinfer_gdn_verify,
     fused_sigmoid_gating_delta_rule_update,
 )
 from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import use_cpp_mamba_cache_manager
@@ -667,6 +669,46 @@ def forward_decode(
 
             a = a.reshape(num_decodes, draft_token_num, -1)
             b = b.reshape(num_decodes, draft_token_num, -1)
+
+            # Prefer the FlashInfer MTP kernel (raw a/b gating in-kernel,
+            # initial state gathered from the pool via cache indices, per-step
+            # intermediate states written to the batch-scoped [:num_decodes]
+            # prefix consumed by update_mamba_states()); fall back to the
+            # Triton recurrent kernel when unavailable.
+            if _can_use_flashinfer_gdn_verify(
+                ssm_states, self.head_k_dim, self.head_v_dim, draft_token_num
+            ):
+                output_d = None
+                if output is not None:
+                    output_d = output.view(
+                        num_decodes,
+                        draft_token_num,
+                        self.num_v_heads // self.attn_tp_size,
+                        self.head_v_dim,
+                    )
+                return _flashinfer_gdn_verify(
+                    A_log=self.A_log,
+                    a=a,
+                    dt_bias=self.dt_bias,
+                    softplus_beta=1.0,
+                    softplus_threshold=20.0,
+                    q=query,
+                    k=key,
+                    v=value,
+                    b=b,
+                    initial_state_source=ssm_states,
+                    initial_state_indices=cache_indices[:num_decodes],
+                    intermediate_states_buffer=intermediate_ssm_states[:num_decodes],
+                    scale=self.head_k_dim**-0.5,
+                    use_qk_l2norm_in_kernel=True,
+                    output=output_d,
+                ).view(
+                    1,
+                    num_decodes * draft_token_num,
+                    self.num_v_heads // self.attn_tp_size,
+                    self.head_v_dim,
+                )
+
             beta = b.sigmoid()
             g = fused_gdn_gating(
                 self.A_log,
@@ -921,41 +963,67 @@ def forward_extend(
 
             a_d = a_d.reshape(num_decodes, draft_token_num, -1)
             b_d = b_d.reshape(num_decodes, draft_token_num, -1)
-            beta_d = b_d.sigmoid()
-            g_d = fused_gdn_gating(
-                self.A_log,
-                a_d.view(num_decodes * draft_token_num, -1),
-                self.dt_bias,
-            ).reshape(num_decodes, draft_token_num, -1)
-
-            recurrent_state_source = ssm_states[state_indices_d]
-            recurrent_state_indices = torch.arange(
-                num_decodes, dtype=torch.int32, device=state_indices_d.device
-            )
+            out_v_heads = self.num_v_heads // self.attn_tp_size
 
             output_d = None
             if output is not None:
                 output_d = output[:, num_prefill_tokens:, :, :].view(
                     num_decodes,
                     draft_token_num,
-                    self.num_v_heads // self.attn_tp_size,
+                    out_v_heads,
                     self.head_v_dim,
                 )
 
-            attn_out_decode = fused_recurrent_gated_delta_rule_update(
-                q=query_d,
-                k=key_d,
-                v=value_d,
-                g=g_d,
-                beta=beta_d,
-                initial_state_source=recurrent_state_source,
-                initial_state_indices=recurrent_state_indices,
-                use_qk_l2norm_in_kernel=True,
-                disable_state_update=True,
-                intermediate_states_buffer=intermediate_ssm_states,
-                cache_steps=draft_token_num,
-                output=output_d,
-            ).view(1, num_decode_tokens, self.num_v_heads // self.attn_tp_size, self.head_v_dim)
+            if _can_use_flashinfer_gdn_verify(
+                ssm_states, self.head_k_dim, self.head_v_dim, draft_token_num
+            ):
+                # FI gathers the initial state from the pool via state_indices_d
+                # (no host gather) and writes batch-scoped intermediate states;
+                # the [:num_decodes] prefix matches update_mamba_states()'s rows.
+                attn_out_decode = _flashinfer_gdn_verify(
+                    A_log=self.A_log,
+                    a=a_d,
+                    dt_bias=self.dt_bias,
+                    softplus_beta=1.0,
+                    softplus_threshold=20.0,
+                    q=query_d,
+                    k=key_d,
+                    v=value_d,
+                    b=b_d,
+                    initial_state_source=ssm_states,
+                    initial_state_indices=state_indices_d,
+                    intermediate_states_buffer=intermediate_ssm_states[:num_decodes],
+                    scale=self.head_k_dim**-0.5,
+                    use_qk_l2norm_in_kernel=True,
+                    output=output_d,
+                ).reshape(1, num_decode_tokens, out_v_heads, self.head_v_dim)
+            else:
+                beta_d = b_d.sigmoid()
+                g_d = fused_gdn_gating(
+                    self.A_log,
+                    a_d.view(num_decodes * draft_token_num, -1),
+                    self.dt_bias,
+                ).reshape(num_decodes, draft_token_num, -1)
+
+                recurrent_state_source = ssm_states[state_indices_d]
+                recurrent_state_indices = torch.arange(
+                    num_decodes, dtype=torch.int32, device=state_indices_d.device
+                )
+
+                attn_out_decode = fused_recurrent_gated_delta_rule_update(
+                    q=query_d,
+                    k=key_d,
+                    v=value_d,
+                    g=g_d,
+                    beta=beta_d,
+                    initial_state_source=recurrent_state_source,
+                    initial_state_indices=recurrent_state_indices,
+                    use_qk_l2norm_in_kernel=True,
+                    disable_state_update=True,
+                    intermediate_states_buffer=intermediate_ssm_states,
+                    cache_steps=draft_token_num,
+                    output=output_d,
+                ).view(1, num_decode_tokens, out_v_heads, self.head_v_dim)
 
             if output is not None:
                 return output