Skip to content

Feature Request: Native Support for DFlash in Omlx #762

Description

@Revive-Curiosity

Is your feature request related to a problem? Please describe.
Currently, Omlx does not support [DFlash: Block Diffusion for Flash Speculative Decoding](https://arxiv.org/html/2602.06036v1) natively. This limits speculative decoding efficiency, resulting in lower acceptance lengths and throughput compared to what DFlash can achieve. For workloads requiring high performance inference, this becomes a bottleneck.

Describe the solution you'd like
Add native support for DFlash within the Omlx project. This would allow acceptance lengths to increase to ~6–8 tokens per token verified by target model and throughput improvements of up to 3.5×–6× compared to standard autoregressive decoding.

Describe alternatives you've considered

  • Using standard autoregressive decoding (baseline, but limited to acceptance length of 1).
  • Implementing speculative decoding with other draft models, though these approaches typically yield lower acceptance lengths and less throughput improvement compared to DFlash.

Additional context

  • DFlash paper: https://arxiv.org/html/2602.06036v1
  • Benchmarks show DFlash consistently achieves acceptance lengths between 6–8 tokens per block, leading to significant speedups in inference throughput.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions