Add segment_ids support to pallas flash attention on GPU.#193

copybara-service[bot] wants to merge 1 commit intomainjax-ml/jax-triton:mainfrom

test_547592877jax-ml/jax-triton:test_547592877Copy head branch name to clipboard