DFlash on DGX — Every Token Rate Scaled by Six?

based on baristankut’s Spark image, inspired by shahizat.

Still subject to further testing — does it respond with DFlash the same way as without? But the token rate is outstanding.

Setup Overview

Component Details
Base Image lmsysorg/sglang:spark (arm64, SM121, PyTorch 2.9.0, sgl-kernel 0.3.16)
SGLang Version 0.5.6 from PR #16818 (DFlash Speculative Decoding)
flashinfer 0.6.3 (pure Python, JIT via TVM/Triton)
Target Model Qwen3-Coder-30B-A3B-Instruct BF16 (57 GB, 16 Shards)
Draft Model z-lab/Qwen3-Coder-30B-A3B-DFlash (~900 MB, 0.5B Diffusion Drafter)
Performance 50.1 tok/s on DGX Spark
SM121 Workarounds DeepGEMM disabled, flashinfer instead of FA3, BF16 instead of FP8

Dockerfile

FROM docker.io/lmsysorg/sglang:spark

# DFlash on top of lmsysorg/sglang:spark (arm64, GB10/SM121)
# sgl-kernel + Triton already compiled for DGX Spark

# Remove old flashinfer-cubin (0.5.0, conflicts with new flashinfer)
# Upgrade flashinfer 0.6.3 (pure Python wheel, kernels via TVM/Triton JIT)
RUN pip uninstall -y --break-system-packages flashinfer-cubin && \
    pip install --no-cache-dir --break-system-packages \
      "flashinfer-python>=0.5.3"

# Install DFlash SGLang PR (overwrites stock SGLang)
RUN pip install --no-cache-dir --no-deps --break-system-packages \
      "git+https://github.com/sgl-project/sglang.git@refs/pull/16818/head#subdirectory=python"

# Patch: fused_qk_norm_rope missing in sgl-kernel 0.3.16 (spark)
# SGLang 0.5.6 imports it top-level → entire Qwen3MoE module fails
# Fix: try/except so the unfused fallback path is used
RUN python3 -c "\
p = '/usr/local/lib/python3.12/dist-packages/sglang/srt/models/qwen3_moe.py'; \
t = open(p).read(); \
t = t.replace( \
    '    from sgl_kernel import fused_qk_norm_rope', \
    '    try:\\n        from sgl_kernel import fused_qk_norm_rope\\n    except ImportError:\\n        fused_qk_norm_rope = None'); \
open(p,'w').write(t)"

WORKDIR /

Launch Script

#!/bin/bash
set -euo pipefail

TARGET_MODEL="/data/tensordata/Qwen3-Coder-30B-A3B-Instruct"
DRAFT_MODEL="/data/tensordata/Qwen3-Coder-30B-A3B-DFlash"
PORT=8011
CONTAINER_NAME="dflash"

podman run -d \
  --replace \
  --name "$CONTAINER_NAME" \
  --device nvidia.com/gpu=all \
  --security-opt=label=disable \
  --hooks-dir=/usr/share/containers/oci/hooks.d \
  --ipc=host \
  --network host \
  -v /data/tensordata:/data/tensordata \
  -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
  -e SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 \
  -e SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 \
  -e SGLANG_ENABLE_JIT_DEEPGEMM=false \
  localhost/dflash \
  python3 -m sglang.launch_server \
    --model-path "$TARGET_MODEL" \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path "$DRAFT_MODEL" \
    --host 0.0.0.0 \
    --port "$PORT" \
    --served-model-name qwen3-coder-30b-dflash \
    --tp-size 1 \
    --dtype bfloat16 \
    --attention-backend flashinfer \
    --mem-fraction-static 0.75 \
    --trust-remote-code

The key insight: DFlash speculative decoding with a tiny 0.5B diffusion drafter “gives” roughly a 6x speedup on the Spark. All the SM121 pain points (no FA3, no FP8 GEMM, no DeepGEMM) are worked around — flashinfer JIT + BF16 does the job.


BF16 Memory-Bound Throughput Calculation

Parameter Value
Active parameters per token 3B (MoE, A3B)
Bytes per parameter (BF16) 2
Weights per token 3B × 2 = 6 GB
Memory bandwidth GB10 273 GB/s

Theoretical maximum: 273 GB/s ÷ 6 GB/token = 45.5 tok/s

This lines up well with what we’re seeing:

Mode tok/s vs. Theory
Theory (100% BW) 45.5 100%
DFlash (measured) 50.1 110%
Vanilla (estimated) ~35–40 ~80–85%

DFlash exceeds the single-token limit because speculative decoding validates multiple tokens per target forward pass — the drafter proposes e.g. 4–8 tokens, the target verifies them in a single batch. This means >1 token is generated per weight load cycle.

With vanilla autoregressive inference (no speculation), you’d realistically expect ~38–40 tok/s (85–90% bandwidth utilization, plus attention/embedding overhead).

So DFlash delivers roughly 25–30% speedup beyond the pure memory bandwidth limit.


Lets find the missing 470%

2 Likes