Enable validated prefix caching and record DGX Spark optimization benchmarks

experiments/graph/Dockerfile

@@ -0,0 +1,2 @@
+FROM local/qwen38-flash-spark:nightly-0bfc7a15
+COPY spark_ngram_adapter.py /usr/local/lib/python3.12/dist-packages/

experiments/graph/README.md

@@ -0,0 +1,23 @@
+# CUDA Graph 实验：保留证据，未采用为默认
+
+当前 nightly 的 Qwen4Exp 使用 runtime breakable CUDA Graph。原始分割参数不能将 CPU 查表排除，
+随后只排除查表又暴露了 capture-time 哈希缓冲区尚未填充的问题。
+此目录的适配将原始哈希计算和 mmap 查表合并为一个 eager-break custom op，输出原地写入。
+
+该实现通过小型 GPU 捕获/重放测试，以及完整模型启动、短问答和大部分长输入测试。
+单流中位数 31.799 tokens/s，对比 eager 30.502，约 +4.3%。
+32K 重复请求用尽了最初设置的 256 输出 token（全部为 reasoning），未产生最终答案，
+导致该轮完整验收未通过并自动恢复 eager。不能据此断言模型算错或图执行有错误；
+图候选没有在提高输出预算后重跑。因此保留它作为实验，不宣传为已通过的优化。
+
+原始结果：../../docs/results/graph2.jsonl。方法与限制：../../docs/benchmark-method.md。
+保持与主部署隔离，默认启动脚本不会启用此目录。
+
+复现实验需先有原始稳定镜像 local/qwen38-flash-spark:nightly-0bfc7a15：
+
+    docker build -t local/qwen38-flash-spark:graph2-0bfc7a15 experiments/graph
+    docker run --rm -i --gpus all -e VLLM_USE_BREAKABLE_CUDAGRAPH=1 \
+      --entrypoint python3 local/qwen38-flash-spark:graph2-0bfc7a15 -u - < experiments/graph/test_disk_adapter.py
+
+运行完整服务时将本目录 compose.yaml 作为根 compose.yaml 的 override；这会替换现有服务，
+不适合与其并行运行。实验前先保留原始配置，完成后恢复默认部署。

experiments/graph/compose.yaml

@@ -0,0 +1,21 @@
+# Experimental: only promote after benchmark and correctness checks.
+services:
+  vllm:
+    image: local/qwen38-flash-spark:graph2-0bfc7a15
+    environment:
+      VLLM_USE_BREAKABLE_CUDAGRAPH: "1"
+    command:
+      - >-
+        exec vllm serve /root/.cache/huggingface/hub/models--nvidia--Qwen3.8-Flash-Next-NVFP4/snapshots/fc694b54fb0174e0913e6adf86691ef85a4ead47
+        --served-model-name qwen3.8-flash-next
+        --host 0.0.0.0 --port 8000
+        --tensor-parallel-size 1
+        --dtype bfloat16 --kv-cache-dtype auto
+        --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION:-0.80}
+        --max-model-len 262144 --max-num-seqs ${MAX_NUM_SEQS:-4} --max-num-batched-tokens 2048
+        --enable-chunked-prefill --no-enable-prefix-caching
+        --speculative-config '{"method":"mtp","num_speculative_tokens":2}'
+        --compilation-config '{"mode":3,"cudagraph_mode":"PIECEWISE","splitting_ops":["vllm::unified_attention_with_output","vllm::unified_mla_attention_with_output","vllm::mamba_mixer2","vllm::mamba_mixer","vllm::short_conv","vllm::qwen4_exp_ple_short_conv","vllm::qwen4_exp_qsa_with_output","vllm::linear_attention","vllm::qwen_gdn_attention_core","vllm::qwen_gdn_attention_core_fused_norm_packed","vllm::gdn_attention_core_xpu","vllm::olmo_hybrid_gdn_full_forward","vllm::sparse_attn_indexer","vllm::rocm_aiter_sparse_attn_indexer","vllm::deepseek_v4_attention","vllm::hpc_rope_norm_forward","vllm::unified_kv_cache_update","vllm::unified_mla_kv_cache_update","vllm::spark_ple_lookup"],"cudagraph_capture_sizes":[1,2,4,8,12]}' --no-enable-flashinfer-autotune
+        --load-format safetensors
+        --reasoning-parser qwen3 --tool-call-parser qwen3_xml --enable-auto-tool-choice
+        --api-key "$$(cat /run/secrets/qwen_api_key)"

experiments/graph/spark_ngram_adapter.py

@@ -0,0 +1,94 @@
+"""PLE mmap adapter for pinned nightly 0bfc7a15, single GPU only.
+
+Uses the reviewed blazux mmap reader, retaining upstream hashing and dequantization.
+"""
+import torch
+from vllm_ple_mmap import _MmapNgramEmbedding, _setup_table_v029, _REGISTRY, _register_op
+
+
+def _lookup_with_hash(input_ids: torch.Tensor, query_start_loc: torch.Tensor,
+                      ngram_context: torch.Tensor, output: torch.Tensor,
+                      layer_name: str) -> None:
+    # Capture records preceding kernels without executing them. CPU lookup must
+    # therefore compute its hash IDs in this eager segment as well.
+    from vllm_ple_mmap import _lookup_ids_impl
+    layer = _REGISTRY[layer_name]
+    ids = layer.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
+    _lookup_ids_impl(ids, output, layer_name)
+
+
+def _lookup_with_hash_fake(input_ids: torch.Tensor, query_start_loc: torch.Tensor,
+                           ngram_context: torch.Tensor, output: torch.Tensor,
+                           layer_name: str) -> None:
+    pass
+
+
+class DiskEmbedding(_MmapNgramEmbedding):
+    supports_prefetch = False
+
+    def __init__(self, n, d, **kwargs):
+        super().__init__(n, d)
+        self.register_buffer("weight", torch.empty(0, dtype=torch.float8_e4m3fn), persistent=False)
+
+    def dequantize(self, embeddings, output_dtype):
+        if self.table is None:
+            raise RuntimeError("PLE disk table not loaded")
+        return embeddings.to(output_dtype) * self.weight_scale.to(output_dtype)
+
+    def start_prefetch(self, *args):
+        pass
+
+
+def apply(cls):
+    import sys
+    from vllm.config import get_current_vllm_config
+    from vllm.distributed import get_etp_group
+    mod = sys.modules[cls.__module__]
+    original_init, original_load = cls.__init__, cls.load_weights
+
+    def init(self, *args, **kwargs):
+        if get_etp_group().world_size != 1:
+            raise RuntimeError("Spark disk adapter supports only ETP=1")
+        device_cls = mod.Qwen4ExpPLEDeviceEmbedding
+        host_cls = mod.Qwen4ExpPLEPinnedHostEmbedding
+        mod.Qwen4ExpPLEDeviceEmbedding = mod.Qwen4ExpPLEPinnedHostEmbedding = DiskEmbedding
+        try:
+            original_init(self, *args, **kwargs)
+        finally:
+            mod.Qwen4ExpPLEDeviceEmbedding, mod.Qwen4ExpPLEPinnedHostEmbedding = device_cls, host_cls
+        self._ple_mmap_prefix = kwargs["prefix"]
+        self._ple_mmap_model_path = get_current_vllm_config().model_config.model
+        _REGISTRY[self._ple_mmap_prefix] = self
+
+    def load(self, weights):
+        loaded = set()
+        def filtered():
+            for name, tensor in weights:
+                if name.startswith("ngram_embedding.shard_") and name.endswith(".weight"):
+                    continue
+                if name == "ngram_embedding.weight_scale":
+                    self.register_buffer("_offload_weight_scale", tensor.detach().to("cuda"), persistent=False)
+                    continue
+                yield name, tensor
+        loaded.update(original_load(self, filtered()))
+        _setup_table_v029(self)
+        self.ngram_embedding.weight_scale = self._offload_weight_scale
+        return loaded
+
+    def forward(self, hidden_states, input_ids, query_start_loc, ngram_context):
+        output = torch.empty((input_ids.shape[0], self.embedding_dim),
+                             dtype=torch.float8_e4m3fn, device=input_ids.device)
+        torch.ops.vllm.spark_ple_lookup(input_ids, query_start_loc, ngram_context,
+                                       output, self._ple_mmap_prefix)
+        return output
+
+    from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+    from vllm.utils.torch_utils import direct_register_custom_op
+    _register_op()
+    if not hasattr(torch.ops.vllm, "spark_ple_lookup"):
+        direct_register_custom_op(
+            op_name="spark_ple_lookup",
+            op_func=eager_break_during_capture(_lookup_with_hash),
+            mutates_args=["output"], fake_impl=_lookup_with_hash_fake,
+        )
+    cls.__init__, cls.load_weights, cls.forward = init, load, forward

experiments/graph/test_disk_adapter.py

@@ -0,0 +1,79 @@
+import os
+os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
+import tempfile
+from types import SimpleNamespace
+from pathlib import Path
+import torch
+import vllm.config
+import vllm.distributed
+from safetensors.torch import save_file
+
+with tempfile.TemporaryDirectory() as folder:
+    fake_config = SimpleNamespace(engram_config=None, model_config=SimpleNamespace(model=folder))
+    vllm.config.get_current_vllm_config = lambda: fake_config
+    vllm.distributed.get_etp_group = lambda: SimpleNamespace(world_size=1)
+    from vllm.models.qwen4_exp.nvidia.ngram_embedding import Qwen4ExpNGramEmbedding
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+    cfg = SimpleNamespace(ngram_size=3, heads_per_ngram=2, eos_token_id=0,
+        vocab_size=100, split_ngram_parts=2, seed=1234,
+        ngram_vocab_size_base=17, make_ngram_vocab_size_divisible_by=8)
+    prefix = 'model.language_model.layers.0.ple.ple_embedding'
+    with torch.device('cuda'):
+        layer = Qwen4ExpNGramEmbedding(cfg, 8, 0, 32,
+            data_parallel_rank=0, prefix=prefix,
+            quant_config=Fp8Config(is_checkpoint_fp8_serialized=True))
+    count = layer.ngram_embedding.org_vocab_size
+    full = (torch.arange(count * 2).reshape(count, 2) % 13).to(torch.float8_e4m3fn)
+    tensors = {}
+    for i, part in enumerate(full.chunk(2)):
+        tensors[f'{prefix}.ngram_embedding.shard_{i}.weight'] = part.contiguous()
+    tensors[f'{prefix}.ngram_embedding.weight_scale'] = torch.tensor(0.25)
+    save_file(tensors, str(Path(folder) / 'ple.safetensors'))
+    layer.load_weights((name.removeprefix(prefix + '.'), value) for name, value in tensors.items())
+    ids = torch.tensor([[0, 1, 1, count - 1], [5, 2, 9, 7]], device='cuda')
+    out = layer.ngram_embedding(ids)
+    expected = full.view(torch.uint8)[ids.cpu()].view(torch.float8_e4m3fn).cuda()
+    assert torch.equal(out.view(torch.uint8), expected.view(torch.uint8))
+    dequant = layer.ngram_embedding.dequantize(out, torch.bfloat16)
+    assert torch.equal(dequant, expected.to(torch.bfloat16) * 0.25)
+    assert layer.ngram_embedding.weight.numel() == 0
+    tokens = torch.tensor([3, 7, 9], device='cuda')
+    starts = torch.tensor([0, 3], dtype=torch.int32, device='cuda')
+    context = torch.tensor([[0, 0]], device='cuda')
+    hashed = layer.compute_ngram_ids(tokens, starts, context)
+    actual = layer(None, tokens, starts, context)
+    reference = full.view(torch.uint8)[hashed.cpu()].view(torch.float8_e4m3fn).cuda().flatten(-2)
+    assert torch.equal(actual.view(torch.uint8), reference.view(torch.uint8))
+    print('PASS: disk rows, repeated IDs, boundary IDs, FP8 bytes and scaling; no resident PLE table')
+compiled = torch.compile(lambda t, s, c: layer(None, t, s, c), fullgraph=True)
+actual_compiled = compiled(tokens, starts, context)
+assert torch.equal(actual_compiled.view(torch.uint8), reference.view(torch.uint8))
+print('COMPILE_PASS: full graph matches exact FP8 table lookup')
+
+# Runtime graph capture must break around the CPU mmap operation, and replay
+# must consume changed inputs instead of reusing capture-time row values.
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+original_hash = layer.compute_ngram_ids
+def checked_hash(*args):
+    assert not torch.cuda.is_current_stream_capturing(), "CPU-dependent hash captured"
+    return original_hash(*args)
+layer.compute_ngram_ids = checked_hash
+stream = torch.cuda.Stream()
+stream.wait_stream(torch.cuda.current_stream())
+with torch.cuda.stream(stream):
+    layer(None, tokens, starts, context)
+    torch.cuda.synchronize()
+    capture = BreakableCUDAGraphCapture()
+    with capture:
+        captured = layer(None, tokens, starts, context)
+        downstream = captured.to(torch.float32) * 0.25
+    assert capture.num_eager_breaks >= 1
+    for values in ([4, 6, 8], [9, 2, 5], [3, 7, 9]):
+        tokens.copy_(torch.tensor(values, device="cuda"))
+        capture.replay()
+        expected_ids = layer.compute_ngram_ids(tokens, starts, context)
+        expected_bytes = full.view(torch.uint8)[expected_ids.cpu()].cuda().flatten(-2)
+        assert torch.equal(captured.view(torch.uint8), expected_bytes)
+        expected_fp8 = expected_bytes.view(torch.float8_e4m3fn)
+        assert torch.equal(downstream, expected_fp8.to(torch.float32) * 0.25)
+print("BREAKABLE_GRAPH_PASS: CPU lookup excluded, changed-input replay byte-exact")