feat: DFlash support for tensor-parallel Qwen3.6 targets by jkyamog · Pull Request #533 · Luce-Org/lucebox

jkyamog · 2026-07-16T15:02:16Z

Extend Qwen3.6 tensor parallelism from target-only autoregressive generation
to DFlash speculative decoding, including the device-local state and transfer
optimizations previously planned as a separate follow-up.

The draft remains on its explicit local device while target verification runs
through the TP meta backend. The change gathers complete target features for
the draft and extends rollback, recurrent state, snapshots, and prefix-cache
restore to split tensors.

Changes

Snapshot and restore recurrent TP state directly between each rank's GPU
tensors instead of gathering it through host memory.
Restore split SSM and convolution state on-device and synchronize once per
rank after rollback.
Decide DDTree eligibility before projection, avoiding duplicate LM-head
execution.
Run draft top-K on rank 0's mirrored GPU logits.
Use the optimized rollback threshold for TP while preserving upstream's
configurable rollback policy for single-device and layer-split execution.
Batch feature-ring synchronization with one get_2d and one set_2d per
non-wrapping run. Long prefills and prefix restores no longer perform two
host/device transfers and stream synchronizations per token.
Add opt-in aggregate phase timing with DFLASH_TP_PROFILE=1.

The stacked target branch also contains the review fixes for invalid device
counts, repeated split-state propagation and copies, multi-segment debug
output, legacy callback compatibility, and NaN-safe fallback all-reduce.

Verification

Full native CUDA build and 19/19 CTest targets pass on main@f258c4a.
Chain speculative generation, DDTree, and fast rollback pass.
Streaming forced-function output produces structured tool_calls.
A 26,704-token prefix snapshot is committed, found, and restored for a
26,758-token OpenClaw request. The repeated request's prefill drops from
21.4 s to 0.1 s and produces byte-identical output.
The deterministic 1,024-token OpenClaw control at 98,304 context produces
SHA-256 3614d727a950d5416a42b5e0e3713054e59ddc9a0365fb4ac1a2cbd198f5964f,
matching the saved pre-optimization and target-only continuation.
The final full matrix completed without warnings for every viable row. TP
plus draft peaked at 15,208/12,640 MiB for Q5 and 20,742/18,176 MiB for Q6
on the two 3090s.

Sustained performance

Hardware: RTX 5090 plus an NVLink-connected RTX 3090 pair at stock power.
Workload: 26,758 prompt tokens, 1,024 generated tokens, 98,304 context, Q8 KV,
Q8 draft, DDTree budget 16, no prefix cache, and a 60-second cooldown between
rows. These are single runs on the final rebased head. Every row uses 98,304
context except the explicitly marked Q6/5090 capacity row.

Model / configuration	DFlash TG	AR TG	Relative	Prefill	Decode	Total
Q5, RTX 5090	126.2 tok/s	59.5 tok/s	2.12x	10.3 s	8.1 s	18.4 s
Q5, 2x3090 layer split	28.7 tok/s	20.6 tok/s	1.39x	44.7 s	35.7 s	81.1 s
Q5, 2x3090 TP + NVLink	86.5 tok/s	40.4 tok/s	2.14x	20.7 s	11.8 s	32.5 s
Q5, 2x3090 TP, P2P disabled	75.5 tok/s	38.9 tok/s	1.94x	29.2 s	13.6 s	42.8 s
Q6, RTX 5090, 32,768 context	80.1 tok/s	48.7 tok/s	1.64x	10.3 s	12.8 s	23.1 s
Q6, 2x3090 layer split	21.8 tok/s	17.7 tok/s	1.23x	46.0 s	47.0 s	93.8 s
Q6, 2x3090 TP + NVLink	60.5 tok/s	34.8 tok/s	1.74x	21.4 s	16.9 s	38.4 s
Q6, 2x3090 TP, P2P disabled	55.2 tok/s	34.3 tok/s	1.61x	29.9 s	18.6 s	48.5 s

Q5 performs 192 target forwards for 192 DDTree steps; Q6 performs 190 for
190. The optimized Q5 NVLink phase profile is:

draft=2.379s project=0.407s snapshot=0.354s verify=8.430s
rollback=0.192s replay=0.000s feature=0.061s accounted=11.823s

Q5 TP is 3.01x faster than Q5 layer split and Q6 TP is 2.78x faster than Q6
layer split. Direct-P2P and P2P-disabled TP runs for both quants produced the
same byte-identical output hash shown above.

Q5 on one 24 GB 3090 failed while allocating a 1,411 MiB rollback-cache buffer
after reaching 23,286 MiB. The server logged prefill_failed, but the HTTP
layer returned 200 with empty assistant content and finish_reason=stop; that
error-mapping behavior is separate from this TP PR. Q6 on the 5090 aborts with
CUDA OOM at both 98,304 and 65,536 context, but succeeds at 32,768 context with
a 31,569 MiB peak. That constrained run produced the same assistant-content
hash as Q5 on the 5090:
ac5279693a48ecc1b5e19daeadda5872f03bb47e40d589c6a40524db0596ef34.

HIP was not compiled on this NVIDIA host because ROCm is not installed. The
new pointer-query calls use the existing CUDA/HIP compatibility header.

cubic-dev-ai

All reported issues were addressed across 24 files

_{Reply with feedback, questions, or to request a fix.

Re-trigger cubic}

jkyamog · 2026-07-18T07:56:25Z

fyi: I have been using this for more than a day, likely used around 20-50M tokens already. So far its been good, saw a parser bug but not related to this PR.

Hardware: 2x RTX 3090 24 GB
Power limit: 300 W per GPU
Target: Qwen3.6-27B-Q6_K.gguf
Draft: dflash-draft-3.6-q8_0.gguf
Target execution: tensor parallel across both GPUs
CUDA peer access/NVLink: enabled
Context: 262,144 tokens
Default max output: 16,384 tokens
Thinking limit: 15,488 tokens
KV cache: Q8_0 K and V
DFlash/DDTree budget: 22
Draft SWA: 2,048
Flash-attention window: disabled/full-context attention
Prefix-cache slots: 4
Prefill-cache slots: 0
Fast rollback: enabled

This not too far from my 5090, but the 2x 3090 being 48gb has 2x the context size.

howard0su · 2026-07-20T00:30:26Z

    std::fprintf(stderr, "[backend_factory] detected arch=%s\n", arch.c_str());

-    if (args.ds4_prefill_mode_set && arch != "deepseek4") {
+    if (args.device.is_tensor_parallel() && arch != "qwen35") {


shall we check device number here as well?

Backend factory device validation: tensor mode now validates its
placement at the factory boundary, including the minimum device count,
duplicate devices, and incompatible split weights. Runtime-visible CUDA
device IDs and the 16-device meta-backend limit remain validated in
Qwen35TensorParallelContext::create, where the CUDA registry is available.

howard0su · 2026-07-20T00:34:27Z

+    // Keep the configurable upstream policy for single-device and layer-split
+    // paths, and lower only the effective TP threshold.
+    ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy();
+    if (cfg_.device.is_tensor_parallel()) {


why this logic is not inside resolve_chain_rollback_policy?

Rollback policy: moved the TP threshold selection into
resolve_chain_rollback_policy(bool tensor_parallel). The backend now only
passes its placement mode. Tests cover TP with both default and F32
checkpoint policies.

howard0su · 2026-07-20T00:37:25Z

+}
+
+bool set_device_for_tensor(const ggml_tensor * tensor) {
+    if (!tensor || !tensor->data) return false;


what is the case tensor is NULL here? I worry this check hides the bugs.

Null tensor checks: tensor, buffer, data, and rank-local tensor
preconditions are now GGML_ASSERTs. These are internal invariants rather
than recoverable runtime failures. The helpers still return false for
actual CUDA API failures.

howard0su · 2026-07-20T00:37:34Z

+namespace {
+
+bool is_meta_tensor(const ggml_tensor * tensor) {
+    return tensor && tensor->buffer &&


what is the case tensor is NULL here? I worry this check hides the bugs.

Null tensor checks: tensor, buffer, data, and rank-local tensor
preconditions are now GGML_ASSERTs. These are internal invariants rather
than recoverable runtime failures. The helpers still return false for
actual CUDA API failures.

howard0su · 2026-07-20T00:39:47Z

+            ? src->ne[0] * src->ne[1] * src->ne[2]
+            : 0;
+        if (!src || !dst || dst->type != GGML_TYPE_F32 ||
+            capture_index < 0 || capture_index >= src->ne[3] ||


capture_index can be verified out of the loop. please double check so much protective checks. some checks looks like ASSERT to me.

Repeated validation: capture_index and convolution source-slot
validation now happen once before the rank loop. Shape, type, and bounds
relationships are asserted.

howard0su · 2026-07-20T00:40:13Z

+        ggml_tensor * dst = ggml_backend_meta_simple_tensor(destination, rank);
+        if (!src || !dst || source_slots.empty() ||
+            src->type != GGML_TYPE_F32 ||
+            dst->type != GGML_TYPE_F32 ||


why this have to be F32?

Convolution F32 restriction: removed. Convolution rollback now supports
any matching non-quantized scalar type and derives offsets and widths from
the tensor element size. SSM destination remains F32 because that path
explicitly converts captured F16/F32 state into the live F32 recurrent
state.

howard0su · 2026-07-20T00:41:56Z

@@ -0,0 +1,31 @@
+#include "ggml-backend.h"


what this test really tested? seems not related to the changed code of ggml in this PR?

CUDA communicator test: this test is inherited from stacked PR #532. It
calls ggml_backend_cuda_allreduce_tensor to verify the public compatibility
symbol still links, then passes the same CUDA backend twice to verify
duplicate ranks are rejected before NCCL communicator creation. The test now
documents both purposes inline.

jkyamog · 2026-07-20T09:57:17Z

@howard0su Thank you very much for reviewing this PR. This code base is definitely not something I can fully understand, I really appreciate your time in looking over it. All of the code is written by codex, I direct it and ask questions. I then test is on my own workload which hopefully can surface some major issues. Over the weekend 3 lucebox (Qwen 27B) was working and did all of the coding workload.

cubic-dev-ai Bot reviewed Jul 16, 2026

View reviewed changes

jkyamog force-pushed the feat-qwen35-dflash-tensor-parallel branch from 2d32fb7 to 9279619 Compare July 16, 2026 20:13

jkyamog added 4 commits July 17, 2026 08:14

vendor(ggml): add selected-device tensor parallel support

5ea373d

feat(qwen35): add target tensor parallelism

e862351

fix(meta): preserve repeated split layouts

975fcc1

fix(qwen35): validate tensor parallel placement

3eb9d70

jkyamog force-pushed the feat-qwen35-dflash-tensor-parallel branch from 9279619 to c8b3fbc Compare July 16, 2026 20:23

jkyamog added 4 commits July 19, 2026 05:58

fix(cuda): validate communicator devices and preserve API

82db48c

feat(dflash): support tensor-parallel targets

2d7699f

perf(dflash): keep TP speculative state on device

495eb5e

perf(dflash): batch TP feature mirror transfers

a3590be

jkyamog force-pushed the feat-qwen35-dflash-tensor-parallel branch from c8b3fbc to a3590be Compare July 19, 2026 05:59

howard0su suggested changes Jul 20, 2026

View reviewed changes

fix(dflash): address tensor parallel review feedback

b1634b9

Uh oh!

Conversation

jkyamog commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Verification

Sustained performance

Uh oh!

cubic-dev-ai Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jkyamog commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jkyamog commented Jul 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jkyamog commented Jul 16, 2026 •

edited

Loading

cubic-dev-ai Bot left a comment •

edited

Loading

jkyamog commented Jul 18, 2026 •

edited

Loading