Skip to content

Add opt-in MI350 pack_segments fused kernel - #6263

Open
Ruishenl wants to merge 1 commit into
pytorch:mainfrom
Ruishenl:export-D118393046
Open

Add opt-in MI350 pack_segments fused kernel#6263
Ruishenl wants to merge 1 commit into
pytorch:mainfrom
Ruishenl:export-D118393046

Conversation

@Ruishenl

@Ruishenl Ruishenl commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

Summary:
Add an opt-in fused prefix-and-pack kernel for gfx950 inputs with at most 128 segments. FBGEMM_ENABLE_PACK_SEGMENTS_SMALL_N=1 removes the separate exclusive-cumsum launch while preserving the legacy CUDA_KERNEL_ASSERT input invariant. The path remains disabled by default pending broader workload validation.

{F1995545397}

Differential Revision: D118393046

Summary:
Add an opt-in fused prefix-and-pack kernel for gfx950 inputs with at most 128 segments. `FBGEMM_ENABLE_PACK_SEGMENTS_SMALL_N=1` removes the separate exclusive-cumsum launch while preserving the legacy `CUDA_KERNEL_ASSERT` input invariant. The path remains disabled by default pending broader workload validation.

 {F1995545397}

Differential Revision: D118393046
@meta-codesync meta-codesync Bot changed the title Fuse MI350 pack_segments scan into packing kernel Add opt-in MI350 pack_segments fused kernel Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant