DFlash split 2: drafter module + tap APIs - #206
Conversation
|
Warning Review limit reached
More reviews will be available in 50 minutes and 30 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughAdds a new ChangesDFlash Speculative Decoding Drafter
CI Runner Pin
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested Reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 3
🤖 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 `@crates/higgs-models/src/dflash.rs`:
- Around line 34-53: Add rustdoc comments for the new public API surface in
DFlashConfig and DFlashDrafter::config, including the struct itself and each
public field introduced in dflash.rs. Use the existing identifiers DFlashConfig,
its fields, and DFlashDrafter::config to locate the changes, and document any
new user-facing behavior so the public config surface is covered consistently
with the Rust docs guidelines.
- Around line 457-459: In DFlash::forward, validate that the cache length
exactly matches self.layers before iterating, since zipping
self.layers.iter_mut() with cache.iter_mut() can silently skip extra decoder
layers or ignore surplus cache entries. Add an explicit length check near the
loop and return an error if the sizes differ, then keep the existing forward
pass logic unchanged once the lengths are confirmed.
- Around line 671-724: In `loads_modal_drafter_against_real_weights` and the
surrounding cache assertions, remove the clippy blockers by replacing the
`as`-based size conversion in `zeros` with a checked integer conversion,
avoiding direct indexing into `cache` by using safe accessors or pattern
matching for the two entries, and eliminating `expect()` in favor of explicit
error handling or assertions. Update the affected test helpers in `dflash.rs` so
the same symbols (`zeros`, `cache`, `load_dflash_drafter`) still drive the
checks without triggering clippy.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 1067480d-8eb5-4975-aaac-6d7fceaabb3a
📒 Files selected for processing (4)
.github/workflows/ci.ymlcrates/higgs-models/src/dflash.rscrates/higgs-models/src/lib.rscrates/higgs-models/src/qwen3_next.rs
| pub struct DFlashConfig { | ||
| pub hidden_size: i32, | ||
| pub num_hidden_layers: i32, | ||
| pub num_attention_heads: i32, | ||
| pub num_key_value_heads: i32, | ||
| #[serde(default = "default_head_dim")] | ||
| pub head_dim: i32, | ||
| pub intermediate_size: i32, | ||
| #[serde(default = "default_rms_norm_eps")] | ||
| pub rms_norm_eps: f32, | ||
| #[serde(default = "default_rope_theta")] | ||
| pub rope_theta: f32, | ||
| #[serde(default = "default_block_size")] | ||
| pub block_size: i32, | ||
| pub vocab_size: i32, | ||
| #[serde(default)] | ||
| pub layer_types: Option<Vec<String>>, | ||
| #[serde(default)] | ||
| pub sliding_window: Option<i32>, | ||
| dflash_config: DFlashSubConfig, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add rustdoc for the public config surface.
DFlashConfig, its public fields, and DFlashDrafter::config are new public API. As per coding guidelines, "**/*.rs: Add doc comments on public structs/fields in Rust when changing user-facing behavior".
Also applies to: 371-385
🤖 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 `@crates/higgs-models/src/dflash.rs` around lines 34 - 53, Add rustdoc comments
for the new public API surface in DFlashConfig and DFlashDrafter::config,
including the struct itself and each public field introduced in dflash.rs. Use
the existing identifiers DFlashConfig, its fields, and DFlashDrafter::config to
locate the changes, and document any new user-facing behavior so the public
config surface is covered consistently with the Rust docs guidelines.
Source: Coding guidelines
| let mut h = noise.clone(); | ||
| for (layer, lc) in self.layers.iter_mut().zip(cache.iter_mut()) { | ||
| h = layer.forward(&h, &target_hidden, lc, cache_offset)?; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate cache length before zipping.
A short cache silently skips decoder layers; a long cache is partially ignored.
Proposed fix
if taps.len() != self.config.num_taps() {
return Err(Exception::custom(format!(
"expected {} taps, got {}",
self.config.num_taps(),
taps.len()
)));
}
+
+ if cache.len() != self.layers.len() {
+ return Err(Exception::custom(format!(
+ "expected {} cache layers, got {}",
+ self.layers.len(),
+ cache.len()
+ )));
+ }
// Cache offset = max cached seq length (0 on first round)📝 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.
| let mut h = noise.clone(); | |
| for (layer, lc) in self.layers.iter_mut().zip(cache.iter_mut()) { | |
| h = layer.forward(&h, &target_hidden, lc, cache_offset)?; | |
| if taps.len() != self.config.num_taps() { | |
| return Err(Exception::custom(format!( | |
| "expected {} taps, got {}", | |
| self.config.num_taps(), | |
| taps.len() | |
| ))); | |
| } | |
| if cache.len() != self.layers.len() { | |
| return Err(Exception::custom(format!( | |
| "expected {} cache layers, got {}", | |
| self.layers.len(), | |
| cache.len() | |
| ))); | |
| } | |
| let mut h = noise.clone(); | |
| for (layer, lc) in self.layers.iter_mut().zip(cache.iter_mut()) { | |
| h = layer.forward(&h, &target_hidden, lc, cache_offset)?; | |
| } |
🤖 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 `@crates/higgs-models/src/dflash.rs` around lines 457 - 459, In
DFlash::forward, validate that the cache length exactly matches self.layers
before iterating, since zipping self.layers.iter_mut() with cache.iter_mut() can
silently skip extra decoder layers or ignore surplus cache entries. Add an
explicit length check near the loop and return an error if the sizes differ,
then keep the existing forward pass logic unchanged once the lengths are
confirmed.
f515b86 to
1237cfa
Compare
1237cfa to
b12e1ed
Compare
1 participant
Second split of the DFlash stack (from #204): the block-diffusion drafter module — drafter architecture, config parsing, safetensors loading hook, cache-crop helper, and
accept_prefixtests.Stacked on #205 (split-1). GitHub cross-repo PRs can only target
main, so until #205 merges, this diff is cumulative (it shows split-1 + split-2). Review the top commit (f515b868) only — the diff narrows to just the drafter once #205 lands.Stack order: #205 (foundation) → this (drafter) → engine draft-verify loop → config/doctor/docs → per-request selection + streaming.
Validation
cargo clippy+cargo fmt --check— cleancargo test -p higgs-models(drafter +accept_prefix) — pass🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Chores
macos-15GitHub-hosted runner for improved consistency.