Feat/fast dllm v2 dual cache#44
Merged
Merged
Conversation
📝 WalkthroughWalkthroughThis PR adds a ChangesFast-dLLM v2 strategy and runtime wiring
Sequence Diagram(s)sequenceDiagram
participant FastDLLMV2Scheduler
participant FastDLLMV2Req
participant FastDLLMV2ModelRunner
participant _chunked_prefill_attn_unified_kernel
FastDLLMV2Scheduler->>FastDLLMV2Req: read fdv2_mode during schedule()
FastDLLMV2ModelRunner->>FastDLLMV2Req: prepare slot_mapping and fdv2_cache_only
FastDLLMV2ModelRunner->>_chunked_prefill_attn_unified_kernel: launch with FDV2_CACHE_ONLY
Estimated code review effort🎯 5 (Critical) | ⏱️ ~90+ minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@diffulex/sampler/fast_dllm_v2.py`:
- Around line 62-63: The FDV2 final-commit seeding in fast_dllm_v2.py is
bypassing the normal sampler constraints, so the next token is chosen with a raw
argmax instead of the configured temperature/top-p/top-k and forbidden token
filtering. Update the logic in the final-commit path of the sampler to select
fdv2_pending_next_token_id using the same constrained sampling flow used
elsewhere in the sampler, so the seeded token respects the active sampling
settings before it is written into the next block.
In `@diffulex/strategy/fast_dllm_v2/__init__.py`:
- Around line 3-15: The package root for diffulex.strategy.fast_dllm_v2 is only
re-exporting config and engine symbols, so add the attention helpers to its
public surface as well. Update the __init__.py imports and the __all__ list to
include the appropriate attention module symbols alongside
FastDLLMV2StrategyConfig, FastDLLMV2KVCacheManager, FastDLLMV2ModelRunner,
FastDLLMV2Req, and FastDLLMV2Scheduler. Make sure the new exports are accessible
directly from diffulex.strategy.fast_dllm_v2 without requiring callers to import
from the internal attention submodule.
In `@diffulex/strategy/fast_dllm_v2/attention/metadata.py`:
- Around line 20-56: The metadata setter currently hardcodes the eager/graph
mode state, so FastDLLMV2AttnMetaData can be built with the wrong runtime
configuration. Update set_fast_dllm_v2_attn_metadata to accept the real
enforce_eager value from the caller and pass it through when constructing
FastDLLMV2AttnMetaData instead of forcing False. Make sure any call sites that
populate FDV2 attention metadata thread this flag consistently so the metadata
matches the actual execution mode.
In `@diffulex/strategy/fast_dllm_v2/engine/model_runner.py`:
- Around line 220-222: The CUDA-graph padding logic in model_runner’s replay
path only carries forward cu_seqlens, leaving other padded-row state stale.
Update the padding section around the cu_seqlens_q/cu_seqlens_k loop to also
neutralize valid_slices, status_table, page_tables, and any trailing
slot_mapping entries for rows beyond num_reqs up to captured_num_seqs, so
replayed smaller batches cannot reactivate stale rows in the Triton path.
In `@diffulex/strategy/fast_dllm_v2/engine/request.py`:
- Around line 21-27: The Request.__init__ constructor is using a shared
SamplingParams() object as a default argument, which should be replaced with
None so each call gets a fresh instance. Update the Request initializer to
accept sampling_params as optional, create a new SamplingParams inside the
constructor when it is not provided, and keep the super().__init__(token_ids,
sampling_params) flow unchanged.
In `@diffulex/strategy/fast_dllm_v2/engine/scheduler.py`:
- Around line 23-30: The batching logic in scheduler.py is only grouping
scheduled requests by fdv2_mode in the request-splitting path, which can mix
cache-only and non-cache SUB_BLOCK_REFINE requests in the same batch. Update the
grouping/splitting in the scheduler method that builds kept and deferred so
cache-only and non-cache requests are separated into distinct batches, using the
relevant request mode/cache-only signal from _req_mode or the request object
itself before handing batches to the runner. Keep the change localized around
the scheduled loop that populates kept and deferred so the later any(...) on
fdv2_cache_only cannot flip an entire mixed batch onto the wrong kernel path.
In `@test/python/kernel/test_vllm_attention_perf.py`:
- Around line 341-342: The sliding-window configuration is being read from
data.case even though the function already has a case parameter and
AttentionBenchData does not define case here. Update the SLIDING_WINDOW
assignment in the attention benchmark setup to use the existing case object
directly, keeping the logic in the same test helper that builds the vLLM
attention perf settings.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: ef4c2916-e68a-4c3f-98d7-0b6336cc6b44
📒 Files selected for processing (21)
diffulex/config.pydiffulex/sampler/fast_dllm_v2.pydiffulex/strategy/fast_dllm_v2/__init__.pydiffulex/strategy/fast_dllm_v2/attention/__init__.pydiffulex/strategy/fast_dllm_v2/attention/metadata.pydiffulex/strategy/fast_dllm_v2/config.pydiffulex/strategy/fast_dllm_v2/engine/__init__.pydiffulex/strategy/fast_dllm_v2/engine/kv_cache_manager.pydiffulex/strategy/fast_dllm_v2/engine/model_runner.pydiffulex/strategy/fast_dllm_v2/engine/request.pydiffulex/strategy/fast_dllm_v2/engine/scheduler.pydiffulex_bench/arg_parser.pydiffulex_bench/config.pydiffulex_bench/configs/fast_dllm_v2_gsm8k.ymldiffulex_bench/configs/fast_dllm_v2_multibd_gsm8k.ymldiffulex_bench/main.pydiffulex_kernel/python/chunked_prefill_grouped_triton.pydiffulex_kernel/python/chunked_prefill_triton.pytest/python/engine/test_fast_dllm_v2_strategy.pytest/python/kernel/test_dllm_flash_attn_chunked_prefill_unified_kernel.pytest/python/kernel/test_vllm_attention_perf.py
Comment on lines
+62
to
+63
| if int(getattr(req, "fdv2_mode", -1)) == _FDV2_FINAL_COMMIT: | ||
| req.fdv2_pending_next_token_id = int(torch.argmax(req_logits[-1, ...]).item()) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Apply the normal sampling constraints when seeding the next block.
Line 63 bypasses temperature/top-p/top-k and forbidden_token_ids, then the pending token is written into the next block. This can seed a mask token or make FDV2 final-commit behavior inconsistent with the configured sampler.
Proposed fix
if int(getattr(req, "fdv2_mode", -1)) == _FDV2_FINAL_COMMIT:
- req.fdv2_pending_next_token_id = int(torch.argmax(req_logits[-1, ...]).item())
+ _, next_tokens, _ = self.sample_tokens(
+ req_logits[-1:, ...],
+ temperature_value,
+ top_p=top_p,
+ top_k=top_k,
+ neg_entropy=(neg_entropy == "neg_entropy"),
+ margin_confidence=(margin_confidence == "margin_confidence"),
+ forbidden_token_ids=[int(req.mask_token_id)],
+ )
+ req.fdv2_pending_next_token_id = int(next_tokens[0].item())
true_local_ids_map[req_id_str] = true_local_ids_sub_map📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if int(getattr(req, "fdv2_mode", -1)) == _FDV2_FINAL_COMMIT: | |
| req.fdv2_pending_next_token_id = int(torch.argmax(req_logits[-1, ...]).item()) | |
| if int(getattr(req, "fdv2_mode", -1)) == _FDV2_FINAL_COMMIT: | |
| _, next_tokens, _ = self.sample_tokens( | |
| req_logits[-1:, ...], | |
| temperature_value, | |
| top_p=top_p, | |
| top_k=top_k, | |
| neg_entropy=(neg_entropy == "neg_entropy"), | |
| margin_confidence=(margin_confidence == "margin_confidence"), | |
| forbidden_token_ids=[int(req.mask_token_id)], | |
| ) | |
| req.fdv2_pending_next_token_id = int(next_tokens[0].item()) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@diffulex/sampler/fast_dllm_v2.py` around lines 62 - 63, The FDV2 final-commit
seeding in fast_dllm_v2.py is bypassing the normal sampler constraints, so the
next token is chosen with a raw argmax instead of the configured
temperature/top-p/top-k and forbidden token filtering. Update the logic in the
final-commit path of the sampler to select fdv2_pending_next_token_id using the
same constrained sampling flow used elsewhere in the sampler, so the seeded
token respects the active sampling settings before it is written into the next
block.
Comment on lines
+3
to
+15
| from diffulex.strategy.fast_dllm_v2.config import FastDLLMV2StrategyConfig | ||
| from diffulex.strategy.fast_dllm_v2.engine.kv_cache_manager import FastDLLMV2KVCacheManager | ||
| from diffulex.strategy.fast_dllm_v2.engine.model_runner import FastDLLMV2ModelRunner | ||
| from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Req | ||
| from diffulex.strategy.fast_dllm_v2.engine.scheduler import FastDLLMV2Scheduler | ||
|
|
||
| __all__ = [ | ||
| "FastDLLMV2StrategyConfig", | ||
| "FastDLLMV2KVCacheManager", | ||
| "FastDLLMV2ModelRunner", | ||
| "FastDLLMV2Req", | ||
| "FastDLLMV2Scheduler", | ||
| ] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Re-export the attention helpers from the package root.
diffulex.strategy.fast_dllm_v2 currently exposes config and engine symbols only, so callers still need to know the internal attention submodule layout. That leaves the new public export surface incomplete.
Suggested fix
+from diffulex.strategy.fast_dllm_v2.attention import (
+ FastDLLMV2AttnMetaData,
+ fetch_fast_dllm_v2_attn_metadata,
+ reset_fast_dllm_v2_attn_metadata,
+ set_fast_dllm_v2_attn_metadata,
+)
from diffulex.strategy.fast_dllm_v2.config import FastDLLMV2StrategyConfig
from diffulex.strategy.fast_dllm_v2.engine.kv_cache_manager import FastDLLMV2KVCacheManager
from diffulex.strategy.fast_dllm_v2.engine.model_runner import FastDLLMV2ModelRunner
from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Req
from diffulex.strategy.fast_dllm_v2.engine.scheduler import FastDLLMV2Scheduler
__all__ = [
+ "FastDLLMV2AttnMetaData",
"FastDLLMV2StrategyConfig",
"FastDLLMV2KVCacheManager",
"FastDLLMV2ModelRunner",
"FastDLLMV2Req",
"FastDLLMV2Scheduler",
+ "fetch_fast_dllm_v2_attn_metadata",
+ "reset_fast_dllm_v2_attn_metadata",
+ "set_fast_dllm_v2_attn_metadata",
]As per PR objectives, the Fast-dLLM v2 package should export config, engine, and attention helpers.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| from diffulex.strategy.fast_dllm_v2.config import FastDLLMV2StrategyConfig | |
| from diffulex.strategy.fast_dllm_v2.engine.kv_cache_manager import FastDLLMV2KVCacheManager | |
| from diffulex.strategy.fast_dllm_v2.engine.model_runner import FastDLLMV2ModelRunner | |
| from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Req | |
| from diffulex.strategy.fast_dllm_v2.engine.scheduler import FastDLLMV2Scheduler | |
| __all__ = [ | |
| "FastDLLMV2StrategyConfig", | |
| "FastDLLMV2KVCacheManager", | |
| "FastDLLMV2ModelRunner", | |
| "FastDLLMV2Req", | |
| "FastDLLMV2Scheduler", | |
| ] | |
| from diffulex.strategy.fast_dllm_v2.attention import ( | |
| FastDLLMV2AttnMetaData, | |
| fetch_fast_dllm_v2_attn_metadata, | |
| reset_fast_dllm_v2_attn_metadata, | |
| set_fast_dllm_v2_attn_metadata, | |
| ) | |
| from diffulex.strategy.fast_dllm_v2.config import FastDLLMV2StrategyConfig | |
| from diffulex.strategy.fast_dllm_v2.engine.kv_cache_manager import FastDLLMV2KVCacheManager | |
| from diffulex.strategy.fast_dllm_v2.engine.model_runner import FastDLLMV2ModelRunner | |
| from diffulex.strategy.fast_dllm_v2.engine.request import FastDLLMV2Req | |
| from diffulex.strategy.fast_dllm_v2.engine.scheduler import FastDLLMV2Scheduler | |
| __all__ = [ | |
| "FastDLLMV2AttnMetaData", | |
| "FastDLLMV2StrategyConfig", | |
| "FastDLLMV2KVCacheManager", | |
| "FastDLLMV2ModelRunner", | |
| "FastDLLMV2Req", | |
| "FastDLLMV2Scheduler", | |
| "fetch_fast_dllm_v2_attn_metadata", | |
| "reset_fast_dllm_v2_attn_metadata", | |
| "set_fast_dllm_v2_attn_metadata", | |
| ] |
🧰 Tools
🪛 Ruff (0.15.18)
[warning] 9-15: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@diffulex/strategy/fast_dllm_v2/__init__.py` around lines 3 - 15, The package
root for diffulex.strategy.fast_dllm_v2 is only re-exporting config and engine
symbols, so add the attention helpers to its public surface as well. Update the
__init__.py imports and the __all__ list to include the appropriate attention
module symbols alongside FastDLLMV2StrategyConfig, FastDLLMV2KVCacheManager,
FastDLLMV2ModelRunner, FastDLLMV2Req, and FastDLLMV2Scheduler. Make sure the new
exports are accessible directly from diffulex.strategy.fast_dllm_v2 without
requiring callers to import from the internal attention submodule.
Comment on lines
+20
to
+56
| def set_fast_dllm_v2_attn_metadata( | ||
| is_prefill: list[bool] | bool, | ||
| cu_seqlens_q: torch.Tensor | None = None, | ||
| cu_seqlens_k: torch.Tensor | None = None, | ||
| max_seqlen_q: int = 0, | ||
| max_seqlen_k: int = 0, | ||
| slot_mapping: torch.Tensor | None = None, | ||
| need_kv_cache_store: bool | None = None, | ||
| context_lens: torch.Tensor | None = None, | ||
| page_tables: torch.Tensor | None = None, | ||
| page_size: int = 32, | ||
| block_size: int = 32, | ||
| kv_cache_layout: str = "unified", | ||
| fdv2_cache_only: bool = False, | ||
| fdv2_mode: int = 0, | ||
| ) -> None: | ||
| global _FAST_DLLM_V2_ATTN_METADATA | ||
| has_prefill, all_prefill = infer_prefill_flags(is_prefill) | ||
| _FAST_DLLM_V2_ATTN_METADATA = FastDLLMV2AttnMetaData( | ||
| is_prefill=is_prefill if isinstance(is_prefill, list) else [bool(is_prefill)], | ||
| enforce_eager=False, | ||
| cu_seqlens_q=cu_seqlens_q, | ||
| cu_seqlens_k=cu_seqlens_k, | ||
| max_seqlen_q=max_seqlen_q, | ||
| max_seqlen_k=max_seqlen_k, | ||
| slot_mapping=slot_mapping, | ||
| need_kv_cache_store_static=need_kv_cache_store, | ||
| has_prefill_static=has_prefill, | ||
| all_prefill_static=all_prefill, | ||
| context_lens=context_lens, | ||
| page_tables=page_tables, | ||
| page_size=page_size, | ||
| block_size=block_size, | ||
| kv_cache_layout=kv_cache_layout, | ||
| fdv2_cache_only=bool(fdv2_cache_only), | ||
| fdv2_mode=int(fdv2_mode), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Thread the real enforce_eager flag into this metadata setter.
Line 40 pins enforce_eager to False, so any FDV2 path that runs in eager mode will publish incorrect attention metadata. This helper should accept the runtime value instead of forcing graph mode.
Suggested fix
def set_fast_dllm_v2_attn_metadata(
is_prefill: list[bool] | bool,
+ enforce_eager: bool = False,
cu_seqlens_q: torch.Tensor | None = None,
cu_seqlens_k: torch.Tensor | None = None,
@@
_FAST_DLLM_V2_ATTN_METADATA = FastDLLMV2AttnMetaData(
is_prefill=is_prefill if isinstance(is_prefill, list) else [bool(is_prefill)],
- enforce_eager=False,
+ enforce_eager=bool(enforce_eager),
cu_seqlens_q=cu_seqlens_q,📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def set_fast_dllm_v2_attn_metadata( | |
| is_prefill: list[bool] | bool, | |
| cu_seqlens_q: torch.Tensor | None = None, | |
| cu_seqlens_k: torch.Tensor | None = None, | |
| max_seqlen_q: int = 0, | |
| max_seqlen_k: int = 0, | |
| slot_mapping: torch.Tensor | None = None, | |
| need_kv_cache_store: bool | None = None, | |
| context_lens: torch.Tensor | None = None, | |
| page_tables: torch.Tensor | None = None, | |
| page_size: int = 32, | |
| block_size: int = 32, | |
| kv_cache_layout: str = "unified", | |
| fdv2_cache_only: bool = False, | |
| fdv2_mode: int = 0, | |
| ) -> None: | |
| global _FAST_DLLM_V2_ATTN_METADATA | |
| has_prefill, all_prefill = infer_prefill_flags(is_prefill) | |
| _FAST_DLLM_V2_ATTN_METADATA = FastDLLMV2AttnMetaData( | |
| is_prefill=is_prefill if isinstance(is_prefill, list) else [bool(is_prefill)], | |
| enforce_eager=False, | |
| cu_seqlens_q=cu_seqlens_q, | |
| cu_seqlens_k=cu_seqlens_k, | |
| max_seqlen_q=max_seqlen_q, | |
| max_seqlen_k=max_seqlen_k, | |
| slot_mapping=slot_mapping, | |
| need_kv_cache_store_static=need_kv_cache_store, | |
| has_prefill_static=has_prefill, | |
| all_prefill_static=all_prefill, | |
| context_lens=context_lens, | |
| page_tables=page_tables, | |
| page_size=page_size, | |
| block_size=block_size, | |
| kv_cache_layout=kv_cache_layout, | |
| fdv2_cache_only=bool(fdv2_cache_only), | |
| fdv2_mode=int(fdv2_mode), | |
| ) | |
| def set_fast_dllm_v2_attn_metadata( | |
| is_prefill: list[bool] | bool, | |
| enforce_eager: bool = False, | |
| cu_seqlens_q: torch.Tensor | None = None, | |
| cu_seqlens_k: torch.Tensor | None = None, | |
| max_seqlen_q: int = 0, | |
| max_seqlen_k: int = 0, | |
| slot_mapping: torch.Tensor | None = None, | |
| need_kv_cache_store: bool | None = None, | |
| context_lens: torch.Tensor | None = None, | |
| page_tables: torch.Tensor | None = None, | |
| page_size: int = 32, | |
| block_size: int = 32, | |
| kv_cache_layout: str = "unified", | |
| fdv2_cache_only: bool = False, | |
| fdv2_mode: int = 0, | |
| ) -> None: | |
| global _FAST_DLLM_V2_ATTN_METADATA | |
| has_prefill, all_prefill = infer_prefill_flags(is_prefill) | |
| _FAST_DLLM_V2_ATTN_METADATA = FastDLLMV2AttnMetaData( | |
| is_prefill=is_prefill if isinstance(is_prefill, list) else [bool(is_prefill)], | |
| enforce_eager=bool(enforce_eager), | |
| cu_seqlens_q=cu_seqlens_q, | |
| cu_seqlens_k=cu_seqlens_k, | |
| max_seqlen_q=max_seqlen_q, | |
| max_seqlen_k=max_seqlen_k, | |
| slot_mapping=slot_mapping, | |
| need_kv_cache_store_static=need_kv_cache_store, | |
| has_prefill_static=has_prefill, | |
| all_prefill_static=all_prefill, | |
| context_lens=context_lens, | |
| page_tables=page_tables, | |
| page_size=page_size, | |
| block_size=block_size, | |
| kv_cache_layout=kv_cache_layout, | |
| fdv2_cache_only=bool(fdv2_cache_only), | |
| fdv2_mode=int(fdv2_mode), | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@diffulex/strategy/fast_dllm_v2/attention/metadata.py` around lines 20 - 56,
The metadata setter currently hardcodes the eager/graph mode state, so
FastDLLMV2AttnMetaData can be built with the wrong runtime configuration. Update
set_fast_dllm_v2_attn_metadata to accept the real enforce_eager value from the
caller and pass it through when constructing FastDLLMV2AttnMetaData instead of
forcing False. Make sure any call sites that populate FDV2 attention metadata
thread this flag consistently so the metadata matches the actual execution mode.
Comment on lines
+220
to
+222
| for i in range(num_reqs, captured_num_seqs): | ||
| graph_vars["cu_seqlens_q"][i + 1] = graph_vars["cu_seqlens_q"][i] | ||
| graph_vars["cu_seqlens_k"][i + 1] = graph_vars["cu_seqlens_k"][i] |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Neutralize all padded CUDA-graph rows, not just cu_seqlens.
Line 220 pads only the cumulative sequence lengths. After replaying a larger batch, stale valid_slices, status_table, page_tables, and trailing slot_mapping can remain for padded rows; the Triton path computes validity from valid_slice - q_start, so those rows may become active again and use stale KV slots.
Proposed fix
+ if num_tokens < captured_num_tokens:
+ graph_vars["input_ids"][num_tokens:captured_num_tokens].zero_()
+ graph_vars["positions"][num_tokens:captured_num_tokens].zero_()
+ graph_vars["slot_mapping"][num_tokens:captured_num_tokens].fill_(-1)
+
for i in range(num_reqs, captured_num_seqs):
graph_vars["cu_seqlens_q"][i + 1] = graph_vars["cu_seqlens_q"][i]
graph_vars["cu_seqlens_k"][i + 1] = graph_vars["cu_seqlens_k"][i]
+ graph_vars["valid_slices"][i] = graph_vars["cu_seqlens_q"][i]
+ graph_vars["context_lens"][i] = 0
+ graph_vars["status_table"][i] = 0
+ graph_vars["prefix_lens"][i] = 0
+ graph_vars["padded_prefix_lens"][i] = 0
+ graph_vars["page_tables"][i].fill_(-1)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| for i in range(num_reqs, captured_num_seqs): | |
| graph_vars["cu_seqlens_q"][i + 1] = graph_vars["cu_seqlens_q"][i] | |
| graph_vars["cu_seqlens_k"][i + 1] = graph_vars["cu_seqlens_k"][i] | |
| if num_tokens < captured_num_tokens: | |
| graph_vars["input_ids"][num_tokens:captured_num_tokens].zero_() | |
| graph_vars["positions"][num_tokens:captured_num_tokens].zero_() | |
| graph_vars["slot_mapping"][num_tokens:captured_num_tokens].fill_(-1) | |
| for i in range(num_reqs, captured_num_seqs): | |
| graph_vars["cu_seqlens_q"][i + 1] = graph_vars["cu_seqlens_q"][i] | |
| graph_vars["cu_seqlens_k"][i + 1] = graph_vars["cu_seqlens_k"][i] | |
| graph_vars["valid_slices"][i] = graph_vars["cu_seqlens_q"][i] | |
| graph_vars["context_lens"][i] = 0 | |
| graph_vars["status_table"][i] = 0 | |
| graph_vars["prefix_lens"][i] = 0 | |
| graph_vars["padded_prefix_lens"][i] = 0 | |
| graph_vars["page_tables"][i].fill_(-1) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@diffulex/strategy/fast_dllm_v2/engine/model_runner.py` around lines 220 -
222, The CUDA-graph padding logic in model_runner’s replay path only carries
forward cu_seqlens, leaving other padded-row state stale. Update the padding
section around the cu_seqlens_q/cu_seqlens_k loop to also neutralize
valid_slices, status_table, page_tables, and any trailing slot_mapping entries
for rows beyond num_reqs up to captured_num_seqs, so replayed smaller batches
cannot reactivate stale rows in the Triton path.
Comment on lines
+21
to
+27
| def __init__( | ||
| self, | ||
| token_ids: list[int], | ||
| sampling_params: SamplingParams = SamplingParams(), | ||
| config: Config | None = None, | ||
| ): | ||
| super().__init__(token_ids, sampling_params) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor
🧩 Analysis chain
🏁 Script executed:
fd -g "*sampling*" SamplingParams | head -5Repository: SJTU-DENG-Lab/Diffulex
Length of output: 262
🏁 Script executed:
grep -r "class SamplingParams" --include="*.py"Repository: SJTU-DENG-Lab/Diffulex
Length of output: 210
🏁 Script executed:
cat -n diffulex/sampling_params.py | head -100Repository: SJTU-DENG-Lab/Diffulex
Length of output: 964
🏁 Script executed:
grep -r "\.sampling_params\." --include="*.py" | head -20Repository: SJTU-DENG-Lab/Diffulex
Length of output: 942
🏁 Script executed:
grep -r "\.sampling_params\.max_tokens\s*=" --include="*.py"Repository: SJTU-DENG-Lab/Diffulex
Length of output: 458
🏁 Script executed:
grep -rE "\.sampling_params\s*=\s*" --include="*.py" diffulex/Repository: SJTU-DENG-Lab/Diffulex
Length of output: 160
🏁 Script executed:
grep -rE "sampling_params\[" --include="*.py" diffulex/Repository: SJTU-DENG-Lab/Diffulex
Length of output: 160
🏁 Script executed:
grep -r "sampling_params\." --include="*.py" diffulex/ | grep "="Repository: SJTU-DENG-Lab/Diffulex
Length of output: 572
Avoid using a mutable default argument for SamplingParams.
Line 24 defaults sampling_params to a single SamplingParams() instance. While the current implementation copies fields to instance variables rather than mutating the shared object, this pattern risks unintended side effects if the logic changes or if the class definition evolves to include mutable state.
Use None as the default to create a fresh instance per call.
Proposed fix
def __init__(
self,
token_ids: list[int],
- sampling_params: SamplingParams = SamplingParams(),
+ sampling_params: SamplingParams | None = None,
config: Config | None = None,
):
+ if sampling_params is None:
+ sampling_params = SamplingParams()
super().__init__(token_ids, sampling_params)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def __init__( | |
| self, | |
| token_ids: list[int], | |
| sampling_params: SamplingParams = SamplingParams(), | |
| config: Config | None = None, | |
| ): | |
| super().__init__(token_ids, sampling_params) | |
| def __init__( | |
| self, | |
| token_ids: list[int], | |
| sampling_params: SamplingParams | None = None, | |
| config: Config | None = None, | |
| ): | |
| if sampling_params is None: | |
| sampling_params = SamplingParams() | |
| super().__init__(token_ids, sampling_params) |
🧰 Tools
🪛 Ruff (0.15.18)
[warning] 24-24: Do not perform function call SamplingParams in argument defaults; instead, perform the call within the function, or read the default from a module-level singleton variable
(B008)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@diffulex/strategy/fast_dllm_v2/engine/request.py` around lines 21 - 27, The
Request.__init__ constructor is using a shared SamplingParams() object as a
default argument, which should be replaced with None so each call gets a fresh
instance. Update the Request initializer to accept sampling_params as optional,
create a new SamplingParams inside the constructor when it is not provided, and
keep the super().__init__(token_ids, sampling_params) flow unchanged.
Source: Linters/SAST tools
Comment on lines
+23
to
+30
| target_mode = self._req_mode(scheduled[0]) | ||
| kept = [] | ||
| deferred = [] | ||
| for req in scheduled: | ||
| if self._req_mode(req) == target_mode: | ||
| kept.append(req) | ||
| else: | ||
| deferred.append(req) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep cache-only and non-cache sub-block requests in separate batches.
Line 23 groups only by fdv2_mode, but the model runner later turns fdv2_cache_only on for the whole batch with any(...). A mixed SUB_BLOCK_REFINE batch can therefore force non-cache requests through the cache-only kernel path and skip their current-buffer K/V attention.
Proposed fix
- def _req_mode(req) -> FastDLLMV2Mode:
- return getattr(req, "fdv2_mode", FastDLLMV2Mode.FULL_BUFFER_INIT)
+ def _req_batch_key(req) -> tuple[FastDLLMV2Mode, bool]:
+ mode = getattr(req, "fdv2_mode", FastDLLMV2Mode.FULL_BUFFER_INIT)
+ use_block_cache = True
+ if mode == FastDLLMV2Mode.SUB_BLOCK_REFINE:
+ use_block_cache = bool(getattr(req, "fdv2_use_block_cache", True))
+ return mode, use_block_cache
@@
- target_mode = self._req_mode(scheduled[0])
+ target_key = self._req_batch_key(scheduled[0])
@@
- if self._req_mode(req) == target_mode:
+ if self._req_batch_key(req) == target_key:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| target_mode = self._req_mode(scheduled[0]) | |
| kept = [] | |
| deferred = [] | |
| for req in scheduled: | |
| if self._req_mode(req) == target_mode: | |
| kept.append(req) | |
| else: | |
| deferred.append(req) | |
| target_key = self._req_batch_key(scheduled[0]) | |
| kept = [] | |
| deferred = [] | |
| for req in scheduled: | |
| if self._req_batch_key(req) == target_key: | |
| kept.append(req) | |
| else: | |
| deferred.append(req) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@diffulex/strategy/fast_dllm_v2/engine/scheduler.py` around lines 23 - 30, The
batching logic in scheduler.py is only grouping scheduled requests by fdv2_mode
in the request-splitting path, which can mix cache-only and non-cache
SUB_BLOCK_REFINE requests in the same batch. Update the grouping/splitting in
the scheduler method that builds kept and deferred so cache-only and non-cache
requests are separated into distinct batches, using the relevant request
mode/cache-only signal from _req_mode or the request object itself before
handing batches to the runner. Keep the change localized around the scheduled
loop that populates kept and deferred so the later any(...) on fdv2_cache_only
cannot flip an entire mixed batch onto the wrong kernel path.
Comment on lines
+341
to
+342
| FDV2_CACHE_ONLY=False, | ||
| SLIDING_WINDOW=int(data.case.sliding_window or 0), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Use the case parameter for sliding-window settings.
Line 342 reads data.case.sliding_window, but AttentionBenchData does not define case in the shown fields and this function already receives case.
Proposed fix
PREFIX_CAUSAL=bool(getattr(metadata, "prefix_causal", False)),
FDV2_CACHE_ONLY=False,
- SLIDING_WINDOW=int(data.case.sliding_window or 0),
+ SLIDING_WINDOW=int(getattr(case, "sliding_window", 0) or 0),
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| FDV2_CACHE_ONLY=False, | |
| SLIDING_WINDOW=int(data.case.sliding_window or 0), | |
| FDV2_CACHE_ONLY=False, | |
| SLIDING_WINDOW=int(getattr(case, "sliding_window", 0) or 0), |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/kernel/test_vllm_attention_perf.py` around lines 341 - 342, The
sliding-window configuration is being read from data.case even though the
function already has a case parameter and AttentionBenchData does not define
case here. Update the SLIDING_WINDOW assignment in the attention benchmark setup
to use the existing case object directly, keeping the logic in the same test
helper that builds the vLLM attention perf settings.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR integrates Fast-dLLM v2 dual-cache decoding into Diffulex.
Main changes:
fast_dllm_v2decoding strategy with FDv2-style block/sub-block execution.multi_bdavailable as a separate Diffulex baseline for Fast-dLLM v2.Evaluation setup
Model:
Efficient-Large-Model/Fast_dLLM_v2_7BMax generation length:
max_new_tokens=2048Precision: BF16
Decoding: greedy / deterministic
Metric: Diffulex lm-eval exact match unless otherwise noted
Hardware: single Nvidia A100 per run
Main runner settings:
Diffulex Dual Cache:decoding_strategy=fast_dllm_v2,block_size=8,buffer_size=4,accept_threshold=1.0,use_block_cache=true,max_num_reqs=32Diffulex MultiBD:decoding_strategy=multi_bd,block_size=8,buffer_size=4,accept_threshold=0.95,token_stability_threshold=0.0,max_num_reqs=32FDv2 Native: original Fast-dLLM v2 repo,use_block_cache=true,bd_size=32,small_block_size=8,threshold=1,batch_size=32Main results
GSM8K MultiBD speed/quality note
The main MultiBD row uses a more conservative
block_size=8, buffer_size=4, accept_threshold=0.95setting for cross-task quality. A speed-oriented MultiBD setting gives much higher throughput on GSM8K, with lower full-set accuracy:block_size=8, buffer_size=4, threshold=0.95block_size=32, buffer_size=2, threshold=0.90block_size=32, buffer_size=2, threshold=0.90, limit=128Task notes
gsm8k_diffulex, full 1319 samples unless noted.humaneval_diffulex, full 164 samples.mbpp_diffulex, full 500 samples.TPFand throughput columns are reported from Diffulex benchmark statistics for Diffulex runs and from the original Fast-dLLM v2 benchmark output for native runs.Summary by CodeRabbit
New Features
fast_dllm_v2decoding strategy across configuration, runtime scheduling, sampling, request handling, and benchmark options.Bug Fixes