[JAX] Fix get_seqlens_and_offsets() to accept vmapped seg ids and non vmapped seg offsets#2692

Merged

KshitijLakhani merged 13 commits intoNVIDIA:mainfrom

KshitijLakhani:klakhani/fix/vmap-get-seg-ids-pos

Mar 11, 2026

Commits on Feb 19, 2026

Fix batcher for when segment ids received are batched/vmapped whereas the TE constructed segment pos are not thereby causing mismatches in impl()
KshitijLakhani
committed

Commits on Feb 24, 2026

Commits on Feb 25, 2026

Fix the batcher logic to check for q and kv seg ids separately
KshitijLakhani
committed

Commits on Feb 27, 2026

Commits on Mar 10, 2026

nit: Use partial instead of a single use function
KshitijLakhani
committed