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

runtime/patch_unify2.py

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