Refactor: initialize DeepSeek V4 scratch on allocation

[zhangqi-chen]

Summary

• Size QKV projection scratch buffers from the padded runtime token count instead of a static decode/prefill maximum
• Initialize split-K accumulation and prefill padding buffers through pl.create_tensor, removing their explicit seed scopes
• Initialize SWA dummy compressed-attention metadata during allocation, including the -1 invalid-index sentinel

[coderabbitai]

📝 Walkthrough

Walkthrough

DeepSeek-V4 kernels replace explicit buffer-seeding scopes and loops with allocation-time initialization, remove the static T_MAX sizing dependency from QKV intermediates, and update related task-graph documentation.

Changes

DeepSeek-V4 initialization updates

Layer / File(s)	Summary
hc_pre accumulation initialization models/deepseek/v4-flash/hc_pre.py	mixes_raw is zero-initialized at allocation, and the standalone seed scope and corresponding documented dependency are removed.
Prefill buffer initialization models/deepseek/v4-flash/prefill_attention_swa.py, models/deepseek/v4-flash/prefill_compressor_ratio128.py	Dummy attention metadata and pooled padding buffers use allocation-time initialization instead of explicit core-group initialization loops.
QKV runtime-sized intermediates models/deepseek/v4-flash/qkv_proj_rope.py	QKV intermediates are sized from runtime token dimensions, with split-K accumulators initialized to zero at creation.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• hw-native-sys/pypto-lib#804: Modifies the hc_pre linear split-K scheduling around the shared mixes_raw accumulation.
• hw-native-sys/pypto-lib#811: Updates related hc_pre separate-variant task scheduling and seed dependencies.
• hw-native-sys/pypto-lib#822: Removes a dedicated seed stage for split-K accumulation in another DeepSeek pipeline.

Poem

A bunny sees buffers bloom,
Zeroed neatly in their room.
Seeds hop out of task-flow trails,
Runtime sizing fills the sails.
QKV and prefill gleam—
Clean kernels dance like a dream!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: moving DeepSeek V4 scratch-buffer initialization to allocation time.
Description check	✅ Passed	The description is directly related to the changes and matches the refactor details in the PR.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

models/deepseek/v4-flash/qkv_proj_rope.py (1)

41-41: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Keep the Q KV scratch allocation capacity compile-time fixed.

t_matmul = pl.max(t_dim, MATMUL_T_TILE) makes the first dimension of qr_fp32, qr_i8_matmul, q_proj_i32, and kv_fp32 depend on the runtime number of tokens. Use a compile-time scratch capacity, such as MATMUL_T_TILE, and add a pl.max(t_dim, MATMUL_T_TILE) range check if the dynamic token count cannot exceed it; the matmul kernels can still use the actual t_dim for valid-shape masking.

Also applies to: 154-158, 216-216, 286-290

🤖 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 `@models/deepseek/v4-flash/qkv_proj_rope.py` at line 41, Keep scratch-buffer
capacities in the Q/KV projection paths compile-time fixed by using
MATMUL_T_TILE rather than the runtime t_dim-derived value for qr_fp32,
qr_i8_matmul, q_proj_i32, and kv_fp32 allocations. Add a pl.max(t_dim,
MATMUL_T_TILE) range check where needed to reject unsupported dynamic token
counts, while preserving actual t_dim for matmul masking.

Source: Learnings

🤖 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.

Outside diff comments:
In `@models/deepseek/v4-flash/qkv_proj_rope.py`:
- Line 41: Keep scratch-buffer capacities in the Q/KV projection paths
compile-time fixed by using MATMUL_T_TILE rather than the runtime t_dim-derived
value for qr_fp32, qr_i8_matmul, q_proj_i32, and kv_fp32 allocations. Add a
pl.max(t_dim, MATMUL_T_TILE) range check where needed to reject unsupported
dynamic token counts, while preserving actual t_dim for matmul masking.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 5412f19d-2f85-4382-ab16-3a166f7c8998

📥 Commits

Reviewing files that changed from the base of the PR and between 3f67f0d and 5816cff.

📒 Files selected for processing (4)

• models/deepseek/v4-flash/hc_pre.py
• models/deepseek/v4-flash/prefill_attention_swa.py
• models/deepseek/v4-flash/prefill_compressor_ratio128.py
• models/deepseek/v4-flash/qkv_proj_rope.py