SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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