[Bug]: GDN FlashInfer decode: initial_state_indices.int() is a no-op for int32 and leaves misaligned index slices (Misaligned Tensor data, alignment=32 bytes)

System Info

• Observed on TensorRT-LLM v1.3.0rc22 (PyTorch backend), NVIDIA H100, Linux x86_64.
• Bug is version-independent (pointer-offset arithmetic); the affected code is unchanged on current main (c45ad83).

Who can help?

@nv-guomingz (authored the sibling fix in #15975)

Information

• The official example scripts
• My own modified scripts

Reproduction

_flashinfer_gdn_decode in tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py forwards its index tensor to the FlashInfer CuTe-DSL kernel as:

initial_state_indices=initial_state_indices.int(),

Tensor.int() is a no-op on an already-int32 tensor: it returns the same storage and the same (possibly misaligned) data pointer. The FlashInfer CuTe-DSL GDN kernels assert 32-byte data alignment on every tensor argument (enforced in build_memref_desc), so any int32 index tensor that is a sliced view at an element offset not divisible by 8 is rejected at runtime with a hard error:

ValueError: Misaligned Tensor data on argument #10 ... expected data alignment=32 bytes

There is no fallback: _can_use_flashinfer_gdn_decode gates dispatch on dtype/arch/shape only and never checks alignment, so the request crashes instead of falling back to the Triton kernel.

Trigger conditions. The in-tree default callers currently pass zero-offset views (state_indices[:batch_size]), so standard serving does not hit this. It triggers as soon as a caller hands fused_sigmoid_gating_delta_rule_update / _flashinfer_gdn_decode an offset slice — e.g. the natural decode half of a mixed prefill+decode batch, state_indices[num_prefills:] (the exact pattern gdn_mixer.py uses for its verify path) — whenever num_prefills % 8 != 0. It is dormant when num_prefills is 0 or a multiple of 8, which makes it intermittent and batch-composition-dependent in practice. I hit it on v1.3.0rc22 with a hybrid GDN model while routing the decode half of mixed batches through this API.

CPU-side demonstration of the pointer math (no GPU needed):

>>> import torch
>>> t = torch.zeros(9, dtype=torch.int32)   # allocator-aligned base
>>> t.data_ptr() % 32
0
>>> v = t[1:]                               # slice view, 4-byte storage offset
>>> v.data_ptr() % 32
4
>>> v.int().data_ptr() == v.data_ptr()      # .int() is a no-op: same misaligned pointer
True
>>> v.is_contiguous()                       # .contiguous() would be a no-op too
True
>>> v.clone().data_ptr() % 32               # clone realigns
0

Inconsistency with the rest of the file. This exact hazard has already been fixed twice around this function:

• _flashinfer_gdn_verify (added in #15975) realigns the same argument: .int() followed by if initial_state_indices.data_ptr() % 32 != 0: initial_state_indices = initial_state_indices.clone(), with regression test test_fi_mtp_verify_misaligned_index_slice.
• #15194 added the analogous % 32 clone guards for the a/b activation slices in _flashinfer_gdn_decode itself — but left initial_state_indices unguarded, still behind the misleading .int().

Fix. Mirror the verify-path guard in _flashinfer_gdn_decode: materialize the index slice with .clone() only when misaligned (zero-copy in the common aligned case). PR incoming.

Found while profiling hybrid-model serving; broader context in #16976.

Expected behavior

_flashinfer_gdn_decode accepts any valid int32 index tensor, including sliced views, like the Triton path and _flashinfer_gdn_verify do.

actual behavior

Hard ValueError: Misaligned Tensor data on argument #10 ... expected data alignment=32 bytes from the FlashInfer CuTe-DSL kernel when the index tensor is a slice at a non-32-byte-aligned storage offset; no Triton fallback.

additional notes

The .int() call reads as if it sanitizes the tensor for the kernel, which makes the latent misalignment easy to miss in review — worth fixing even though no default in-tree path currently passes an offset slice to the decode entry point.

Start in tensorrt_llm/_torch/modules/fla/fused_sigmoid_gating_recurrent.py at _flashinfer_gdn_decode, and compare its initial_state_indices handling with _flashinfer_gdn_verify. Run test_fi_mtp_verify_misaligned_index_slice as the existing regression example. Done means valid misaligned int32 sliced indices are accepted by the decode path without the FlashInfer alignment error, while aligned inputs remain zero-copy.