SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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 -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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