vllm-omni/docs/design/diffusion_ops.md

<!-- markdownlint-disable MD013 -->

# RFC: Establish an operator boundary for vLLM-Omni diffusion

> **Status:** Draft for discussion
>
> **vLLM-Omni snapshot:** [`92534733`](https://github.com/vllm-project/vllm-omni/commit/925347332f8190d172dfa8ac0c307a5cfa2320cb)
>
> **Upstream vLLM snapshot:** [`fa2a2589`](https://github.com/vllm-project/vllm/commit/fa2a2589bd2a1fce0851df7fd42ffb54b6195f04)

## Summary

vLLM has an operator boundary, most visibly through `_custom_ops.py` and its
kernel-provider interfaces. Model code can use optimized kernels without
depending on native extension names, packed layouts, allocation details, or
platform-specific dispatch.

vLLM-Omni already reuses much of that machinery. However, diffusion support has
also accumulated local Triton kernels, `torch.library` custom ops, platform
calls, third-party providers, and compatibility patches. These cover semantics
that are sometimes different from autoregressive inference:

- N-dimensional and non-contiguous activations;
- routed mixture-of-transformers (MoT) computation;
- diffusion-specific attention and quantized attention;
- platform-specific microscaling formats; and
- stateful diffusion cache operations.

This RFC asks whether vLLM-Omni should define an explicit diffusion operator
boundary. “Independent” means independently defined semantics, not independent
reimplementation: generic kernels should continue to come from upstream vLLM.

Approval would only establish ownership principles. It would not approve a
refactor, public API, file layout, or migration plan.

## Motivation

### Diffusion needs a stronger tensor contract

Diffusion linears commonly receive tensors shaped `(..., K)` and must return
`(..., N)`, preserving all leading dimensions. Inputs may also be
non-contiguous.

The current code adapts upstream FP8 paths by:

- making ScaledMM activations contiguous before a two-dimensional `view`; and
- restoring leading dimensions when the FlashInfer provider returns a
  two-dimensional result.

Several quantization methods independently perform similar flatten/restore
logic. These are provider-contract requirements, not new GEMM algorithms.

### Some operations are genuinely diffusion-specific

Examples include:

- MoT kernels that route text and VAE tokens to different weights inside one
  fused GEMM;
- diffusion attention with its own layouts, masks, parallel strategies, and
  optional quantization;
- NPU FP8 attention that rotates and block-quantizes Q/K/V on each denoising
  step; and
- experimental paged attention that writes into a process-owned diffusion KV
  pool.

These semantics do not map cleanly to generic dense-linear or autoregressive
attention interfaces.

### Quantization mixes policy, checkpoint handling, and kernels

Current quantization code combines:

1. deciding which stages and layers to quantize;
2. loading, remapping, and packing weights and scales; and
3. selecting and invoking a hardware kernel.

The first two are often Omni-specific. The third is frequently generic and
already supported upstream. Without a clear boundary, it is difficult to tell
when Omni is extending vLLM and when it is duplicating it.

## Current status

The audited snapshot has no project-owned C++, CUDA, or HIP extension sources.
Its optimized operator surface is Python- and provider-driven:

| Mechanism | Current status |
| --- | --- |
| Upstream vLLM ops/providers | Default path for generic quantization, GEMM, normalization, and related primitives |
| Diffusion `CustomOp` | Construction-time platform dispatch for norm, RoPE, AdaLayerNorm, and MoT RMSNorm |
| Attention registry | Nine default diffusion attention backends with platform overrides |
| Local Triton | Seven decorated functions, primarily MoT plus one model-specific fusion |
| `torch.library` | Three custom ops: SageAttention 3, ROCm MXFP4, and experimental paged write-attention |
| External providers | Direct use of `torch_npu`, AITER, FlashInfer, SageAttention, BitsAndBytes, and Quack |
| Compatibility patches | FP8 contiguity/output shape, NVFP4 scale sanitization, Quack selection, and scoped provider forcing |

### Quantization snapshot

| Path | Kernel source | Omni-specific behavior |
| --- | --- | --- |
| FP8 / ModelOpt / NVFP4 | Upstream vLLM | Checkpoint adaptation, diffusion shape fixes, HSDP weight views, scale sanitization |
| CUDA INT8 | Upstream vLLM | Configuration and online/offline weight handling |
| NPU INT8 and MXFP8 | `torch_npu` | Layouts, scales, bias/dtype handling, and shape restoration |
| XPU MXFP8 | Upstream vLLM | Local N-dimensional wrapper |
| NPU MXFP4 | `torch_npu` | Single-scale and diffusion-oriented DualScale formats |
| ROCm MXFP4 | Local AITER custom op | Overlaps with an AITER provider now present upstream |
| BitsAndBytes W4 | BitsAndBytes | Load-time quantization and shape adaptation |
| Quack FP8 | Quack CuteDSL | Installed by replacing the selected FlashInfer method |
| MoT FP8/INT8 | Local Triton plus upstream FP8 activation quantization | Routed text/VAE semantics |

Upstream now has provider registries for FP8, INT8, NVFP4, MXFP8, and MXFP4
across supported platforms, plus `register_linear_kernel()`. This should remain
the preferred home for model-agnostic linear kernels.

### Problems exposed by the current structure

- Tensor shape, contiguity, dtype, bias, packing, and scale contracts are not
  expressed uniformly across providers.
- Global patches solve real problems but make provider ownership, ordering, and
  version compatibility difficult to inspect.
- FakeTensor, `torch.compile`, fallback, mutation, and aliasing behavior varies
  by integration.
- Platform capability information is duplicated across code and documentation
  and has already drifted.
- Tests cover configuration and weight layouts well, but quantized MoT, Quack
  execution, real ROCm MXFP4, and cross-provider compile behavior remain thin.

## Direction proposed for discussion

```text
diffusion model or layer
        |
        v
diffusion semantic contract, only where upstream is insufficient
        |
        +----> upstream vLLM provider
        +----> Omni-owned diffusion provider
        +----> platform or out-of-tree provider
```

Suggested principles:

1. **Reuse upstream by default.** Do not wrap or fork a generic kernel whose
   upstream contract already works for diffusion.
2. **Fix generic contracts upstream.** Shape preservation, non-contiguous input
   handling, generic packing, and provider registration should live in vLLM
   when they benefit more than diffusion.
3. **Own diffusion semantics in Omni.** MoT routing, quantized diffusion
   attention, and diffusion cache behavior are reasonable Omni responsibilities.
4. **Separate policy from execution.** Stage selection, ignored layers,
   checkpoint adaptation, and quality-sensitive BF16 fallback are not kernel
   responsibilities.
5. **Treat hardware integrations as providers.** Model code should not depend
   directly on vendor ABI, packed layouts, or optional package internals.
6. **Make compilation and state explicit.** Opaque kernels need deliberate fake
   behavior, mutation/aliasing semantics, and fallback rules.
7. **Prefer scoped selection over global replacement.** Active providers and
   fallbacks should be inspectable and limited to the intended operation or
   stage.
8. **Do not require one monolithic file.** The boundary could be a small package
   or extensions to existing registries; copying upstream `_custom_ops.py` is
   not the goal.

Likely upstream-owned:

- generic dense-linear quantization kernels and activation quantization;
- provider-independent shape/dtype contracts; and
- generic provider registration and selection.

Likely Omni-owned:

- routed MoT operators;
- dynamically quantized diffusion attention;
- diffusion-specific cache operations; and
- formats or semantics not represented upstream, such as the current NPU
  MXFP4 DualScale path.

## Non-goals

- Defining a detailed refactor or migration sequence.
- Choosing final module names or public Python signatures.
- Reimplementing vLLM's generic custom ops or linear kernels.
- Adding an Omni native extension loader without a demonstrated need.
- Changing supported hardware or claiming unvalidated quantization support.
- Removing compatibility patches before replacement behavior is validated.

## Open questions

1. Is an explicit diffusion operator boundary needed, or are the current layer
   and backend interfaces sufficient?
2. Should the boundary remain diffusion-only or eventually cover other
   non-autoregressive Omni stages?
3. What minimum shape, layout, dtype, bias, and compilation contract should
   every provider satisfy?
4. Which current adaptations are generic upstream fixes versus
   diffusion-specific behavior?
5. Should quantized diffusion attention be an attention backend, an operator
   provider, or both?
6. Should platform providers live in Omni core, platform projects, or
   out-of-tree packages?
7. Is upstream `register_linear_kernel()` sufficient for provider priority,
   diagnostics, explicit selection, and fallback?
8. How should stateful operators declare process-global mutation to eager mode
   and `torch.compile`?
9. What validation should be required before an experimental provider becomes
   a supported default?

## Feedback requested

Feedback is requested from upstream kernel, diffusion, quantization, platform,
attention, and `torch.compile` maintainers. Particularly useful comments would
identify:

- an operation that clearly belongs upstream or clearly belongs in Omni;
- a missing provider contract;
- a compatibility patch that should remain local or move upstream; or
- a platform integration that needs an out-of-tree extension point.

If there is agreement on the ownership principles, a follow-up design can
compare concrete API options. If not, this RFC should document why the existing
abstractions are sufficient.

## References

- [vLLM `_custom_ops.py`](https://github.com/vllm-project/vllm/blob/fa2a2589bd2a1fce0851df7fd42ffb54b6195f04/vllm/_custom_ops.py)
- [vLLM linear-kernel registry](https://github.com/vllm-project/vllm/blob/fa2a2589bd2a1fce0851df7fd42ffb54b6195f04/vllm/model_executor/kernels/linear/__init__.py)
- [vLLM-Omni quantization factory](https://github.com/vllm-project/vllm-omni/blob/925347332f8190d172dfa8ac0c307a5cfa2320cb/vllm_omni/quantization/factory.py)
- [Diffusion attention registry](https://github.com/vllm-project/vllm-omni/blob/925347332f8190d172dfa8ac0c307a5cfa2320cb/vllm_omni/diffusion/attention/backends/registry.py)
- [Diffusion `CustomOp`](https://github.com/vllm-project/vllm-omni/blob/925347332f8190d172dfa8ac0c307a5cfa2320cb/vllm_omni/diffusion/layers/custom_op.py)
- [Compatibility patches](https://github.com/vllm-project/vllm-omni/blob/925347332f8190d172dfa8ac0c307a5cfa2320cb/vllm_omni/patch.py)
- [MoT GEMM](https://github.com/vllm-project/vllm-omni/blob/925347332f8190d172dfa8ac0c307a5cfa2320cb/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py)
- [NPU FP8 diffusion attention](https://github.com/vllm-project/vllm-omni/blob/925347332f8190d172dfa8ac0c307a5cfa2320cb/vllm_omni/platforms/npu/quant/kv_quant_npu.py)