diff --git a/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py b/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py index f9f21a911a95..e819f1690728 100644 --- a/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py +++ b/tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py @@ -16,10 +16,13 @@ (the FlashInfer prefill kernel does NOT apply L2 norm internally; the ``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, 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. + * Pre-gather and post-scatter of indexed SSM state into FlashInfer's packed + ``[num_seqs, H, V, K]`` layout. TRT-LLM's GDN state pool uses the same + ``[N, H, V, K]`` logical layout, so the adapter gathers/scatters without + transposing the last two dims. The SM100/SM103 kernel carries the recurrent + state in fp32 in TMEM regardless of the initial/output-state I/O dtype, so + the round-trip stays in the native pool dtype (bf16/fp16) there with no + precision change; only SM90/SM120 need an fp32 up-cast/down-cast. 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. @@ -34,6 +37,7 @@ gather_cast_vk_to_fp32_vk, ) from tensorrt_llm._torch.modules.fla.l2norm import l2norm_fwd +from tensorrt_llm._utils import is_sm_100f # Mirror the @torch.compiler.disable on the legacy Triton wrapper @@ -103,10 +107,18 @@ def chunk_gated_delta_rule( q3 = l2norm_fwd(q3) k3 = l2norm_fwd(k3) - # --- Step 4: gather initial state and cast dtype --------------------- + # --- Step 4: gather initial state (+ cast dtype only when required) --- # 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) + # The SM100/SM103 kernel carries the recurrent state in fp32 in TMEM + # regardless of the initial/output-state I/O dtype (the state tensors are + # only the gmem load/store format), so passing bf16/fp16 state there is + # numerically identical to the fp32 round-trip while moving half the bytes. + # SM90/SM120 still require fp32 state. Fuse gather (+ optional cast) and + # contiguous into a single Triton kernel. + state_dtype = initial_state.dtype if is_sm_100f() else torch.float32 + gathered_init = gather_cast_vk_to_fp32_vk( + initial_state, initial_state_indices, out_dtype=state_dtype + ) # --- Step 5+6: call FlashInfer with pre-allocated output/state buffers # FI 0.6.10 accepts `output=` / `output_state=`; pre-allocating skips its @@ -126,7 +138,10 @@ def chunk_gated_delta_rule( ) if need_state: num_seqs = cu_seqlens.shape[0] - 1 - state_buf = q3.new_empty(num_seqs, num_o_heads, head_size, head_size, dtype=torch.float32) + # Match the initial-state dtype (native bf16/fp16 on SM100/SM103, else + # fp32); FlashInfer writes the final state in this dtype and the scatter + # below adapts to the destination pool dtype without an extra cast. + state_buf = q3.new_empty(num_seqs, num_o_heads, head_size, head_size, dtype=state_dtype) out_packed, out_state = flashinfer.chunk_gated_delta_rule( q=q3, k=k3, @@ -158,8 +173,9 @@ def chunk_gated_delta_rule( ) out_state = None - # --- Step 7: cast state back, scatter / return --------------------- - # Fuse cast (fp32 -> initial_state.dtype) + optional indexed scatter into a + # --- Step 7: cast state back (if needed), scatter / return --------- + # Fuse cast (out_state.dtype -> destination dtype; a no-op on SM100/SM103 + # where both are the native pool 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: diff --git a/tensorrt_llm/_torch/modules/fla/fused_state_io.py b/tensorrt_llm/_torch/modules/fla/fused_state_io.py index c5885f5b6011..ee420993fec3 100644 --- a/tensorrt_llm/_torch/modules/fla/fused_state_io.py +++ b/tensorrt_llm/_torch/modules/fla/fused_state_io.py @@ -79,8 +79,16 @@ def _gather_cast_vk_to_fp32_vk_kernel( def gather_cast_vk_to_fp32_vk( initial_state: torch.Tensor, initial_state_indices: Optional[torch.Tensor], + out_dtype: Optional[torch.dtype] = None, ) -> torch.Tensor: - """Fused ``initial_state[indices].to(fp32).contiguous()`` for ``[N, H, V, K]`` state.""" + """Fused ``initial_state[indices].to(out_dtype).contiguous()`` for ``[N, H, V, K]`` state. + + ``out_dtype`` defaults to ``torch.float32``, the dtype the SM90/SM120 + FlashInfer GDN prefill kernels require. On the SM100/SM103 kernel, which + reads native bf16/fp16 state and casts to fp32 internally, pass + ``initial_state.dtype`` to gather without an up-cast (paired with a matching + scatter that skips the down-cast). + """ assert initial_state.dim() == 4, f"initial_state must be 4D, got {initial_state.shape}" n_pool, h, v, k = initial_state.shape if initial_state_indices is not None: @@ -93,7 +101,9 @@ def gather_cast_vk_to_fp32_vk( # K and V are typically 128 in GDN; one (BLOCK_K, BLOCK_V) tile covers the full K and V dimensions. # entire (K, V) plane per (seq, head). Larger tiles save grid overhead; # smaller tiles improve occupancy at small num_seqs * H. - output = torch.empty(num_seqs, h, v, k, dtype=torch.float32, device=initial_state.device) + if out_dtype is None: + out_dtype = torch.float32 + output = torch.empty(num_seqs, h, v, k, dtype=out_dtype, device=initial_state.device) block_v = min(v, 128) block_k = min(k, 128) num_v_blocks = triton.cdiv(v, block_v)