fix: remove over-strict K%4 assert in get_shuffle_matrix_sf_a_row_indices (#3163)
## Summary Removes the `assert K % 4 == 0` guard from `get_shuffle_matrix_sf_a_row_indices` in `flashinfer/utils.py`. The assertion rejects valid inputs (e.g. `K=90` for GPT-OSS-120B MXFP4 weights), even though the downstream kernel handles them correctly. **Why it's safe to remove:** - `get_shuffle_matrix_sf_a_row_indices` only computes row permutation indices over the **M dimension**. K is unpacked from the shape but never used. - K is padded up to a multiple of 4 in the downstream path: - **C++ launcher** (`csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp:197`): `cols_padded = PadUpFn(cols, 4)` and passes `n_padded` to the kernel. - **CUDA kernel** (`csrc/nv_internal/cpp/kernels/quantization.cu:400-418`): loops up to `numColsPadded`, writes `0` for padded slots (`T sf = 0; if (colIdx < numCols) sf = SFIn[...];`). - **Python** (`flashinfer/quantization/fp4_quantization.py:51-54, 297-304`): `_compute_swizzled_layout_sf_size` uses `round_up(total_column, 4)` to size the output buffer to match. Fixes #2122. ## Test plan - [ ] Verify `get_shuffle_matrix_sf_a_row_indices` succeeds for `K=90` (the GPT-OSS-120B case that triggered the original assertion). - [ ] Verify `K=1`, `K=4`, `K=88` (previously valid) still return correct row indices. - [ ] Verify the full `shuffle_matrix_sf_a` pipeline (row shuffle → `block_scale_interleave`) produces correctly-sized outputs for non-aligned K. - [ ] Run the GPT-OSS-120B MXFP4 MoE repro from #2122 and confirm it no longer asserts. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * Reduced strictness around matrix dimension requirements; the system now tolerates sizes previously rejected and relies on internal padding where needed, improving compatibility and reducing manual configuration. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Alex Yang <aleyang@nvidia.com>
J
Jimmy Zhou committed
6a09cc0a49b7ddd0e0b6863b83172214c8dce5ab
Parent: 9f4d161
Committed by GitHub <noreply@github.com>
on 5/12/2026, 4:27:12 AM