feat: support Qwen3.5 DFlash and its optimizations#510
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6dc54f7688
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| # Sync BEFORE base writeback to flush any residual async error | ||
| # from previous draft writeback on this stream. | ||
| torch.cuda.synchronize() |
There was a problem hiding this comment.
When KVStore writeback is enabled, this unconditional CUDA sync runs before every base writeback, and the added debug block also synchronizes through .item()/.tolist() and logs at warning on the same hot path. That serializes all pending GPU work and removes the intended overlap of host/device KV transfers, so decode throughput will collapse in offload deployments. Please gate this diagnostics path or remove the synchronizations/logging before shipping.
Useful? React with 👍 / 👎.
| f"outputs. Got invalid ids {invalid}; valid range is " | ||
| f"[0, {num_layers - 2}] for {num_layers} target layers." | ||
| ) | ||
| self.model.layers_to_capture = {val + 1 for val in layer_ids} |
There was a problem hiding this comment.
When a DFlash checkpoint's dflash_config.target_layer_ids is not already sorted, this set loses the order declared by the checkpoint; the forward loop appends aux_hidden_states in ascending layer order and LogitsProcessor concatenates them before the DFlash fc. That permutes the features relative to the checkpoint, producing wrong draft K/V. Preserve the configured order while still checking membership.
Useful? React with 👍 / 👎.
8075bdc to
46059aa
Compare
0331070 to
5052889
Compare
Signed-off-by: Moustafa Saleh <8815169+mesaleh@users.noreply.github.com>
…raph warmup prefix hack
2054474 to
04c4aea
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 04c4aea31c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._ensure_fused_workspace(total_ctx, ctx_hidden.dtype) | ||
|
|
||
| proj_out_2d = self._fused_kv_proj_workspace[:total_ctx] | ||
| torch.mm(ctx_hidden, self._fused_kv_flat_weight_t, out=proj_out_2d) |
There was a problem hiding this comment.
When a DFlash draft checkpoint has attention_bias=True, this fused decode path computes K/V with a bare torch.mm, so the sliced qkv_proj.bias is dropped; the sequential kv_proj_only() path includes that bias via F.linear(..., bias). Because decode-only runs enable the fused path, biased draft attention layers will materialize different K/V into the draft cache and produce incorrect draft tokens. Please either fold the per-layer KV biases into proj_out_2d or disable the fused path for biased projections.
Useful? React with 👍 / 👎.
| else: | ||
| k_buf_ptr = tl.load(k_buf_ptrs_ptr + layer_id).to(tl.pointer_type(tl.bfloat16)) | ||
| v_buf_ptr = tl.load(v_buf_ptrs_ptr + layer_id).to(tl.pointer_type(tl.bfloat16)) |
There was a problem hiding this comment.
When fused DFlash KV is used with a non-FP8, non-BF16 draft KV cache (for example a float16 cache on the trtllm draft backend, which this PR no longer forces to BF16), this branch still casts the raw cache pointers to tl.bfloat16. The stores then write BF16-encoded bits into a float16 buffer, so later attention reads corrupted K/V. Please gate the fused scatter to BF16/E4M3 caches or pass the actual cache dtype into the Triton pointer/store path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9a60cd5929
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| num_kv_heads=self.num_kv_heads, | ||
| layer_id=layer_id, | ||
| sliding_window_size=sliding_window, | ||
| causal=causal, |
There was a problem hiding this comment.
When a DFlash draft config has both full_attention and sliding_attention and the flat scheduler ext is enabled, MHATokenToKVPool publishes multiple paged-cache groups from layer_types, and validate_paged_cache_group_ids() rejects any PagedAttention whose group_id is empty. This constructor now reads the layer type for sliding-window behavior but still leaves group_id at its default, so Qwen3.5 DFlash draft models with mixed layer types fail during executor initialization instead of serving; pass the layer type/derived group id into PagedAttention as the other grouped-cache models do.
Useful? React with 👍 / 👎.
lightseek-bot
changed the title
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 43107d56ac
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._qk_norm_eps = eps | ||
| rope_parameters = getattr(config, "rope_parameters", None) | ||
| if rope_parameters is not None: | ||
| rope_theta = float(rope_parameters["rope_theta"]) |
There was a problem hiding this comment.
When a DFlash draft checkpoint provides rope_parameters in the Qwen3.5/HF format for default or scaled RoPE, rope_theta can remain only on the top-level config while the dict contains fields like rope_type/factor; this direct index then raises KeyError during DFlashAttention construction and prevents those checkpoints from loading. Please read rope_theta from the dict only when present and otherwise fall back to config.rope_theta/the default.
Useful? React with 👍 / 👎.
| decode_only | ||
| and self._fused_kv_enabled | ||
| and self._kv_aux_stream is not None | ||
| and (capturing or not get_is_cuda_graph_phase()) |
There was a problem hiding this comment.
When incremental projection is enabled during CUDA-graph warmup, the target forward still schedules _on_capture_slot_ready on the aux stream, but this graph-phase guard makes run() take the sequential path because the stream is not yet capturing. That sequential path calls _write_native_cache_fused on the main stream using the same _fused_kv_proj_workspace and KV slots without waiting for the aux stream, so warmup can race with the aux materialization and record/read corrupted draft KV. Please either disable the incremental callbacks during graph warmup or join/reuse the overlap path here.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 57007b9708
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| flat_ids = block_ids.reshape(-1) | ||
| input_embeds = self.embed_tokens(flat_ids) | ||
| input_embeds = self.embed_tokens(flat_ids, reduce_results=False) |
There was a problem hiding this comment.
When DFLASH is run with different attention and dense TP sizes, this skips the VocabParallelEmbedding all-reduce over the embedding/attention TP group and relies on the first draft layer's residual path to reduce the partial embeddings over mapping.dense.tp_group instead. That group can be different (the CLI supports separate --attn-tp-size/--dense-tp-size, and resolve_communication() only disables fusion when they differ), so tokens whose embedding row lives on another attention-TP rank enter the draft model as zeros/partials and produce incorrect draft logits. Keep the embedding reduction here or explicitly reduce over the embedding's TP group before handing the tensor to the dense-parallel draft layer.
Useful? React with 👍 / 👎.
minedec
changed the title
3 participants
Summary
Support DFlash in qwen3.5 and DFlash perf optimize
Support Qwen3.5 native DFLASH : Add native DFLASH integration for Qwen3.5 by capturing hidden states from selected target layers and feeding them into the DFLASH drafter.
Support FA4 backend in decode path : Allow the DFLASH drafter to run with --drafter-attention-backend fa4, and keep the attention handling consistent with the trtllm backend — i.e., process the attention through the decode path, but make every query token visible to all KV tokens, thereby achieving a non-causal attention pattern.
Support FP8 KV cache for DFLASH : Add support for DFLASH draft KV cache in
fp8_e4m3, including fused KVmaterialization paths for the TRTLLM MHA backend (FA4 backend not support yet).
Overlap KV projection with target forward: Introduce incremental KV projection, split the DFLASH FC projection by captured layers and run partial GEMMs on an auxiliary CUDA stream as soon as each hidden state is available, overlapping KV projection with the remaining target forward work.
Fuse DFLASH prepare-decode operations : Add a fused Triton prepare-decode kernel that combines current-token selection, draft sequence length update, block position generation, output cache location computation, and block id update into one launch.
(Fix) Fix rope theta : Fix the RoPE theta config loading issue, which improves the DFLASH acceptance rate.
Test Plan
Test the following server configuration on GB200.
Perf
Test on GB200: compare optimized DFlash against the main branch which pick basic Qwen3.5 DFlash support commits, evaluated on GSM8K with batch size 1. Collect performance metrics on the first 10 samples of the dataset.
Accuracy