From 61dbc4f041c5c3fac1fcedf104f8ca3320a219bb Mon Sep 17 00:00:00 2001 From: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> Date: Mon, 6 Jul 2026 13:47:22 +0800 Subject: [PATCH] [None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> --- .../fla/fused_sigmoid_gating_recurrent.py | 108 ++++++++ .../_torch/modules/mamba/gdn_mixer.py | 120 +++++++-- .../mamba/test_flashinfer_gdn_verify.py | 254 ++++++++++++++++++ 3 files changed, 456 insertions(+), 26 deletions(-) create mode 100644 tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py diff --git a/tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py b/tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py index a37f281e3a4c..961189d97346 100644 --- a/tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py +++ b/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, diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 88fe4451675e..f81f92cbf945 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/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 diff --git a/tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py b/tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py new file mode 100644 index 000000000000..b7b30c6abcac --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_flashinfer_gdn_verify.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parity test for the GDN MTP *verify* path: FlashInfer ``gated_delta_rule_mtp`` +vs the Triton ``fused_recurrent_gated_delta_rule_update`` reference. + +The verify path (speculative decoding) runs the gated delta rule over +``draft_token_num`` tokens per sequence and must write the SSM state *after each +draft token* into an ``intermediate_states_buffer`` (so the cache manager can +later select the state at the accepted position), without updating the live +state pool (``disable_state_update=True``). + +This test asserts the FlashInfer MTP kernel produces the same attention output +AND the same per-step intermediate states as the Triton kernel, so the verify +branch in ``gdn_mixer`` can dispatch to FlashInfer with a Triton fallback. +""" + +import pytest +import torch + + +def _fi_mtp_available() -> bool: + if not torch.cuda.is_available(): + return False + from tensorrt_llm._utils import is_flashinfer_gdn_supported_arch + + if not is_flashinfer_gdn_supported_arch(): + return False + try: + from flashinfer.gdn_kernels.gdn_decode_bf16_state import gated_delta_rule_mtp # noqa: F401 + except Exception: + return False + return True + + +skip_unsupported = pytest.mark.skipif( + not _fi_mtp_available(), + reason="Requires SM90/SM100/SM103 and a FlashInfer build with " + "gdn_decode_bf16_state.gated_delta_rule_mtp", +) + + +@skip_unsupported +@pytest.mark.parametrize("draft_token_num", [1, 2, 3, 4, 5, 6, 7, 8]) +@pytest.mark.parametrize("num_decodes", [1, 3]) +@pytest.mark.parametrize("H,HV", [(4, 8), (2, 2)]) +def test_fi_mtp_verify_matches_triton(draft_token_num, num_decodes, H, HV): + from flashinfer.gdn_kernels.gdn_decode_bf16_state import gated_delta_rule_mtp + + from tensorrt_llm._torch.modules.fla.fused_recurrent import ( + fused_recurrent_gated_delta_rule_update, + ) + from tensorrt_llm._torch.modules.mamba.gdn_mixer import fused_gdn_gating + + torch.manual_seed(0) + dev = "cuda" + N, T, K, V = num_decodes, draft_token_num, 128, 128 + scale = K**-0.5 + + q = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + k = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + v = (torch.randn(N, T, HV, V, device=dev) * 0.1).to(torch.bfloat16) + a = torch.randn(N, T, HV, device=dev) * 0.1 + b = torch.randn(N, T, HV, device=dev) * 0.1 + A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() + dt_bias = torch.randn(HV, device=dev) * 0.1 + state_pool = (torch.randn(N, HV, V, K, device=dev) * 0.1).to(torch.bfloat16) + idx = torch.arange(N, device=dev, dtype=torch.int32) + + # --- Triton reference (gdn_mixer is_target_verify branch) --- + g = fused_gdn_gating(A_log, a.view(N * T, HV), dt_bias).view(N, T, HV) + beta = b.sigmoid() + buf_tri = torch.zeros(N, T, HV, V, K, device=dev, dtype=torch.bfloat16) + out_tri = fused_recurrent_gated_delta_rule_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state_source=state_pool.clone(), + initial_state_indices=idx, + use_qk_l2norm_in_kernel=True, + disable_state_update=True, + intermediate_states_buffer=buf_tri, + cache_steps=T, + ) + + # --- FlashInfer MTP verify --- + buf_fi = torch.zeros(N, T, HV, V, K, device=dev, dtype=torch.bfloat16) + out_fi = q.new_empty(N, T, HV, V) + gated_delta_rule_mtp( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=state_pool.clone(), + initial_state_indices=idx, + intermediate_states_buffer=buf_fi, + disable_state_update=True, + use_qk_l2norm_in_kernel=True, + scale=scale, + output=out_fi, + ) + + # bf16 recurrent accumulation: use bf16-appropriate tolerance. + torch.testing.assert_close(out_fi.float(), out_tri.float(), rtol=2e-2, atol=2e-2) + torch.testing.assert_close(buf_fi.float(), buf_tri.float(), rtol=2e-2, atol=2e-2) + + +@skip_unsupported +def test_fi_mtp_verify_buffer_is_batch_scoped(): + """FI requires ``intermediate_states_buffer`` to be batch-scoped (dim0 == B). + + Unlike the Triton kernel (which indexes a pool-scoped buffer by + ``initial_state_indices``), FI writes batch row ``i`` to buffer row ``i`` and + rejects a pool-sized buffer. gdn_mixer therefore passes the + ``[:num_decodes]`` prefix slice. This guards that contract so a future FI + bump that silently accepts a larger buffer (writing the wrong rows) is + caught. + """ + from flashinfer.gdn_kernels.gdn_decode_bf16_state import gated_delta_rule_mtp + + dev = "cuda" + N, T, H, HV, K, V = 2, 4, 4, 8, 128, 128 + pool = 8 # pool-scoped buffer larger than the batch + torch.manual_seed(0) + q = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + k = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + v = (torch.randn(N, T, HV, V, device=dev) * 0.1).to(torch.bfloat16) + a = torch.randn(N, T, HV, device=dev) * 0.1 + b = torch.randn(N, T, HV, device=dev) * 0.1 + A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() + dt_bias = torch.randn(HV, device=dev) * 0.1 + state_pool = (torch.randn(N, HV, V, K, device=dev) * 0.1).to(torch.bfloat16) + idx = torch.arange(N, device=dev, dtype=torch.int32) + buf_pool = torch.zeros(pool, T, HV, V, K, device=dev, dtype=torch.bfloat16) + out = q.new_empty(N, T, HV, V) + + with pytest.raises(AssertionError): + gated_delta_rule_mtp( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=state_pool, + initial_state_indices=idx, + intermediate_states_buffer=buf_pool, # dim0=pool != B=N + disable_state_update=True, + use_qk_l2norm_in_kernel=True, + scale=K**-0.5, + output=out, + ) + + +@skip_unsupported +def test_fi_mtp_verify_misaligned_index_slice(): + """Index slices with a non-32B-aligned storage offset must be realigned. + + In the mixed prefill+decode verify path, gdn_mixer passes + ``state_indices_d = cache_indices[num_prefills:]`` — an int32 view whose + 4*num_prefills-byte storage offset violates the FI kernel's 32-byte + alignment assert (``Misaligned Tensor data on argument`` at runtime). + ``_flashinfer_gdn_verify`` must copy such views before dispatch. + """ + from tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent import ( + _flashinfer_gdn_verify, + ) + + torch.manual_seed(0) + dev = "cuda" + N, T, H, HV, K, V = 2, 4, 4, 8, 128, 128 + q = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + k = (torch.randn(N, T, H, K, device=dev) * 0.1).to(torch.bfloat16) + v = (torch.randn(N, T, HV, V, device=dev) * 0.1).to(torch.bfloat16) + a = torch.randn(N, T, HV, device=dev) * 0.1 + b = torch.randn(N, T, HV, device=dev) * 0.1 + A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() + dt_bias = torch.randn(HV, device=dev) * 0.1 + state_pool = (torch.randn(N, HV, V, K, device=dev) * 0.1).to(torch.bfloat16) + buf = torch.zeros(N, T, HV, V, K, device=dev, dtype=torch.bfloat16) + + # int32 slice with a 4-byte storage offset (mimics cache_indices[1:]). + idx_buf = torch.arange(N + 1, device=dev, dtype=torch.int32) - 1 + idx_misaligned = idx_buf[1:] + assert idx_misaligned.data_ptr() % 32 != 0 + + out_mis = _flashinfer_gdn_verify( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=state_pool, + initial_state_indices=idx_misaligned, + intermediate_states_buffer=buf, + scale=K**-0.5, + use_qk_l2norm_in_kernel=True, + ) + + out_ref = _flashinfer_gdn_verify( + A_log=A_log, + a=a, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=b, + initial_state_source=state_pool, + initial_state_indices=idx_misaligned.clone(), + intermediate_states_buffer=buf.clone(), + scale=K**-0.5, + use_qk_l2norm_in_kernel=True, + ) + torch.testing.assert_close(out_mis.float(), out_ref.float()) + + +@skip_unsupported +def test_fi_verify_gate_env_killswitch(monkeypatch): + """The dispatch gate honors the disable env vars and shape constraints.""" + import tensorrt_llm._torch.modules.fla.fused_sigmoid_gating_recurrent as fsg + + pool = torch.zeros(4, 8, 128, 128, device="cuda", dtype=torch.bfloat16) + assert fsg._can_use_flashinfer_gdn_verify(pool, 128, 128, 4) + + monkeypatch.setenv("TRTLLM_FLA_DISABLE_FLASHINFER_GDN_VERIFY", "1") + assert not fsg._can_use_flashinfer_gdn_verify(pool, 128, 128, 4) + monkeypatch.delenv("TRTLLM_FLA_DISABLE_FLASHINFER_GDN_VERIFY") + + monkeypatch.setenv("TRTLLM_FLA_DISABLE_FLASHINFER_GDN", "1") + assert not fsg._can_use_flashinfer_gdn_verify(pool, 128, 128, 4) + monkeypatch.delenv("TRTLLM_FLA_DISABLE_FLASHINFER_GDN") + + # Shape/dtype constraints + assert not fsg._can_use_flashinfer_gdn_verify(pool.float(), 128, 128, 4) + assert not fsg._can_use_flashinfer_gdn_verify(pool, 64, 128, 4) + assert not fsg._can_use_flashinfer_gdn_verify(pool, 128, 128, 0) + assert not fsg._can_use_flashinfer_gdn_verify(pool, 128, 128, fsg._FI_GDN_MAX_MTP_T + 1) + assert not fsg._can_use_flashinfer_gdn_verify(None, 128, 128, 4)