Add torch.compile-compatible custom op for fp4_quantize (#3081)
## Summary
- Registers `fp4_quantize` as a `torch.library.custom_op`
(`flashinfer::fp4_quantize`) so `torch.compile` / dynamo treats it as
opaque and never traces into the JIT/subprocess internals
- Adds a `register_fake` meta kernel with correct shape inference for
all scale-factor layout combinations (swizzled 128x4, swizzled 8x4, and
linear)
- Callers inside `torch.compile` regions use
`torch.ops.flashinfer.fp4_quantize(...)` instead of `fp4_quantize()`
directly
Closes #2999
## Test plan
- [ ] Verify `torch.ops.flashinfer.fp4_quantize(x, scale)` produces
identical outputs to `fp4_quantize(x, scale)` in eager mode
- [ ] Verify `torch.compile(fullgraph=True)` succeeds when calling
`torch.ops.flashinfer.fp4_quantize`
- [x] Test with `sf_vec_size ∈ {16, 32}`, `is_sf_swizzled_layout ∈
{True, False}`, `is_sf_8x4_layout ∈ {True, False}`
- [x] Tested on DGX Spark (SM121) with SGLang 0.5.10
`torch.compile(fullgraph=True)` — shapes and values match eager
execution
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Introduced FP4 quantization support as a new operator exposed to
PyTorch.
* Added a meta/fake implementation for shape inference and allocation so
packed FP4 outputs and their scale-factor tensors are created with
correct shapes and layout handling (row-/column-major and swizzled
layouts) for downstream workflows.
<!-- review_stack_entry_start -->
[](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3081?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)
<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Co-authored-by: Brian K. Ryu <bryu@nvidia.com> S
Serge Panev committed
77477e275c74a860c63931aca83fa5ad36d5d719
Parent: 41e5aa2
Committed by GitHub <noreply@github.com>
on 5/22/2026, 4:22:07 PM