docker-images: qwen38-flashnext-nvfp4

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

docker-images/qwen38-flashnext-nvfp4/build.sh

@@ -0,0 +1,30 @@
+#!/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} ==="
+# Persistent buildkit cache (per node, shared across image dirs): the
+# version-independent stages (apt, torch, CUDA base) are reused across
+# rebuilds — only the pinned source stages (vllm/b12x/…) re-download and
+# re-compile. Override the dir with SM121_BUILDKIT_CACHE if desired.
+CACHE_DIR="${SM121_BUILDKIT_CACHE:-$HOME/.cache/sm121-buildkit}"
+mkdir -p "${CACHE_DIR}"
+docker buildx build "${ARGS[@]}" --cache-from "type=local,src=${CACHE_DIR}" \
+    --cache-to "type=local,dest=${CACHE_DIR},mode=max" \
+    --progress=plain --load -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

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

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

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