Procházet zdrojové kódy

fix(recipes): switch Qwen3.8-Flash-Next to NVFP4 for 2x Spark

FP8 cannot serve on 2x DGX Spark: MoE weight_block_size=[128,128] +
moe_intermediate_size=640 violates SGLang's (intermediate/tp) % block_n == 0
at tp=2 (sgl-project/sglang#21810); only tp in {1,5} align, tp=1 can't hold
125B. NVFP4 (RadixArk) is the community-verified 2x Spark path.

Also vendor the SM121 QSA decode kernel fix as a docker-images/ overlay:
the stock lmsysorg/sglang:qwen38flashnext cannot serve qwen4_exp on GB10
(QSA varlen decode resolves to flash-attn-4 CuTe-DSL kernels that fail MLIR
on SM121). The overlay swaps in a CUDA-graph-safe Triton fallback, verified
end-to-end (patch check/apply/check + behavior probe).

- recipe renamed qwen38-flash-next-fp8 -> qwen38-flash-next-nvfp4 (NVFP4 flags,
  NEXTN spec decode from the in-checkpoint MTP, flashinfer, PLE auto-offload)
- docker-images/qwen38-flashnext-nvfp4/ (Dockerfile, build.sh, versions.env,
  patches/patch_qsa_sm121.py + qsa_fa_fallback.py)
- provenance: MiaAI-Lab/Qwen3.8-Flash-Next-Dual-DGX-Sparks (validated TP2/RoCE)
Alexander Korolev před 1 měsícem
rodič
revize
5ee2326c40

+ 32 - 0
docker-images/qwen38-flashnext-nvfp4/Dockerfile

@@ -0,0 +1,32 @@
+# syntax=docker/dockerfile:1.7
+# Qwen3.8-Flash-Next-NVFP4 on DGX Spark (SM121 / GB10).
+#
+# Base: lmsysorg/sglang:qwen38flashnext — SGLang day-0 for qwen4_exp (arm64,
+# model support = sgl-project/sglang PR #36497 baked in). It CANNOT serve this
+# model on GB10 stock: QSA's packed varlen decode resolves to flash-attn-4
+# CuTe-DSL kernels that fail MLIR compilation on SM121.
+#
+# This layer swaps in a Triton FlashDecoding-style fallback (CUDA-graph-safe)
+# for non-SM100 devices. Provenance + validation:
+#   https://github.com/MiaAI-Lab/Qwen3.8-Flash-Next-Dual-DGX-Sparks
+#   (validated end-to-end TP=2 over RoCE on this cluster; single + 2-node
+#   dummy boots with prefill, decode CUDA graphs and NEXTN spec decode).
+
+ARG BASE_IMAGE
+FROM ${BASE_IMAGE}
+
+# The Triton fallback kernel (installed next to the QSA backend).
+COPY patches/qsa_fa_fallback.py \
+    /sgl-workspace/sglang/python/sglang/srt/layers/attention/qsa_fa_fallback.py
+
+# Inject the SM121 early-return into the QSA varlen resolver.
+COPY patches/patch_qsa_sm121.py /tmp/patch_qsa_sm121.py
+RUN set -eux; \
+    SRC="$(python3 -c 'import sglang, os; print(sglang.__path__[0])')"; \
+    python3 /tmp/patch_qsa_sm121.py --check; \
+    python3 /tmp/patch_qsa_sm121.py; \
+    python3 /tmp/patch_qsa_sm121.py --check; \
+    python3 -m py_compile \
+        "$SRC/srt/layers/attention/qwen_sparse_attn_backend.py" \
+        "$SRC/srt/layers/attention/qsa_fa_fallback.py"; \
+    rm -f /tmp/patch_qsa_sm121.py

+ 22 - 0
docker-images/qwen38-flashnext-nvfp4/build.sh

@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+# Build the SM121-patched Qwen3.8-Flash-Next-NVFP4 image.
+#
+# Generic driver: forwards every KEY=VAL pin in versions.env as a --build-arg
+# to the Dockerfile. TAG is set by the sm121 launcher to the recipe's
+# docker_image (manual build: TAG=qwen38-flashnext-nvfp4:local bash build.sh).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-qwen38-flashnext-nvfp4:local}"
+
+ARGS=()
+while IFS= read -r line; do
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"; val="${line#*=}"
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+DOCKER_BUILDKIT=1 docker build "${ARGS[@]}" --progress=plain -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

+ 60 - 0
docker-images/qwen38-flashnext-nvfp4/patches/patch_qsa_sm121.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+"""Patch SGLang's Qwen sparse-attention (QSA) varlen-decode path for SM121
+(DGX Spark / GB10).
+
+The stock `lmsysorg/sglang:qwen38flashnext` resolves QSA's packed varlen decode
+to flash-attn-4's CuTe-DSL kernels, which fail MLIR compilation on SM121
+(layout-congruence error). This injects an early-return that routes non-SM100
+devices to the Triton FlashDecoding-style fallback in `qsa_fa_fallback.py`
+(CUDA-graph-safe: cu_seqlens are read on-device, so graph replay works).
+
+Anchor-based, refuses on drift, idempotent (--check passes on pristine and on
+already-patched source). See the vllm-mod-patches skill.
+"""
+import os
+import sys
+
+import sglang
+
+FILE = os.path.join(
+    sglang.__path__[0], "srt/layers/attention/qwen_sparse_attn_backend.py"
+)
+# The dense-varlen resolver's first try block (FA2 preferred). Insert the
+# SM121 early-return directly before it.
+OLDER = "    try:\n        from flash_attn import flash_attn_varlen_func"
+NEWER = (
+    "    from sglang.srt.utils import is_sm100_supported\n"
+    "    if not is_sm100_supported():\n"
+    "        from sglang.srt.layers.attention.qsa_fa_fallback import triton_varlen_attn_func\n"
+    "        return triton_varlen_attn_func\n"
+    + OLDER
+)
+
+
+def _consistent(src: str, old: str, new: str) -> bool:
+    co, cn = src.count(old), src.count(new)
+    if cn >= 1:  # patched: new once, and any old sits only inside new
+        return cn == 1 and co <= new.count(old)
+    return co == 1  # pristine: old once
+
+
+def main() -> None:
+    check = "--check" in sys.argv
+    src = open(FILE).read()
+    if check:
+        if not _consistent(src, OLDER, NEWER):
+            raise SystemExit(f"CHECK FAIL: {FILE} is neither pristine nor patched")
+        print(f"check OK: {FILE} state is consistent")
+        return
+    if src.count(OLDER) != 1:
+        raise SystemExit(
+            f"refusing: {FILE} anchor matched {src.count(OLDER)}x; re-sync patch"
+        )
+    out = src.replace(OLDER, NEWER)
+    assert out.count(NEWER) == 1
+    open(FILE, "w").write(out)
+    print(f"patched {FILE} for SM121")
+
+
+if __name__ == "__main__":
+    main()

+ 179 - 0
docker-images/qwen38-flashnext-nvfp4/patches/qsa_fa_fallback.py

@@ -0,0 +1,179 @@
+"""SM121 (DGX Spark / GB10) fallback varlen attention for Qwen sparse attention.
+
+Upstream `qwen_sparse_attn_backend._resolve_flash_attn_varlen_func` prefers
+classic FA2 and otherwise falls back to flash-attn-4's CuTe DSL interface.
+FA4's cute kernels do not compile on SM121 (MLIR layout-congruence error),
+so this module provides a drop-in `flash_attn_varlen_func` backed by a Triton
+FlashDecoding-style kernel specialized for the QSA call contract:
+
+  * every "sequence" has exactly ONE query row (cu_seqlens_q = arange),
+  * variable numbers of gathered KV rows per sequence (<= topk),
+  * GQA (num_q_heads a multiple of num_kv_heads),
+  * cu_seqlens_* are read on-device, so the kernel is CUDA-graph replay safe
+    (the QSA backend rewrites cu_seqlens_k contents during graph replay).
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _varlen_one_q_attn_kernel(
+    q_ptr, k_ptr, v_ptr, o_ptr,
+    cu_seqlens_q_ptr, cu_seqlens_k_ptr,
+    sm_scale,
+    HQ: tl.constexpr, HKV: tl.constexpr,
+    D: tl.constexpr, D_PAD: tl.constexpr,
+    BLOCK_KV: tl.constexpr,
+    q_stride_t: tl.constexpr, q_stride_h: tl.constexpr,
+    k_stride_t: tl.constexpr, k_stride_h: tl.constexpr,
+    v_stride_t: tl.constexpr, v_stride_h: tl.constexpr,
+    o_stride_t: tl.constexpr, o_stride_h: tl.constexpr,
+):
+    """One program per (query row, q head): online-softmax attention over the
+    gathered KV rows of that query, 1-token query, causal is a no-op."""
+    row = tl.program_id(0)
+    head = tl.program_id(1)
+
+    # QSA always emits one query per varlen sequence.
+    q_start = tl.load(cu_seqlens_q_ptr + row)
+    k_start = tl.load(cu_seqlens_k_ptr + row)
+    k_end = tl.load(cu_seqlens_k_ptr + row + 1)
+
+    offs_d = tl.arange(0, D_PAD)
+    mask_d = offs_d < D
+
+    q = tl.load(
+        q_ptr + (q_start * q_stride_t) + (head * q_stride_h) + offs_d,
+        mask=mask_d, other=0.0,
+    ).to(tl.float32)
+
+    kh = head // (HQ // HKV)
+
+    m_i = -float("inf")
+    l_i = 0.0
+    acc = tl.zeros([D_PAD], dtype=tl.float32)
+
+    for k0 in range(k_start, k_end, BLOCK_KV):
+        offs_kv = k0 + tl.arange(0, BLOCK_KV)
+        mask_kv = offs_kv < k_end
+        kv_ptrs = (offs_kv * k_stride_t) + (kh * k_stride_h)
+        k_blk = tl.load(
+            k_ptr + kv_ptrs[:, None] + offs_d[None, :],
+            mask=mask_kv[:, None] & mask_d[None, :], other=0.0,
+        ).to(tl.float32)
+        scores = tl.sum(q[None, :] * k_blk, axis=1) * sm_scale
+        scores = tl.where(mask_kv, scores, -float("inf"))
+
+        m_new = tl.maximum(m_i, tl.max(scores, axis=0))
+        alpha = tl.exp(m_i - m_new)
+        p = tl.exp(scores - m_new)
+        l_i = l_i * alpha + tl.sum(p, axis=0)
+        acc = acc * alpha
+
+        v_ptrs = (offs_kv * v_stride_t) + (kh * v_stride_h)
+        v_blk = tl.load(
+            v_ptr + v_ptrs[:, None] + offs_d[None, :],
+            mask=mask_kv[:, None] & mask_d[None, :], other=0.0,
+        ).to(tl.float32)
+        acc += tl.sum(p[:, None] * v_blk, axis=0)
+        m_i = m_new
+
+    out = acc / tl.where(l_i > 0.0, l_i, 1.0)
+    tl.store(
+        o_ptr + (row * o_stride_t) + (head * o_stride_h) + offs_d,
+        out.to(o_ptr.dtype.element_ty),
+        mask=mask_d,
+    )
+
+
+def _next_pow2(n: int) -> int:
+    return triton.next_power_of_2(max(n, 16))
+
+
+def triton_varlen_attn_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_q: int = 1,
+    max_seqlen_k: int = 0,
+    softmax_scale: float = 1.0,
+    causal: bool = True,
+    **kwargs,
+) -> torch.Tensor:
+    """Drop-in replacement for flash_attn.flash_attn_varlen_func (QSA calls).
+
+    Supports the exact shape the QSA sparse backend issues: one query row per
+    varlen sequence, GQA, matching q/k/v dtypes (bf16/fp16), any head_dim.
+    """
+    if not q.is_cuda:
+        raise RuntimeError("qsa_fa_fallback requires CUDA tensors")
+    if q.dim() != 3 or k.dim() != 3 or v.dim() != 3:
+        raise RuntimeError(f"expected 3D q/k/v, got {q.shape}/{k.shape}/{v.shape}")
+    if k.dtype != q.dtype or v.dtype != q.dtype:
+        raise RuntimeError(
+            f"qsa_fa_fallback: q/k/v dtypes must match "
+            f"({q.dtype}/{k.dtype}/{v.dtype}); keep KV cache in model dtype"
+        )
+    if q.dtype not in (torch.bfloat16, torch.float16):
+        raise RuntimeError(f"qsa_fa_fallback: unsupported dtype {q.dtype}")
+
+    total_q, hq, d = q.shape
+    total_k, hkv, dk = k.shape
+    if dk != d or v.shape[2] != d:
+        raise RuntimeError("head_dim mismatch")
+    if hq % hkv != 0:
+        raise RuntimeError(f"GQA mismatch: {hq} q heads vs {hkv} kv heads")
+
+    num_seqs = total_q
+    if cu_seqlens_k.numel() != num_seqs + 1:
+        raise RuntimeError("cu_seqlens_k size mismatch")
+
+    # Host-side syncs (torch.equal) are illegal inside CUDA graph capture;
+    # the eager path has already validated the one-query-per-sequence contract
+    # before any graph is captured, so skip the check while capturing.
+    if not torch.cuda.is_current_stream_capturing():
+        q_lens_ok = bool(
+            torch.equal(
+                cu_seqlens_q[1:] - cu_seqlens_q[:-1],
+                torch.ones_like(cu_seqlens_q[1:]),
+            )
+        )
+        if not q_lens_ok:
+            raise RuntimeError(
+                "qsa_fa_fallback only supports q_len==1 per varlen sequence "
+                "(the QSA backend's only call shape)"
+            )
+
+    if cu_seqlens_q.dtype != torch.int32:
+        cu_seqlens_q = cu_seqlens_q.to(torch.int32)
+    if cu_seqlens_k.dtype != torch.int32:
+        cu_seqlens_k = cu_seqlens_k.to(torch.int32)
+
+    q_c = q if q.is_contiguous() else q.contiguous()
+    k_c = k if k.is_contiguous() else k.contiguous()
+    v_c = v if v.is_contiguous() else v.contiguous()
+    out = torch.empty_like(q_c)
+
+    BLOCK_KV = 64
+    grid = (num_seqs, hq)
+    _varlen_one_q_attn_kernel[grid](
+        q_c, k_c, v_c, out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        softmax_scale,
+        HQ=hq, HKV=hkv,
+        D=d, D_PAD=_next_pow2(d),
+        BLOCK_KV=BLOCK_KV,
+        q_stride_t=q_c.stride(0), q_stride_h=q_c.stride(1),
+        k_stride_t=k_c.stride(0), k_stride_h=k_c.stride(1),
+        v_stride_t=v_c.stride(0), v_stride_h=v_c.stride(1),
+        o_stride_t=out.stride(0), o_stride_h=out.stride(1),
+        num_warps=4,
+    )
+    return out

+ 4 - 0
docker-images/qwen38-flashnext-nvfp4/versions.env

@@ -0,0 +1,4 @@
+# Authoritative build inputs for the Qwen3.8-Flash-Next-NVFP4 SM121 image.
+# Base image: SGLang day-0 for qwen4_exp (arm64, SM121-compatible). The
+# patches/ dir carries the SM121 QSA decode kernel fix (see Dockerfile).
+BASE_IMAGE=lmsysorg/sglang:qwen38flashnext

+ 0 - 72
recipes/qwen38-flash-next-fp8/recipe.toml

@@ -1,72 +0,0 @@
-# Qwen3.8-Flash-Next-FP8 (SGLang day-0) — 2× DGX Spark / GB10 (sm_121), TP=2.
-#
-# Qwen4-architecture preview: 125B-A6B MoE (512 experts, 10 routed + 1 shared),
-# Gated DeltaNet + Qwen Sparse Attention, 51B n-gram "engram" table, 4B
-# in-checkpoint MTP. FP8 fine-grained (block 128). Text + vision. Native 262K
-# ctx (YaRN to 1M). Model: https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8
-#
-# Engine: SGLang day-0 — the ONLY engine with qwen4_exp today. vLLM support
-# is an open PR (#53896), merged into no GB10 fork image (aidendle94 / eugr /
-# local-inference-lab are all pre-release pins).
-#
-# Image: lmsysorg/sglang:qwen38flashnext — multi-arch incl. arm64 (digest
-# sha256:14ed58251…; re-verify at serve time). Day-0 = boots, not tuned.
-#
-# Memory: 172.8 GiB weights → 86.4 GiB/node. SGLang claims the 51B engram
-# table offloads to host RAM (async prefetch) — if active, the GPU budget
-# shrinks and a larger KV pool / longer ctx is possible.
-#
-# PREREQUISITES:
-#   Weights cached on every node (dgx1/dgx2 ~/.cache/huggingface → NAS):
-#       huggingface-cli download Qwen/Qwen3.8-Flash-Next-FP8
-#
-# REFERENCE:
-#   https://forums.developer.nvidia.com/t/qwen3-8-flash-next/381228
-
-engine = "sglang"
-name = "Qwen3.8-Flash-Next FP8 (SGLang day-0)"
-docker_image = "lmsysorg/sglang:qwen38flashnext"
-model = "Qwen/Qwen3.8-Flash-Next-FP8"
-required_nodes = 2
-tensor_parallel_size = 2
-pipeline_parallel_size = 1
-distributed_backend = "mp"
-shm_size = "32gb"
-
-entrypoint = ["python3", "-m", "sglang.launch_server"]
-
-args = [
-    # Memory — weights ~86.4 GiB of 119.6 GiB/node; the rest goes to KV +
-    # activations. Undervalue to be safe on day-0; raise after measuring.
-    "--mem-fraction-static", "0.80",
-    "--max-running-requests", "10",
-    # GDN (DeltaNet) hybrid state pool
-    "--max-mamba-cache-size", "30",
-    # Attention / KV
-    "--attention-backend", "flashinfer",
-    "--kv-cache-dtype", "fp8_e4m3",
-    "--mamba-ssm-dtype", "bfloat16",
-    "--mamba-full-memory-ratio", "4.21",
-    "--mamba-radix-cache-strategy", "extra_buffer_lazy",
-    # Chunked prefill
-    "--chunked-prefill-size", "8192",
-    # Compile + decode graphs
-    "--enable-torch-compile", "--torch-compile-max-bs", "4",
-    "--cuda-graph-max-bs-decode", "4",
-    "--num-continuous-decode-steps", "2",
-    # Tokenizer / reasoning / tool calling (Qwen3 family)
-    "--reasoning-parser", "qwen3",
-    "--tool-call-parser", "qwen3_coder",
-    # Observability
-    "--enable-metrics",
-    "--enable-cache-report",
-]
-
-[env]
-# Offline — weights must be cached on each node
-HF_HUB_OFFLINE = "1"
-
-# Architecture
-TORCH_CUDA_ARCH_LIST = "12.1a"
-CUTE_DSL_ARCH = "sm_121a"
-OMP_NUM_THREADS = "4"

+ 87 - 0
recipes/qwen38-flash-next-nvfp4/recipe.toml

@@ -0,0 +1,87 @@
+# Qwen3.8-Flash-Next (NVFP4, SGLang) — 2× DGX Spark / GB10 (sm_121), TP=2.
+#
+# Qwen4-architecture preview: 176B total (6B active) GDN + QSA hybrid MoE,
+# 512 routed experts (top-10) + shared expert, 51B n-gram "engram" table
+# (PLE, offloaded to host RAM), in-checkpoint MTP draft. Text + vision.
+# Native 262144 ctx (YaRN to 1M). Always-thinking.
+#
+# QUANT: NVFP4 (RadixArk), NOT the official FP8. FP8 cannot serve on 2×
+# Spark: its MoE weight_block_size=[128,128] + moe_intermediate_size=640
+# violates SGLang's (intermediate/tp) % block_n == 0 at tp=2 (sgl-project/
+# sglang #21810) — only tp in {1,5} align, and tp=1 can't hold 125B. NVFP4
+# is the community-verified 2× Spark path. SGLang's own cookbook lists the
+# FP8 quant's supported HW as H200/B200/B300/GB300 — not GB10.
+#
+# Image: qwen38-flashnext-nvfp4:local — built by the launcher from
+#   docker-images/qwen38-flashnext-nvfp4/ (base lmsysorg/sglang:qwen38flashnext
+#   + a SM121 QSA decode kernel fix). The stock day-0 image cannot serve
+#   qwen4_exp on GB10: QSA's packed varlen decode resolves to flash-attn-4
+#   CuTe-DSL kernels that fail MLIR compilation on SM121; the overlay swaps in
+#   a CUDA-graph-safe Triton fallback. Provenance/validation:
+#     https://github.com/MiaAI-Lab/Qwen3.8-Flash-Next-Dual-DGX-Sparks
+#
+# Model: RadixArk/Qwen3.8-Flash-Next-NVFP4 (~135 GB)
+#   https://huggingface.co/RadixArk/Qwen3.8-Flash-Next-NVFP4
+#
+# REFERENCE:
+#   https://forums.developer.nvidia.com/t/qwen3-8-flash-next/381228
+
+engine = "sglang"
+name = "Qwen3.8-Flash-Next NVFP4 (SGLang, 2× Spark)"
+docker_image = "qwen38-flashnext-nvfp4:local"
+docker_image_dir = "qwen38-flashnext-nvfp4"
+model = "RadixArk/Qwen3.8-Flash-Next-NVFP4"
+required_nodes = 2
+tensor_parallel_size = 2
+pipeline_parallel_size = 1
+distributed_backend = "mp"
+shm_size = "32gb"
+
+entrypoint = ["python3", "-m", "sglang.launch_server"]
+
+args = [
+    "--served-model-name", "Qwen3.8-Flash-Next-NVFP4",
+    "--trust-remote-code",
+    # NVFP4 (ModelOpt) weights
+    "--quantization", "modelopt_fp4",
+    # flashinfer is required on SM120/SM121 (trtllm_mha is SM100-only)
+    "--attention-backend", "flashinfer",
+    # page 64 is pinned by SGLang for compressed QSA (full//ratio addressing)
+    "--page-size", "64",
+    # GDN (DeltaNet) hybrid state pool
+    "--mamba-ssm-dtype", "bfloat16",
+    "--mamba-radix-cache-strategy", "extra_buffer",
+    "--mamba-track-interval", "64",
+    # Memory — 0.70 leaves DRAM for the ~26 GB/rank pinned PLE (engram) table
+    "--mem-fraction-static", "0.70",
+    "--chunked-prefill-size", "4096",
+    "--max-running-requests", "16",
+    "--context-length", "262144",
+    # NEXTN speculative decoding — draft = the in-checkpoint MTP layer
+    "--speculative-algorithm", "NEXTN",
+    "--speculative-num-steps", "3",
+    "--speculative-eagle-topk", "1",
+    "--speculative-num-draft-tokens", "4",
+    # Parsers (auto = resolve from the checkpoint chat template)
+    "--reasoning-parser", "auto",
+    "--tool-call-parser", "auto",
+    "--allow-auto-truncate",
+    # Observability
+    "--enable-metrics",
+    "--enable-cache-report",
+    # SM121 house rule: prefill CUDA graphs unreliable on GB10
+    "--disable-prefill-cuda-graph",
+    # Pin decode-graph batch list (default list OOMs the graph pool on GB10)
+    "--cuda-graph-bs-decode", "1", "2", "3", "4", "5", "6", "7", "8", "10", "12", "14", "16",
+]
+
+[env]
+# Offline — weights must be cached on each node
+HF_HUB_OFFLINE = "1"
+TRANSFORMERS_OFFLINE = "1"
+
+# Architecture (aarch64 / sm_121a DGX Spark)
+TORCH_CUDA_ARCH_LIST = "12.1a"
+FLASHINFER_CUDA_ARCH_LIST = "12.1a"
+CUTE_DSL_ARCH = "sm_121a"
+OMP_NUM_THREADS = "4"