[TRTLLM-14345][feat] Improve the GDN Replay Kernel Under Low Latency

tensorrt_llm/_torch/modules/fla/cached_replay.py

@@ -13,6 +13,9 @@
 from tensorrt_llm._utils import get_sm_version
 
 _SMALL_GRID_HEAD_TILES = 512
+_BV16_MAX_HEAD_TILES = 64
+_BV32_MAX_HEAD_TILES = 128
+_RATIO2_FINE_MAPPING_MAX_HEAD_TILES = 256
 _EIGHT_WARP_COMMIT_HEAD_TILES = 1024
 _PIPELINED_COMMIT_HEAD_TILES = 2048
 _TWO_STAGE_REPLAY_HEAD_TILES = 4096
@@ -21,6 +24,33 @@
 CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE = 16
 
 
+def _default_cached_replay_block_v(
+    use_tuned_bf16_mapping: bool,
+    head_tiles: int,
+    value_dim: int,
+) -> int:
+    """Select the replay value tile without changing its cache layout."""
+    if use_tuned_bf16_mapping:
+        if head_tiles <= _BV16_MAX_HEAD_TILES:
+            return 16
+        if head_tiles <= _BV32_MAX_HEAD_TILES:
+            return 32
+        if head_tiles <= _SMALL_GRID_HEAD_TILES:
+            return 64
+    return triton.next_power_of_2(value_dim)
+
+
+def _supports_fine_grained_replay_tiling(
+    num_key_heads: int,
+    num_value_heads: int,
+    head_tiles: int,
+) -> bool:
+    """Return whether the measured fine-grained value tiling applies."""
+    return num_value_heads == 4 * num_key_heads or (
+        num_value_heads == 2 * num_key_heads and head_tiles <= _RATIO2_FINE_MAPPING_MAX_HEAD_TILES
+    )
+
+
 @triton.jit
 def _gdc_wait_with_memory_clobber():
     tl.inline_asm_elementwise(
@@ -51,6 +81,8 @@ def _cached_replay_kernel(
     pnat,
     scale,
     pool_stride_slot,
+    g_stride_row,
+    beta_stride_row,
     A_log,
     dt_bias,
     T: tl.constexpr,
@@ -192,12 +224,12 @@ def _cached_replay_kernel(
                 other=0,
             )
         b_g = tl.load(
-            g + (i_n * T + o_t) * HV + i_hv,
+            g + (i_n * T + o_t) * g_stride_row + i_hv,
             mask=mask_t,
             other=0.0,
         ).to(tl.float32)
         b_beta = tl.load(
-            beta + (i_n * T + o_t) * HV + i_hv,
+            beta + (i_n * T + o_t) * beta_stride_row + i_hv,
             mask=mask_t,
             other=0.0,
         ).to(tl.float32)
@@ -604,6 +636,8 @@ def commit_gdn_cached_replay_history_layers(
         "q",
         "k",
         "v",
+        "g",
+        "beta",
         "packed_qkv",
         "ssm_states",
         "old_u",
@@ -655,6 +689,12 @@ def fused_recurrent_gated_delta_rule_cached_replay_update(
     HV, V = v.shape[2], v.shape[3]
     assert q.shape == k.shape
     assert v.shape[:2] == (N, T)
+    assert g.shape == (N, T, HV)
+    assert beta.shape == (N, T, HV)
+    assert g.stride(2) == 1, "g head dimension must be contiguous"
+    assert beta.stride(2) == 1, "beta head dimension must be contiguous"
+    assert g.stride(0) == T * g.stride(1), "g token rows must have a uniform stride"
+    assert beta.stride(0) == T * beta.stride(1), "beta token rows must have a uniform stride"
     use_packed_qkv = packed_qkv is not None
     if use_packed_qkv:
         qkv_width = 2 * H * K + HV * V
@@ -669,25 +709,26 @@ def fused_recurrent_gated_delta_rule_cached_replay_update(
         v = v.contiguous()
         packed_qkv = q
     BK = triton.next_power_of_2(K)
-    # GB200 dispatch for the production Qwen3.5 MTP per-CTA shape. Balanced
-    # DEP and TEP runs at the same global batch have the same N * HV head-tile
-    # count, so use that workload measure instead of topology-specific H/HV
-    # values or the per-rank batch alone.
-    use_tuned_bf16_mapping = (
+    # The complete workload mapping is retained for the 4:1 shape; the
+    # 2:1 fine tiling is enabled only through the measured low-batch range.
+    use_production_bf16_shape = (
         T == 4
         and history_size <= 16
-        and HV == 4 * H
         and K == 128
         and V == 128
         and ssm_states.dtype == torch.bfloat16
     )
     head_tiles = N * HV
-    use_small_grid_mapping = use_tuned_bf16_mapping and head_tiles <= _SMALL_GRID_HEAD_TILES
+    use_tuned_bf16_mapping = use_production_bf16_shape and HV == 4 * H
+    use_fine_grained_mapping = use_production_bf16_shape and _supports_fine_grained_replay_tiling(
+        H, HV, head_tiles
+    )
+    use_small_grid_mapping = use_fine_grained_mapping and head_tiles <= _SMALL_GRID_HEAD_TILES
     if block_v is None:
-        block_v = 64 if use_small_grid_mapping else triton.next_power_of_2(V)
+        block_v = _default_cached_replay_block_v(use_fine_grained_mapping, head_tiles, V)
     if num_warps is None:
         num_warps = (
-            2 if use_tuned_bf16_mapping and (use_small_grid_mapping or launch_with_pdl) else 4
+            2 if use_fine_grained_mapping and (use_small_grid_mapping or launch_with_pdl) else 4
         )
     BV = block_v
     use_large_workload_mapping = (
@@ -773,6 +814,8 @@ def launch(
             pnat=prev_num_accepted_tokens,
             scale=scale,
             pool_stride_slot=s_h0_0,
+            g_stride_row=g.stride(1),
+            beta_stride_row=beta.stride(1),
             A_log=A_log,
             dt_bias=dt_bias,
             T=T,

tensorrt_llm/_torch/modules/mamba/gdn_mixer.py

@@ -181,12 +181,12 @@ def __init__(
         replay_enabled = is_gdn_replay_enabled()
         if replay_enabled:
             logger.info_once(
-                "Configured GDN MTP replay implementation: cached",
+                "GDN MTP replay is enabled; set TRTLLM_USE_GDN_REPLAY=0 to disable it",
                 key="gdn_mtp_replay_cached",
             )
         else:
             logger.info_once(
-                "GDN MTP replay disabled; set TRTLLM_USE_GDN_REPLAY=1 to enable it",
+                "GDN MTP replay is disabled; set TRTLLM_USE_GDN_REPLAY=1 to enable it",
                 key="gdn_mtp_replay_disabled",
             )
 

tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py

@@ -289,9 +289,20 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int,
             return
         num_decodes = batch_size - num_contexts
         self.replay_num_decodes = num_decodes
-        self.replay_n_writes.zero_()
         if num_decodes == 0:
             return
+        if getattr(kv_cache_manager, 'use_gdn_cached_replay_all_layer_commit',
+                   False):
+            from tensorrt_llm._torch.modules.fla.cached_replay import \
+                CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE
+
+            # The fused small-batch GDN kernel commits its checkpoint in-layer
+            # and indexes cache metadata directly. Work items are only consumed
+            # by the partitioned replay + all-layer commit path.
+            if num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE:
+                return
+
+        self.replay_n_writes.zero_()
         if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'):
             raise RuntimeError(
                 "Replay state update is enabled, but the KV cache manager "

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2021,10 +2021,8 @@ def _create_kv_cache_manager(
                 "MTP path")
             use_replay = False
 
-        # Replay is opt-in because its end-to-end benefit is workload-dependent.
+        # Replay is enabled by default for eligible GDN MTP workloads.
         if not is_gdn_replay_enabled():
-            logger.info("GDN replay kernel is disabled; set "
-                        "TRTLLM_USE_GDN_REPLAY=1 to enable it")
             use_replay = False
         logger.info("GDN replay state update: " +
                     ("ENABLED" if use_replay else "DISABLED"))