Repository navigation
Description
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
0Inconsistency with the rest of the file. This exact hazard has already been fixed twice around this function:
_flashinfer_gdn_verify(added in [None][feat] Dispatch GDN MTP target-verify to FlashInfer bf16 kernel #15975) realigns the same argument:.int()followed byif initial_state_indices.data_ptr() % 32 != 0: initial_state_indices = initial_state_indices.clone(), with regression testtest_fi_mtp_verify_misaligned_index_slice.- [TRTLLM-13349][perf] Fuse gemma RMSNorm into AllReduce for Qwen3-Next/Qwen3.5… #15194 added the analogous
% 32clone guards for thea/bactivation slices in_flashinfer_gdn_decodeitself — but leftinitial_state_indicesunguarded, 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.
Activity
[NVIDIA#17027][fix] Realign misaligned int32 index slices in FlashInf…
[NVIDIA#17027][fix] Address review: narrow probe exceptions, naming, …
[NVIDIA#17027][fix] Realign misaligned int32 index slices in FlashInf…
[NVIDIA#17027][fix] Address review: narrow probe exceptions, naming, …
[NVIDIA#17027][fix] Address review: shorten probe import, share decod…
[NVIDIA#17027][fix] Realign misaligned int32 index slices in FlashInf…
[NVIDIA#17027][fix] Address review: narrow probe exceptions, naming, …
[NVIDIA#17027][fix] Address review: shorten probe import, share decod…
[NVIDIA#17027][fix] Realign misaligned int32 index slices in FlashInf…
[NVIDIA#17027][fix] Address review: narrow probe exceptions, naming, …
[NVIDIA#17027][fix] Address review: shorten probe import, share decod…
[#17027][fix] Realign misaligned int32 index slices in FlashInfer GDN…