TBNilles/qwen3.5-122B-A10B-on-spark
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

qwen3.5-122B-A10B on DGX Spark: vLLM + DFlash + dense-bandwidth stack, one-shot installer

Browse Source
This commit is contained in:
ent
2026-06-24 13:02:35 +10:00
commit 60bf1b7b02
21 changed files with 2353 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
runtime/mtp_serve.sh
Executable
+25
View File
@@ -0,0 +1,25 @@
#!/bin/bash
# mtp_serve.sh — native qwen3_5 MTP-N head for the comparison (`--profile mtp`).
# The MTP head (1 layer, reuses target KV/embed/lm_head) is in the Intel
# checkpoint (mtp.layers.0) — no separate drafter, no unify patch needed.
# $1 = num_speculative_tokens (default 2 = the "MTP-2" recipe); $2 = backend.
set -euo pipefail
NSPEC="${1:-2}"
BACKEND="${2:-flash_attn}"
MODEL="${MODEL:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
GPU_MEM="${GPU_MEM:-0.8}"
PORT="${PORT:-8000}"
echo "[mtp] qwen3_5_mtp — backend=$BACKEND, num_speculative_tokens=$NSPEC, model=$MODEL"
exec vllm serve "$MODEL" \
--served-model-name qwen \
--host 0.0.0.0 --port "$PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--max-num-seqs 16 \
--max-num-batched-tokens "$MAX_MODEL_LEN" \
--gpu-memory-utilization "$GPU_MEM" \
--no-enable-prefix-caching \
--enable-chunked-prefill \
--trust-remote-code \
--attention-backend "$BACKEND" \
--speculative-config "{\"method\":\"qwen3_5_mtp\",\"num_speculative_tokens\":$NSPEC,\"model\":\"$MODEL\"}"
runtime/patch_fla_shmem.py
+32
View File
@@ -0,0 +1,32 @@
#!/usr/bin/env python3
"""spark-fla-shmem: let the FLA (flash-linear-attention / GDN) Triton kernels use their
BIG tiles on sm121 (GB10 / DGX Spark).
The FLA Backend gate (vllm/model_executor/layers/fla/ops/utils.py, identical in SGLang)
picks big tiles only if device_max_shared_mem >= Backend.DEFAULT (102400 = 100 KiB):
cumsum.py : BS_LIST = [32,64] if check_shared_mem() else [16,32]
chunk_o.py: BKV_LIST = [64,128] if check_shared_mem() else [32,64]
sm121 reports max_shared_mem = 101376 (99 KiB), JUST below 102400 -> check returns False
-> small tiles -> slower GDN/linear-attention (the per-token hot path). But 101376 is
EXACTLY Backend.ADA, and RTX 4090 (ADA, same 99 KiB) runs the big tiles fine, so they
provably fit in 99 KiB. Lower DEFAULT to 101376 so sm121 (and ADA) pass the gate.
"""
import pathlib, sys
P = pathlib.Path(
"/usr/local/lib/python3.12/site-packages/vllm/model_executor/layers/fla/ops/utils.py"
)
src = P.read_text()
if "spark-fla-shmem" in src:
print("[patch_fla_shmem] already patched", flush=True)
sys.exit(0)
OLD = " DEFAULT = 102400 # Default"
NEW = " DEFAULT = 101376 # spark-fla-shmem: 102400->101376 so sm121 GB10 (99 KiB) uses big GDN tiles (already fit on ADA's identical 99 KiB)"
if OLD not in src:
print("[patch_fla_shmem] ERROR: target line not found — FLA source differs", flush=True)
sys.exit(1)
P.write_text(src.replace(OLD, NEW))
print("[patch_fla_shmem] patched OK (DEFAULT 102400 -> 101376; big GDN tiles on sm121)", flush=True)
runtime/patch_inc_hybrid.py
+163
View File
@@ -0,0 +1,163 @@
#!/usr/bin/env python3
"""spark-dflash-hybrid-fp8: port of albond's hybrid INT4+FP8 dispatch
(patches/01-hybrid-int4-fp8/inc.py.patch) onto AEON 0.23's vllm INCConfig.
Adds, to vllm/model_executor/layers/quantization/inc.py:
* INCConfig.fp8_config / fp8_layers fields
* maybe_update_config OVERRIDE — AEON 0.23 signature (model_name, hf_config=None,
revision=None); the base hook is already CALLED from config/vllm.py:634, so we
only supply the override. Scans the checkpoint's safetensors metadata for
float8_e4m3fn weights that have a .weight_scale_inv, builds an Fp8Config, and
records those layer prefixes.
* _is_layer_fp8 — exact + fused + substring match against fp8_layers
* FP8 dispatch at BOTH dense short-circuits: get_quant_method's extra_config
override AND the apply_*_quant_layer not-quantized blocks.
Idempotent; sentinel 'spark-dflash-hybrid-fp8'. Mirrors patch_unify2.py's style.
"""
import sys
import vllm.model_executor.layers.quantization.inc as inc_mod
path = inc_mod.__file__
src = open(path).read()
SENT = "spark-dflash-hybrid-fp8"
if SENT in src:
print(f"[patch_inc_hybrid] already applied: {path}")
sys.exit(0)
# 1. __init__ fields (anchor unique: only INCConfig has pack_factor = Fraction(32,..))
a1 = " self.pack_factor = Fraction(32, weight_bits)\n"
b1 = a1 + (
" # spark-dflash-hybrid-fp8: populated by maybe_update_config\n"
" self.fp8_config = None\n"
" self.fp8_layers = set()\n"
)
assert src.count(a1) == 1, f"anchor1 count={src.count(a1)}"
src = src.replace(a1, b1)
# 2. apply_vllm_mapper: remap fp8_layers (HF->vLLM names) after extra_config remap
a2 = (
" if self.extra_config is not None:\n"
" self.extra_config = hf_to_vllm_mapper.apply_dict(self.extra_config)\n"
)
b2 = a2 + (
" if self.fp8_layers: # spark-dflash-hybrid-fp8\n"
" self.fp8_layers = set(\n"
" hf_to_vllm_mapper.apply_list(list(self.fp8_layers))\n"
" )\n"
)
assert src.count(a2) == 1, f"anchor2 count={src.count(a2)}"
src = src.replace(a2, b2)
# 3. insert maybe_update_config + _is_layer_fp8 before apply_awq_quant_layer
a3 = ' def apply_awq_quant_layer(self, layer, prefix: str, backend: str = "auto"):\n'
methods = ''' def maybe_update_config( # spark-dflash-hybrid-fp8
self,
model_name: str,
hf_config=None,
revision: str | None = None,
):
"""Detect FP8 dense layers in a hybrid INT4+FP8 checkpoint."""
import torch as _torch
from safetensors.torch import _TYPES as _SF
from vllm.transformers_utils.config import get_safetensors_params_metadata
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
metadata = get_safetensors_params_metadata(model_name, revision=revision)
fp8_weights = {}
for pn, info in metadata.items():
ds = info.get("dtype", None)
if ds is None:
continue
if _SF.get(ds) == _torch.float8_e4m3fn and pn.endswith(".weight"):
sn = pn.replace(".weight", ".weight_scale_inv")
if sn in metadata:
fp8_weights[pn] = info
if not fp8_weights:
logger.info("spark-dflash-hybrid-fp8: no FP8 dense layers detected")
return
block_size = None
for pn, info in fp8_weights.items():
sn = pn.replace(".weight", ".weight_scale_inv")
ws = info.get("shape", [])
ss = metadata[sn].get("shape", [])
if len(ws) == 2 and len(ss) == 2:
block_size = [ws[0] // ss[0], ws[1] // ss[1]]
break
if block_size is None:
return
self.fp8_config = Fp8Config(
is_checkpoint_fp8_serialized=True,
activation_scheme="dynamic",
weight_block_size=block_size,
)
self.fp8_layers = {n.rsplit(".weight", 1)[0] for n in fp8_weights}
_sample = sorted(self.fp8_layers)[:3]
logger.info(
"spark-dflash-hybrid-fp8: detected %d FP8 dense layers "
"(block_size=%s) e.g. %s",
len(self.fp8_layers), block_size, _sample,
)
def _is_layer_fp8(self, prefix: str) -> bool: # spark-dflash-hybrid-fp8
if not self.fp8_layers:
return False
if prefix in self.fp8_layers:
return True
fused = getattr(self, "packed_modules_mapping", {})
proj = prefix.split(".")[-1]
if proj in fused:
shards = [prefix.replace(proj, s) for s in fused[proj]]
return all(
any(fl in sp for fl in self.fp8_layers) for sp in shards
)
return any(fl in prefix for fl in self.fp8_layers)
'''
assert src.count(a3) == 1, f"anchor3 count={src.count(a3)}"
src = src.replace(a3, methods + a3)
# 4. FP8 dispatch in the not-quantized blocks (awq/gptq/xpu/cpu are byte-identical;
# guard is inert unless fp8_config set, so patching all is safe)
a4 = (
" if not self.check_quantized(weight_bits):\n"
" if isinstance(layer, (LinearBase, ParallelLMHead)):\n"
" return UnquantizedLinearMethod()\n"
" else:\n"
" return None\n"
)
b4 = (
" if not self.check_quantized(weight_bits):\n"
" if self.fp8_config and self._is_layer_fp8(prefix): # spark-dflash-hybrid-fp8\n"
" from vllm.model_executor.layers.quantization.fp8 import (\n"
" Fp8LinearMethod,\n"
" )\n"
" return Fp8LinearMethod(self.fp8_config)\n"
" if isinstance(layer, (LinearBase, ParallelLMHead)):\n"
" return UnquantizedLinearMethod()\n"
" else:\n"
" return None\n"
)
n4 = src.count(a4)
assert n4 >= 2, f"anchor4 count={n4}"
src = src.replace(a4, b4)
# 5. FP8 dispatch in get_quant_method's extra_config (bits>=16) override
a5 = (
' ) and self.extra_config[layer_name].get("bits", 16) >= 16:\n'
" return UnquantizedLinearMethod()\n"
)
b5 = (
' ) and self.extra_config[layer_name].get("bits", 16) >= 16:\n'
" if self.fp8_config and self._is_layer_fp8(prefix): # spark-dflash-hybrid-fp8\n"
" from vllm.model_executor.layers.quantization.fp8 import (\n"
" Fp8LinearMethod,\n"
" )\n"
" return Fp8LinearMethod(self.fp8_config)\n"
" return UnquantizedLinearMethod()\n"
)
assert src.count(a5) == 1, f"anchor5 count={src.count(a5)}"
src = src.replace(a5, b5)
open(path, "w").write(src)
print(f"[patch_inc_hybrid] applied {SENT} to {path} (not-quant blocks x{n4})")
runtime/patch_int8_lmhead_v3.py
+124
View File
@@ -0,0 +1,124 @@
#!/usr/bin/env python3
"""INT8 W8A16 lm-head v3 for AEON vLLM 0.23 (sm121). Replaces ONLY the
`lm_head.quant_method.apply(...)` call inside LogitsProcessor._get_logits with a
batched int8 GEMV (one kernel launch for any batch), leaving the existing TP
gather + org_vocab_size trim untouched.
Fixes vs the broken v2 port:
* KEEPS the bf16 lm_head weight (DFlash drafter shares it) — does NOT zero it.
Trades the memory saving for correctness; the speed win is the int8 read in
_get_logits, independent of keeping bf16 around.
* Single BATCHED kernel (dot-based, pad B->16) for ALL B — no per-row Python
loop (the v2 B>4 loop was what made spec decode SLOWER).
* Fixed proven config (N128/K128/w4/s3, ~227 GB/s, argmax-exact vs bf16 on the
real [248320,3072] shape) — no autotune (avoids sm121 bad-config miscompiles).
Sentinel DGX_SPARK_INT8_LMHEAD_V3. Verified standalone: int8 3.35ms vs bf16 8.8ms
(B=1) / 6.5ms (B=13); maxerr 3e-4 < quant floor, argmax 100%.
"""
import os
import sys
TARGET = "/usr/local/lib/python3.12/site-packages/vllm/model_executor/layers/logits_processor.py"
ANCHOR = (
" # Get the logits for the next tokens.\n"
" logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)\n"
)
REPLACE = (
" # DGX_SPARK_INT8_LMHEAD_V3: int8 w8a16 GEMV for the huge vocab projection\n"
" logits = _spark_int8_lmhead_apply(self, lm_head, hidden_states, embedding_bias)\n"
)
MODULE_CODE = '''
# ===================== DGX_SPARK_INT8_LMHEAD_V3 =====================
import triton as _spark_triton
import triton.language as _spark_tl
@_spark_triton.jit
def _spark_k_int8(x_ptr, w_ptr, s_ptr, o_ptr, B, N, K,
sxb, sxk, swn, swk, sob, son,
BLOCK_B: _spark_tl.constexpr, BLOCK_N: _spark_tl.constexpr,
BLOCK_K: _spark_tl.constexpr):
pid_n = _spark_tl.program_id(0)
offs_b = _spark_tl.arange(0, BLOCK_B)
offs_n = pid_n * BLOCK_N + _spark_tl.arange(0, BLOCK_N)
offs_k = _spark_tl.arange(0, BLOCK_K)
x_ptrs = x_ptr + offs_b[:, None] * sxb + offs_k[None, :] * sxk
w_ptrs = w_ptr + offs_n[:, None] * swn + offs_k[None, :] * swk
acc = _spark_tl.zeros((BLOCK_B, BLOCK_N), dtype=_spark_tl.float32)
for k in range(0, K, BLOCK_K):
km = (offs_k[None, :] + k) < K
x = _spark_tl.load(x_ptrs, mask=(offs_b[:, None] < B) & km, other=0.0).to(_spark_tl.float16)
w = _spark_tl.load(w_ptrs, mask=(offs_n[:, None] < N) & km, other=0).to(_spark_tl.float16)
acc += _spark_tl.dot(x, w.T)
x_ptrs += BLOCK_K * sxk
w_ptrs += BLOCK_K * swk
s = _spark_tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0).to(_spark_tl.float32)
acc = acc * s[None, :]
o_ptrs = o_ptr + offs_b[:, None] * sob + offs_n[None, :] * son
_spark_tl.store(o_ptrs, acc, mask=(offs_b[:, None] < B) & (offs_n[None, :] < N))
def _spark_int8_gemm(hidden, w_int8, w_scale):
import torch
N, K = w_int8.shape
x = hidden.reshape(-1, K)
B = x.shape[0]
BLOCK_B = max(16, _spark_triton.next_power_of_2(B))
out = torch.empty(B, N, dtype=torch.float32, device=x.device)
xf = x.to(torch.float16)
grid = ((N + 127) // 128,)
_spark_k_int8[grid](xf, w_int8, w_scale, out, B, N, K,
xf.stride(0), xf.stride(1), w_int8.stride(0), w_int8.stride(1),
out.stride(0), out.stride(1),
BLOCK_B=BLOCK_B, BLOCK_N=128, BLOCK_K=128,
num_warps=4, num_stages=3)
return out.reshape(hidden.shape[:-1] + (N,))
def _spark_int8_lmhead_apply(self, lm_head, hidden_states, embedding_bias):
import sys
import torch
if not getattr(lm_head, "_spark_int8_ready", None) is True and \\
not getattr(lm_head, "_spark_int8_disabled", False):
w = getattr(lm_head, "weight", None)
if (w is not None and w.dtype in (torch.bfloat16, torch.float16)
and w.dim() == 2 and w.shape[0] > 100000):
with torch.no_grad():
scales = (w.float().abs().amax(dim=1) / 127.0).clamp(min=1e-12)
w_int8 = (w.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
lm_head._spark_w_int8 = w_int8.contiguous()
lm_head._spark_w_scale = scales.to(torch.float16)
lm_head._spark_int8_ready = True
print("DGX_SPARK_INT8_LMHEAD_V3: lm_head -> int8 (%s), bf16 kept for shared drafter"
% (list(w_int8.shape),), file=sys.stderr, flush=True)
else:
lm_head._spark_int8_disabled = True
if getattr(lm_head, "_spark_int8_ready", False) and embedding_bias is None:
return _spark_int8_gemm(hidden_states, lm_head._spark_w_int8, lm_head._spark_w_scale)
return lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)
# =================== end DGX_SPARK_INT8_LMHEAD_V3 ===================
'''
def main():
if not os.path.exists(TARGET):
print(f"FAIL: {TARGET} not found"); sys.exit(1)
src = open(TARGET).read()
if "DGX_SPARK_INT8_LMHEAD_V3" in src:
print("SKIP: int8 lm-head v3 already applied"); return
if ANCHOR not in src:
print("FAIL: _get_logits apply-anchor not found"); sys.exit(1)
if src.count(ANCHOR) != 1:
print(f"FAIL: anchor count={src.count(ANCHOR)} (expected 1)"); sys.exit(1)
src = src.replace(ANCHOR, REPLACE)
src = src + MODULE_CODE
open(TARGET, "w").write(src)
print("OK: int8 lm-head v3 applied")
if __name__ == "__main__":
main()
runtime/patch_unify2.py
+55
View File
@@ -0,0 +1,55 @@
#!/usr/bin/env python3
"""spark-dflash-unify (CORRECTED): let the DFlash drafter's attention KV spec unify
with the hybrid GDN+mamba target's *padded* page size WITHOUT corrupting the drafter.
Root cause (AEON 0.23.0+aeon.sm121a.dflash): with --mamba-block-size set, vLLM's own
hybrid alignment makes target mamba page == target attention page by PADDING the mamba
page (e.g. "+0.54%"). That padded value becomes max_page_size. The drafter's attention
page is smaller; unify_kv_cache_spec_page_size scales its block_size by
ratio = max_page_size // layer_page_size, but because max_page_size is a *padded* (non
block-size-linear) number, the scaled page lands just under it and
`assert new_spec.page_size_bytes == max_page_size` fires.
The OLD spark-pad-unify patch "fixed" this by `replace(layer_spec,
page_size_padded=max_page_size)` -- but it DROPPED the block_size scaling, leaving the
drafter at block_size=16 behind a max-sized physical page. That mis-strided the drafter
KV and pinned mean accept length at ~1.47 (garbage drafts), so DFlash net-lost.
This corrected patch mirrors what get_kv_cache_groups already does for
HiddenStateCacheSpec layers: keep the SCALED block_size AND pad the <1% remainder.
Strides stay correct (block_size is properly scaled); only the page tail is padded.
"""
import pathlib, sys
P = pathlib.Path("/usr/local/lib/python3.12/site-packages/vllm/v1/core/kv_cache_utils.py")
src = P.read_text()
if "spark-dflash-unify" in src:
print("[patch_unify2] already patched", flush=True)
sys.exit(0)
OLD = ''' new_spec = replace(layer_spec, block_size=new_block_size)
assert new_spec.page_size_bytes == max_page_size
new_kv_cache_spec[layer_name] = new_spec'''
NEW = ''' new_spec = replace(layer_spec, block_size=new_block_size)
if new_spec.page_size_bytes != max_page_size:
# spark-dflash-unify: max_page_size is a *padded* hybrid page; the
# scaled attention page lands just under it. Pad the remainder while
# KEEPING the scaled block_size, exactly as get_kv_cache_groups does
# for HiddenStateCacheSpec. (The old patch dropped the scaling ->
# mis-strided the DFlash drafter -> accept len stuck ~1.47.)
new_spec = replace(
layer_spec,
block_size=new_block_size,
page_size_padded=max_page_size,
)
assert new_spec.page_size_bytes == max_page_size
new_kv_cache_spec[layer_name] = new_spec'''
if OLD not in src:
print("[patch_unify2] ERROR: target block not found — vLLM source differs", flush=True)
sys.exit(1)
P.write_text(src.replace(OLD, NEW))
print("[patch_unify2] patched OK (scaled-block + pad-remainder)", flush=True)
runtime/serve.sh
Executable
+59
View File
@@ -0,0 +1,59 @@
#!/bin/bash
# serve.sh — runs INSIDE the sm121 vLLM container (mounted at /host). Applies the
# runtime monkeypatches, then `vllm serve`s the Qwen3.5-122B-A10B INT4 target with
# the DFlash drafter. Driven by install.sh; can also be run by hand.
#
# args: $1 = num_speculative_tokens (0 = no-spec baseline)
# $2 = target attention backend (flash_attn | FLASHINFER)
# env: MODEL target path/repo (default Intel INT4; /model for hybrid)
# INC_HYBRID=1 apply the hybrid INT4+FP8 dense-expert dispatch patch
# INT8_LMHEAD_V3=1 apply the int8 lm-head GEMV patch
# MAX_MODEL_LEN GPU_MEM PORT
#
# Stack rationale: the DFlash drafter is non-causal -> needs FLASH_ATTN (FA2). The
# hybrid GDN+mamba+MoE target's KV page geometry won't absorb the drafter's
# attention spec without patch_unify2 (scale-block unify) + prefix-caching OFF
# (NoPrefixCache coordinator, dodges the hash assert). See docs/FINDINGS.md.
set -euo pipefail
NSPEC="${1:-12}"
BACKEND="${2:-flash_attn}"
MODEL="${MODEL:-Intel/Qwen3.5-122B-A10B-int4-AutoRound}"
DRAFT="${DRAFT:-z-lab/Qwen3.5-122B-A10B-DFlash}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
GPU_MEM="${GPU_MEM:-0.8}"
PORT="${PORT:-8000}"
# FLA sm121 big-tile shmem fix (prefill/TTFT only on sm121; harmless, free).
echo "[serve] FLA sm121 big-tile shmem patch"
python3 /host/patch_fla_shmem.py || true
if [ "${INC_HYBRID:-0}" = "1" ]; then
echo "[serve] hybrid INT4+FP8 dispatch patch (inc.py)"
python3 /host/patch_inc_hybrid.py
fi
if [ "${INT8_LMHEAD_V3:-0}" = "1" ]; then
echo "[serve] int8 lm-head v3 patch (batched w8a16 GEMV)"
python3 /host/patch_int8_lmhead_v3.py
fi
if [ "$NSPEC" = "0" ]; then
SPEC_ARG=()
echo "[serve] NO-SPEC baseline (identical flags, prefix-off)"
else
SPEC_ARG=(--speculative-config "{\"method\":\"dflash\",\"model\":\"$DRAFT\",\"num_speculative_tokens\":$NSPEC,\"attention_backend\":\"FLASH_ATTN\"}")
echo "[serve] DFlash n=$NSPEC, target-backend=$BACKEND, drafter=FLASH_ATTN, model=$MODEL"
fi
python3 /host/patch_unify2.py || { [ "$NSPEC" = "0" ] && true; }
exec vllm serve "$MODEL" \
--served-model-name qwen \
--host 0.0.0.0 --port "$PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--max-num-seqs 16 \
--max-num-batched-tokens "$MAX_MODEL_LEN" \
--gpu-memory-utilization "$GPU_MEM" \
--no-enable-prefix-caching \
--enable-chunked-prefill \
--trust-remote-code \
--attention-backend "$BACKEND" \
"${SPEC_ARG[@]}"