[None][feat] Enable flashifner gdn decoding kernel for qwen3.5 (#13645)

Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>

Author: nv-guomingz

tensorrt_llm/_torch/modules/fla/chunk.py

@@ -148,7 +148,8 @@ def chunk_gated_delta_rule(
             Scale factor for the RetNet attention scores.
             If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
         initial_state (Optional[torch.Tensor]):
-            Initial state of shape `[N, H, K, V]` for `N` input sequences.
+            Initial state of shape `[N, H, V, K]` (K innermost, matching the
+            state pool) for `N` input sequences.
             For equal-length input sequences, `N` equals the batch size `B`.
             Default: `None`.
         initial_state_indices (Optional[torch.Tensor]):
@@ -159,7 +160,7 @@ def chunk_gated_delta_rule(
             `initial_state` in-place. Callers are responsible for ensuring the
             selected slots are safe to update without aliasing races.
         output_final_state (Optional[bool]):
-            Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.
+            Whether to output the final state of shape `[N, H, V, K]`. Default: `False`.
         cu_seqlens (torch.LongTensor):
             Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
             consistent with the FlashAttention API.
@@ -171,7 +172,7 @@ def chunk_gated_delta_rule(
         o (torch.Tensor):
             Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
         final_state (torch.Tensor):
-            Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.
+            Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`.
 
     Examples::
         >>> import torch

tensorrt_llm/_torch/modules/fla/chunk_delta_h.py

@@ -104,21 +104,23 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
         ht = ht + i_h * K * V
 
     # load initial state
+    # Pool layout [slots, HV, V, K], K innermost: logical block (K, V) but K has
+    # stride 1, V has stride K, order=(0, 1) so Triton treats K as innermost.
     if USE_INITIAL_STATE:
-        p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV),
-                                   (1, 0))
+        p_h0_1 = tl.make_block_ptr(h0, (K, V), (1, K), (0, i_v * BV), (64, BV),
+                                   (0, 1))
         b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32)
         if K > 64:
-            p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV),
-                                       (64, BV), (1, 0))
+            p_h0_2 = tl.make_block_ptr(h0, (K, V), (1, K), (64, i_v * BV),
+                                       (64, BV), (0, 1))
             b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32)
         if K > 128:
-            p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV),
-                                       (64, BV), (1, 0))
+            p_h0_3 = tl.make_block_ptr(h0, (K, V), (1, K), (128, i_v * BV),
+                                       (64, BV), (0, 1))
             b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32)
         if K > 192:
-            p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV),
-                                       (64, BV), (1, 0))
+            p_h0_4 = tl.make_block_ptr(h0, (K, V), (1, K), (192, i_v * BV),
+                                       (64, BV), (0, 1))
             b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32)
 
     # main recurrence
@@ -215,26 +217,26 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
             b_k = tl.load(p_k, boundary_check=(0, 1))
             b_h4 += tl.dot(b_k, b_v_new)
 
-    # epilogue
+    # epilogue — write final state back to pool-layout [slots, HV, V, K].
     if STORE_FINAL_STATE or USE_INDEXED_STATE:
-        p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV),
-                                 (1, 0))
+        p_ht = tl.make_block_ptr(ht, (K, V), (1, K), (0, i_v * BV), (64, BV),
+                                 (0, 1))
         tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
         if K > 64:
-            p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV),
-                                     (64, BV), (1, 0))
+            p_ht = tl.make_block_ptr(ht, (K, V), (1, K), (64, i_v * BV),
+                                     (64, BV), (0, 1))
             tl.store(p_ht,
                      b_h2.to(p_ht.dtype.element_ty),
                      boundary_check=(0, 1))
         if K > 128:
-            p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV),
-                                     (64, BV), (1, 0))
+            p_ht = tl.make_block_ptr(ht, (K, V), (1, K), (128, i_v * BV),
+                                     (64, BV), (0, 1))
             tl.store(p_ht,
                      b_h3.to(p_ht.dtype.element_ty),
                      boundary_check=(0, 1))
         if K > 192:
-            p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV),
-                                     (64, BV), (1, 0))
+            p_ht = tl.make_block_ptr(ht, (K, V), (1, K), (192, i_v * BV),
+                                     (64, BV), (0, 1))
             tl.store(p_ht,
                      b_h4.to(p_ht.dtype.element_ty),
                      boundary_check=(0, 1))
@@ -277,7 +279,9 @@ def chunk_gated_delta_rule_fwd_h(
             "Indexed chunk state updates require inplace_indexed_state_update=True."
         )
     store_final_state_in_kernel = output_final_state and not use_indexed_state
-    final_state = (k.new_empty(N, H, K, V, dtype=torch.float32)
+    # Kernel writes final state in [V, K] layout (K innermost) to match the
+    # pool layout. Allocate accordingly so the tensor's shape reflects memory.
+    final_state = (k.new_empty(N, H, V, K, dtype=torch.float32)
                    if store_final_state_in_kernel else None)
 
     v_new = torch.empty_like(u) if save_new_value else None

tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py

@@ -17,14 +17,9 @@
     ``use_qk_l2norm_in_kernel`` parameter on ``flashinfer.chunk_gated_delta_rule``
     is currently a dead arg, see ``flashinfer/gdn_prefill.py:317-356``).
   * Pre-gather and post-scatter of indexed SSM state (FlashInfer requires
-    packed ``[num_seqs, H, D, D]`` fp32 initial/output state).
-  * State layout: TRT-LLM's Triton uses ``[N, H, K, V]`` ordering of the last
-    two dims; FlashInfer's prefill kernel uses ``[N, H, V, K]`` (per docstring
-    in ``flashinfer/gdn_prefill.py``: "The final state layout is ``[N, H, V, K]``";
-    confirmed by ``flashinfer/tests/gdn/test_prefill_delta_rule.py`` which
-    transposes the last two dims to compare with the reference implementation).
-    The wrapper transposes initial_state on the way in and output_state on the
-    way out so callers see TRT-LLM's ``[N, H, K, V]`` everywhere.
+    packed ``[num_seqs, H, V, K]`` fp32 initial/output state). TRT-LLM's GDN
+    state pool uses the same ``[N, H, V, K]`` logical layout, so the adapter
+    casts/gathers/scatters without transposing the last two dims.
 
 This module is only imported when ``TLLM_USE_FLASHINFER_GDN_PREFILL=1`` is set
 at process start; do not import it lazily inside hot paths.
@@ -35,8 +30,8 @@
 import torch
 
 from tensorrt_llm._torch.modules.fla.fused_state_io import (
-    gather_cast_transpose_kv_to_fp32_vk,
-    transpose_cast_scatter_fp32_vk_to_kv,
+    cast_scatter_fp32_vk_to_vk,
+    gather_cast_vk_to_fp32_vk,
 )
 from tensorrt_llm._torch.modules.fla.l2norm import l2norm_fwd
 
@@ -107,12 +102,10 @@ def chunk_gated_delta_rule(
         q3 = l2norm_fwd(q3)
         k3 = l2norm_fwd(k3)
 
-    # --- Step 4: gather initial state and convert layout -----------------
-    # TRT-LLM stores SSM state with last two dims as (K, V); FlashInfer expects
-    # (V, K). Fuse gather + cast-to-fp32 + transpose + contiguous into a single
-    # Triton kernel so we get one HBM pass and one launch instead of three
-    # (``[indices]`` gather, ``.to(fp32)`` cast, ``.transpose.contiguous()`` copy).
-    gathered_init = gather_cast_transpose_kv_to_fp32_vk(initial_state, initial_state_indices)
+    # --- Step 4: gather initial state and cast dtype ---------------------
+    # TRT-LLM's GDN kernels and FlashInfer both use [N, H, V, K] state layout.
+    # Fuse gather + cast-to-fp32 + contiguous into a single Triton kernel.
+    gathered_init = gather_cast_vk_to_fp32_vk(initial_state, initial_state_indices)
 
     # --- Step 5+6: call FlashInfer with pre-allocated output/state buffers
     # FI 0.6.10 accepts `output=` / `output_state=`; pre-allocating skips its
@@ -160,26 +153,24 @@ def chunk_gated_delta_rule(
         )
         out_state = None
 
-    # --- Step 7: convert state layout back, scatter / return -----------
-    # Fuse transpose + cast (fp32 → initial_state.dtype) + optional indexed
-    # scatter into a single Triton pass, mirroring Step 4. The non-inplace
-    # branch allocates a fresh (K, V) buffer and writes every row; the inplace
-    # branch writes only the slots named by ``initial_state_indices`` and
-    # leaves the rest of ``initial_state`` untouched.
+    # --- Step 7: cast state back, scatter / return ---------------------
+    # Fuse cast (fp32 -> initial_state.dtype) + optional indexed scatter into a
+    # single Triton pass, mirroring Step 4. The inplace branch writes only the
+    # slots named by ``initial_state_indices`` and leaves the rest untouched.
     if inplace_indexed_state_update:
-        transpose_cast_scatter_fp32_vk_to_kv(out_state, initial_state, initial_state_indices)
+        cast_scatter_fp32_vk_to_vk(out_state, initial_state, initial_state_indices)
         final_to_return: Optional[torch.Tensor] = None
     elif output_final_state:
         num_seqs_out, num_h_out, v_out, k_out = out_state.shape
         final_to_return = torch.empty(
             num_seqs_out,
             num_h_out,
-            k_out,
             v_out,
+            k_out,
             dtype=initial_state.dtype,
             device=out_state.device,
         )
-        transpose_cast_scatter_fp32_vk_to_kv(out_state, final_to_return, None)
+        cast_scatter_fp32_vk_to_vk(out_state, final_to_return, None)
     else:
         final_to_return = None
 

tensorrt_llm/_torch/modules/fla/fused_recurrent.py

@@ -74,7 +74,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
 
     b_h = tl.zeros([BK, BV], dtype=tl.float32)
     if USE_INITIAL_STATE:
-        p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :]
+        # Pool layout [N, HV, V, K] with K innermost: offset v*K + k.
+        p_h0 = h0 + i_nh * V * K + o_k[:, None] + o_v[None, :] * K
         b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
 
     for _ in range(0, T):
@@ -110,7 +111,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
         p_beta += HV * (V if IS_BETA_HEADWISE else 1)
 
     if STORE_FINAL_STATE:
-        p_ht = ht + i_nh * K * V + o_k[:, None] * V + o_v[None, :]
+        # Pool layout [N, HV, V, K].
+        p_ht = ht + i_nh * V * K + o_k[:, None] + o_v[None, :] * K
         tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
 
 
@@ -137,7 +139,8 @@ def fused_recurrent_gated_delta_rule_fwd(
 
     o = q.new_empty(NK, *v.shape)
     if output_final_state:
-        final_state = q.new_empty(N, HV, K, V, dtype=torch.float32)
+        # Pool layout: [N, HV, V, K] (K innermost).
+        final_state = q.new_empty(N, HV, V, K, dtype=torch.float32)
     else:
         final_state = None
 
@@ -240,19 +243,19 @@ def fused_recurrent_gated_delta_rule(
             Scale factor for the RetNet attention scores.
             If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
         initial_state (Optional[torch.Tensor]):
-            Initial state of shape `[N, HV, K, V]` for `N` input sequences.
+            Initial state of shape `[N, HV, V, K]` for `N` input sequences.
             For equal-length input sequences, `N` equals the batch size `B`.
             Default: `None`.
         output_final_state (Optional[bool]):
-            Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
+            Whether to output the final state of shape `[N, HV, V, K]`. Default: `False`.
         cu_seqlens (torch.LongTensor):
             Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
             consistent with the FlashAttention API.
     Returns:
         o (torch.Tensor):
             Outputs of shape `[B, T, HV, V]`.
         final_state (torch.Tensor):
-            Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
+            Final state of shape `[N, HV, V, K]` if `output_final_state=True` else `None`.
     Examples::
         >>> import torch
         >>> import torch.nn.functional as F
@@ -265,7 +268,7 @@ def fused_recurrent_gated_delta_rule(
         >>> v = torch.randn(B, T, HV, V, device='cuda')
         >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
         >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
-        >>> h0 = torch.randn(B, HV, K, V, device='cuda')
+        >>> h0 = torch.randn(B, HV, V, K, device='cuda')
         >>> o, ht = fused_gated_recurrent_delta_rule(
             q, k, v, g, beta,
             initial_state=h0,
@@ -387,8 +390,9 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
         idx = tl.load(h0_indices + i_n)
         # Add bounds checking for idx
         if idx >= 0:  # Assuming negative indices are invalid
-            p_h0 = (h0_source + idx * HV * K * V + i_hv * K * V +
-                    o_k[:, None] * V + o_v[None, :])
+            # Pool layout [slots, HV, V, K], K innermost (stride 1).
+            p_h0 = (h0_source + idx * HV * V * K + i_hv * V * K + o_k[:, None] +
+                    o_v[None, :] * K)
             b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
 
     # Prepare intermediate state cache variables if enabled
@@ -427,12 +431,12 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
         # store intermediate states if enabled
         if CACHE_INTERMEDIATE_STATES:
             if cache_idx >= 0:
-                # Compute cache pointer for this step
-                step_offset = step_idx * HV * K * V
+                # Match the SSM pool layout [slots, HV, V, K] with K innermost.
+                step_offset = step_idx * HV * V * K
                 cache_ptr = (intermediate_states_buffer +
-                             cache_idx * cache_steps * HV * K * V +
-                             step_offset + i_hv * K * V + o_k[:, None] * V +
-                             o_v[None, :])
+                             cache_idx * cache_steps * HV * V * K +
+                             step_offset + i_hv * V * K + o_k[:, None] +
+                             o_v[None, :] * K)
                 tl.store(cache_ptr,
                          b_h.to(cache_ptr.dtype.element_ty),
                          mask=mask_h)
@@ -447,12 +451,12 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
         p_beta += HV * (V if IS_BETA_HEADWISE else 1)
 
     # Store final state back to h0_source with bounds checking
-    # ssm states
+    # ssm states (pool layout [slots, HV, V, K], K innermost).
     if not DISABLE_STATE_UPDATE:
         idx = tl.load(h0_indices + i_n)
         if idx >= 0:  # Add bounds checking
-            p_h0 = (h0_source + idx * HV * K * V + i_hv * K * V +
-                    o_k[:, None] * V + o_v[None, :])
+            p_h0 = (h0_source + idx * HV * V * K + i_hv * V * K + o_k[:, None] +
+                    o_v[None, :] * K)
             tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h)