Merge DSV4 compressor ratio modules#730
Conversation
📝 WalkthroughWalkthroughThis PR consolidates DeepSeek-V4 KV compression logic into two new modules (compressor_ratio4.py, compressor_ratio128.py) with shared write-schedule helpers, replacing prior separate decode/prefill compressor files. Downstream attention, indexer, prefill_fwd, and prefill_layer modules are updated to token-major layouts, shared RoPE frequency tables, and consolidated batch-indexed compress_state/block-table contracts. ChangesDeepSeek-V4 Compressor Consolidation
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the DeepSeek-V4 compressor modules by introducing unified ratio-128 and ratio-4 compressor implementations (compressor_ratio128.py and compressor_ratio4.py) along with write-schedule helpers (compressor_schedule.py). It updates various attention and prefill modules to utilize these new token-major interfaces directly, removing obsolete decode-specific compressor files. Feedback on these changes includes removing an unused duplicate function prefill_compressor_ratio128 in compressor_ratio4.py that references undefined variables, fixing a NameError in decode_indexer.py where x_flat is referenced instead of x, and simplifying several manual modulo calculations across the codebase to use the standard % operator.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| return ( | ||
| torch.zeros(HEAD_DIM, dtype=torch.float32, device=x.device), | ||
| torch.full((HEAD_DIM,), float("-inf"), dtype=torch.float32, device=x.device), | ||
| ) | ||
| intra = abs_pos % DECODE_COMPRESS_STATE_BLOCK_SIZE | ||
| return ( | ||
| compress_state[blk_id, intra, HEAD_DIM:OUT_DIM], | ||
| compress_state[blk_id, intra, OUT_DIM + HEAD_DIM:], | ||
| ) | ||
|
|
||
| for b in range(bsz): | ||
| first_pos = int(position_ids[b, 0].item()) | ||
| pre_tokens = min(DECODE_S, ratio - (first_pos % ratio)) | ||
| boundary_s = ratio - 1 - (first_pos % ratio) | ||
| should_compress = 0 <= boundary_s < DECODE_S | ||
| boundary_end = first_pos + pre_tokens - 1 | ||
| cur_window_start = boundary_end - ratio + 1 | ||
| prev_window_start = cur_window_start - ratio | ||
|
|
||
| # Per-token ape add + state scatter through explicit token-major slots. | ||
| for s in range(DECODE_S): | ||
| pos = int(position_ids[b, s].item()) | ||
| token_ape_row = pos % ratio | ||
| score[b, s, :] = score[b, s, :] + ape[token_ape_row] | ||
| state_row = int(state_slot_mapping[b, s].item()) | ||
| if state_row >= 0: | ||
| blk_id = state_row // DECODE_COMPRESS_STATE_BLOCK_SIZE | ||
| intra = state_row % DECODE_COMPRESS_STATE_BLOCK_SIZE | ||
| compress_state[blk_id, intra, :OUT_DIM] = kv[b, s, :] | ||
| compress_state[blk_id, intra, OUT_DIM:] = score[b, s, :] | ||
|
|
||
| if should_compress: | ||
| should_compress_rows[b] = True | ||
| kv_rows = [] | ||
| score_rows = [] | ||
| for s in range(ratio): | ||
| abs_pos = prev_window_start + s | ||
| if abs_pos < 0: | ||
| kv_rows.append(torch.zeros(HEAD_DIM, dtype=torch.float32, device=x.device)) | ||
| score_rows.append(torch.full((HEAD_DIM,), float("-inf"), dtype=torch.float32, device=x.device)) | ||
| continue | ||
| kv_row, score_row = read_front_state(b, abs_pos) | ||
| kv_rows.append(kv_row) | ||
| score_rows.append(score_row) | ||
| for s in range(ratio): | ||
| abs_pos = cur_window_start + s | ||
| kv_row, score_row = read_back_state(b, abs_pos) | ||
| kv_rows.append(kv_row) | ||
| score_rows.append(score_row) | ||
| kvs = torch.stack(kv_rows, dim=0).unsqueeze(0) | ||
| scs = torch.stack(score_rows, dim=0).unsqueeze(0) | ||
| pooled[b : b + 1] = (kvs * scs.softmax(dim=1)).sum(dim=1, keepdim=True) | ||
|
|
||
| tensors["compress_state"][:] = compress_state | ||
|
|
||
| if not bool(should_compress_rows.any()): | ||
| return | ||
|
|
||
| def rmsnorm(x, w): | ||
| x = x.float() | ||
| var = x.square().mean(-1, keepdim=True) | ||
| x = x * torch.rsqrt(var + EPS) | ||
| return w * x | ||
|
|
||
| for b in range(bsz): | ||
| if not bool(should_compress_rows[b]): | ||
| continue | ||
| first_pos = int(position_ids[b, 0].item()) | ||
| boundary_s = ratio - 1 - (first_pos % ratio) | ||
| kv_b = rmsnorm(pooled[b : b + 1], norm_w) | ||
|
|
||
| x_pair = kv_b[..., -rd:].unflatten(-1, (-1, 2)) | ||
| x0, x1 = x_pair[..., 0], x_pair[..., 1] | ||
| # cos/sin from the shared freqs table at the compression-window origin, matching | ||
| # the in-kernel gather (window_start = first_pos - first_pos % ratio). freqs_* is | ||
| # BF16; .float() replicates the kernel's BF16->FP32 cast. | ||
| window_start_b = first_pos - (first_pos % ratio) | ||
| cos_v = freqs_cos[window_start_b, : rd // 2].float().view(-1) | ||
| sin_v = freqs_sin[window_start_b, : rd // 2].float().view(-1) | ||
| y0 = x0 * cos_v - x1 * sin_v | ||
| y1 = x0 * sin_v + x1 * cos_v | ||
|
|
||
| kv_b = torch.cat([kv_b[..., :-rd], torch.stack([y0, y1], dim=-1).flatten(-2)], dim=-1) | ||
|
|
||
| cmp_row = int(cmp_slot_mapping[b, boundary_s].item()) | ||
| if cmp_row >= 0: | ||
| # Kernel writes the committed pooled result to the sequence's first | ||
| # token row (t = b * DECODE_S); non-boundary batches and other token rows | ||
| # stay at the output tensor's zero-init. | ||
| tensors["kv"][b * DECODE_S : b * DECODE_S + 1, :] = kv_b.reshape(1, HEAD_DIM) | ||
| blk_id = cmp_row // BLOCK_SIZE | ||
| cmp_kv_cache[blk_id, cmp_row % BLOCK_SIZE, 0] = kv_b[0, 0] | ||
|
|
||
| tensors["cmp_kv_cache"][:] = cmp_kv_cache | ||
|
|
||
|
|
||
| def build_decode_tensor_specs(start_pos=None): | ||
| import torch # type: ignore[import] | ||
| from decode_metadata import ( | ||
| block_table, | ||
| compressed_slot_mapping, | ||
| csa_decode_start_set, | ||
| position_ids_from_starts, | ||
| resolve_start_positions, | ||
| state_slot_mapping, | ||
| ) | ||
| from golden import TensorSpec | ||
| from rope_tables import build_deepseek_v4_rope_tables | ||
|
|
||
| shared_freqs_cos, shared_freqs_sin = build_deepseek_v4_rope_tables(M, COMPRESS_RATIO, dtype=torch.bfloat16) | ||
|
|
||
| def init_x(): | ||
| return torch.rand(DECODE_B, DECODE_S, D).reshape(DECODE_T, D) | ||
| def init_compress_state(): |
There was a problem hiding this comment.
The function prefill_compressor_ratio128 is a leftover duplicate from compressor_ratio128.py and is completely unused in this ratio-4 module. Furthermore, it references undefined variables HCA_C128_RMS_PAD_ROWS and ROPE_HALF which will break JIT compilation at import/compile time. This function should be removed entirely.
| inner_kv_flat = pl.reshape(inner_kv, [T, INNER_HEAD_DIM]) | ||
| idx_slot_mapping_flat = pl.reshape(idx_slot_mapping, [T]) | ||
| inner_state_slot_mapping_flat = pl.reshape(inner_state_slot_mapping, [T]) | ||
| indexer_compressor( |
There was a problem hiding this comment.
The variable x_flat is not defined in this scope. Since the input parameter x has already been flattened to shape [T, D], you should pass x directly to indexer_compressor instead of x_flat to avoid a runtime NameError.
x, inner_kv_flat,| prev_state_block = pl.cast(prev_abs // RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | ||
| prev_state_intra = pl.cast(prev_abs - prev_state_block * RATIO4_STATE_BLOCK_SIZE, pl.INDEX) |
There was a problem hiding this comment.
The modulo operation can be simplified to use the standard % operator instead of performing manual subtraction and multiplication with mixed types. This improves readability and maintainability.
| prev_state_block = pl.cast(prev_abs // RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | |
| prev_state_intra = pl.cast(prev_abs - prev_state_block * RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | |
| prev_state_block = pl.cast(prev_abs // RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | |
| prev_state_intra = pl.cast(prev_abs % RATIO4_STATE_BLOCK_SIZE, pl.INDEX) |
| cur_state_block = pl.cast(cur_abs // RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | ||
| cur_state_intra = pl.cast(cur_abs - cur_state_block * RATIO4_STATE_BLOCK_SIZE, pl.INDEX) |
There was a problem hiding this comment.
The modulo operation can be simplified to use the standard % operator instead of performing manual subtraction and multiplication with mixed types. This improves readability and maintainability.
| cur_state_block = pl.cast(cur_abs // RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | |
| cur_state_intra = pl.cast(cur_abs - cur_state_block * RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | |
| cur_state_block = pl.cast(cur_abs // RATIO4_STATE_BLOCK_SIZE, pl.INDEX) | |
| cur_state_intra = pl.cast(cur_abs % RATIO4_STATE_BLOCK_SIZE, pl.INDEX) |
| prev_state_block = pl.cast(prev_abs // INNER_STATE_BLOCK_SIZE, pl.INDEX) | ||
| prev_state_intra = pl.cast(prev_abs - prev_state_block * INNER_STATE_BLOCK_SIZE, pl.INDEX) |
There was a problem hiding this comment.
The modulo operation can be simplified to use the standard % operator instead of performing manual subtraction and multiplication with mixed types. This improves readability and maintainability.
| prev_state_block = pl.cast(prev_abs // INNER_STATE_BLOCK_SIZE, pl.INDEX) | |
| prev_state_intra = pl.cast(prev_abs - prev_state_block * INNER_STATE_BLOCK_SIZE, pl.INDEX) | |
| prev_state_block = pl.cast(prev_abs // INNER_STATE_BLOCK_SIZE, pl.INDEX) | |
| prev_state_intra = pl.cast(prev_abs % INNER_STATE_BLOCK_SIZE, pl.INDEX) |
| cur_state_block = pl.cast(cur_abs // INNER_STATE_BLOCK_SIZE, pl.INDEX) | ||
| cur_state_intra = pl.cast(cur_abs - cur_state_block * INNER_STATE_BLOCK_SIZE, pl.INDEX) |
There was a problem hiding this comment.
The modulo operation can be simplified to use the standard % operator instead of performing manual subtraction and multiplication with mixed types. This improves readability and maintainability.
| cur_state_block = pl.cast(cur_abs // INNER_STATE_BLOCK_SIZE, pl.INDEX) | |
| cur_state_intra = pl.cast(cur_abs - cur_state_block * INNER_STATE_BLOCK_SIZE, pl.INDEX) | |
| cur_state_block = pl.cast(cur_abs // INNER_STATE_BLOCK_SIZE, pl.INDEX) | |
| cur_state_intra = pl.cast(cur_abs % INNER_STATE_BLOCK_SIZE, pl.INDEX) |
![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
🧹 Nitpick comments (1)
models/deepseek/v4/compressor_ratio4.py (1)
24-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the shared RMS/RoPE helper out of the ratio-128 module.
Importing
compressor_rmsnorm_ropefromcompressor_ratio128.pymakes ratio-4 import execute ratio-128 module-level config checks and constants. Put this ratio-agnostic helper incompressor_schedule.pyor a small common module to avoid ratio-specific import side effects.🤖 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/compressor_ratio4.py` at line 24, The shared RMS/RoPE helper is being imported from the ratio-128 module, which pulls in ratio-specific module initialization and config checks unnecessarily. Move the ratio-agnostic `compressor_rmsnorm_rope` helper into a common location such as `compressor_schedule.py` or a dedicated shared module, then update `compressor_ratio4` to import it from there so `compressor_ratio128` is no longer used as a dependency for shared functionality.
🤖 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 `@models/deepseek/v4/compressor_ratio4.py`:
- Around line 320-335: The kv_and_cache_write SPMD block is missing an explicit
dependency on the RMS/RoPE producer, so the final write can start before
normed_kv is ready. Update the with pl.spmd(..., name_hint="kv_and_cache_write")
block to include a deps argument that depends on the
compressor_rmsnorm_rope-producing step, using the existing
compressor_rmsnorm_rope and normed_kv flow as the reference point. Keep the rest
of the write logic unchanged; only wire the scheduling dependency so normed_kv
is not consumed early.
In `@models/deepseek/v4/decode_indexer.py`:
- Around line 121-123: `position_ids` is already flattened to a 1D token tensor,
but the remaining reads in `decode_indexer.py` still treat it as 2D with
batch/token indices. Update the code paths around the `position_ids` accesses in
the decode indexing logic to use the flattened token index already available in
scope instead of `[b, s]` or `[batch_idx, token_s]`, and keep the
visibility-bound computation aligned with the 1D rank.
In `@models/deepseek/v4/prefill_layer.py`:
- Around line 308-329: The request-building code in prefill_layer.py references
B in the slice calls and block-count logic, but B is not defined in this module.
Fix this by importing the same batch-size alias used in prefill_fwd.py,
specifically PREFILL_BATCH as B, so the symbols in the request loop and
_req_block_count resolve correctly.
---
Nitpick comments:
In `@models/deepseek/v4/compressor_ratio4.py`:
- Line 24: The shared RMS/RoPE helper is being imported from the ratio-128
module, which pulls in ratio-specific module initialization and config checks
unnecessarily. Move the ratio-agnostic `compressor_rmsnorm_rope` helper into a
common location such as `compressor_schedule.py` or a dedicated shared module,
then update `compressor_ratio4` to import it from there so `compressor_ratio128`
is no longer used as a dependency for shared functionality.
🪄 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: 2f2848a1-abe6-48c8-8623-ab97fb4803b2
📒 Files selected for processing (19)
models/deepseek/v4/compressor_ratio128.pymodels/deepseek/v4/compressor_ratio4.pymodels/deepseek/v4/compressor_schedule.pymodels/deepseek/v4/decode_attention_csa.pymodels/deepseek/v4/decode_attention_hca.pymodels/deepseek/v4/decode_compressor_ratio128.pymodels/deepseek/v4/decode_compressor_ratio4.pymodels/deepseek/v4/decode_indexer.pymodels/deepseek/v4/decode_indexer_compressor.pymodels/deepseek/v4/prefill_attention_csa.pymodels/deepseek/v4/prefill_attention_hca.pymodels/deepseek/v4/prefill_attention_swa.pymodels/deepseek/v4/prefill_compressor_ratio128.pymodels/deepseek/v4/prefill_compressor_ratio4.pymodels/deepseek/v4/prefill_fwd.pymodels/deepseek/v4/prefill_indexer.pymodels/deepseek/v4/prefill_indexer_compressor.pymodels/deepseek/v4/prefill_layer.pymodels/deepseek/v4/prefill_sparse_attn.py
💤 Files with no reviewable changes (4)
- models/deepseek/v4/decode_compressor_ratio128.py
- models/deepseek/v4/decode_compressor_ratio4.py
- models/deepseek/v4/prefill_compressor_ratio128.py
- models/deepseek/v4/prefill_compressor_ratio4.py
| compressor_rmsnorm_rope(pooled_kv, norm_w, cos_b, sin_b, normed_kv) | ||
|
|
||
| with pl.spmd(DECODE_B // DECODE_RMS_TILE, name_hint="kv_and_cache_write") as _write_tid: | ||
| batch_base_idx = pl.tile.get_block_idx() | ||
| pad_base = batch_base_idx * DECODE_RMS_PAD_TILE | ||
| for inner in pl.range(DECODE_RMS_TILE): | ||
| pad_row = pad_base + inner | ||
| cache_row_raw = pl.read(write_dst_map, [0, pad_row]) | ||
| if cache_row_raw >= 0: | ||
| kv_row_raw = pl.read(kv_out_row_map, [0, pad_row]) | ||
| if kv_row_raw >= 0: | ||
| kv_row = pl.cast(kv_row_raw, pl.INDEX) | ||
| cache_row = pl.cast(cache_row_raw, pl.INDEX) | ||
| kv_row_fp32 = normed_kv[pad_row : pad_row + 1, 0 : HEAD_DIM] | ||
| kv_flat[kv_row : kv_row + 1, :] = kv_row_fp32 | ||
| cmp_kv_cache_flat[cache_row : cache_row + 1, :] = pl.cast(kv_row_fp32, target_type=pl.BF16, mode="rint") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect producer/consumer ordering around ratio-4 decode finalization.
rg -n -C6 'softmax_pool|compressor_rmsnorm_rope|kv_and_cache_write' models/deepseek/v4/compressor_ratio4.pyRepository: hw-native-sys/pypto-lib
Length of output: 3248
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- files mentioning pl.spmd with deps ---'
rg -n 'pl\.spmd\(' models -g '*.py' | head -n 200
echo
echo '--- search for DSL docs / definitions ---'
rg -n 'class .*spmd|def spmd|deps=' . -g '*.py' -g '*.md' | head -n 200Repository: hw-native-sys/pypto-lib
Length of output: 27042
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the surrounding implementation around the suspicious decode path.
sed -n '250,350p' models/deepseek/v4/compressor_ratio4.py
echo
echo '--- similar pattern in the prefill path ---'
sed -n '720,790p' models/deepseek/v4/compressor_ratio4.pyRepository: hw-native-sys/pypto-lib
Length of output: 8417
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- compressor_rmsnorm_rope definitions ---'
rg -n -C5 'def compressor_rmsnorm_rope' models/deepseek/v4
echo
echo '--- exact decode-indexer compressor sequencing for comparison ---'
sed -n '180,290p' models/deepseek/v4/decode_indexer_compressor.pyRepository: hw-native-sys/pypto-lib
Length of output: 8725
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- ratio-128 rmsnorm/rope and write chain ---'
sed -n '220,390p' models/deepseek/v4/compressor_ratio128.pyRepository: hw-native-sys/pypto-lib
Length of output: 9911
Add an explicit dependency to kv_and_cache_write. compressor_rmsnorm_rope(...) is not enough to order the final SPMD write; wire this block to the RMS/RoPE producer with deps=[...] so normed_kv can’t be consumed early.
🤖 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/compressor_ratio4.py` around lines 320 - 335, The
kv_and_cache_write SPMD block is missing an explicit dependency on the RMS/RoPE
producer, so the final write can start before normed_kv is ready. Update the
with pl.spmd(..., name_hint="kv_and_cache_write") block to include a deps
argument that depends on the compressor_rmsnorm_rope-producing step, using the
existing compressor_rmsnorm_rope and normed_kv flow as the reference point. Keep
the rest of the write logic unchanged; only wire the scheduling dependency so
normed_kv is not consumed early.
| position_ids: pl.Tensor[[T], pl.INT32], | ||
| idx_slot_mapping: pl.Tensor[[T], pl.INT64], | ||
| inner_state_slot_mapping: pl.Tensor[[T], pl.INT64], |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify no stale 2-D reads remain for flattened position_ids.
rg -nP 'pl\.read\(position_ids,\s*\[[^]]*,[^]]*\]\)' models/deepseek/v4/decode_indexer.py models/deepseek/v4/decode_indexer_compressor.pyRepository: hw-native-sys/pypto-lib
Length of output: 349
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== position_ids declaration and nearby readers =="
nl -ba models/deepseek/v4/decode_indexer.py | sed -n '110,330p'
echo
echo "== all position_ids reads in relevant files =="
rg -n 'position_ids' models/deepseek/v4/decode_indexer.py models/deepseek/v4/decode_indexer_compressor.pyRepository: hw-native-sys/pypto-lib
Length of output: 251
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== decode_indexer.py: position_ids declaration and reads =="
sed -n '115,320p' models/deepseek/v4/decode_indexer.py | cat -n
echo
echo "== decode_indexer_compressor.py: any position_ids usage =="
sed -n '1,260p' models/deepseek/v4/decode_indexer_compressor.py | cat -n | rg -n 'position_ids|pl\.read|pl\.Tensor'Repository: hw-native-sys/pypto-lib
Length of output: 16761
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== decode_indexer.py: target range =="
sed -n '115,320p' models/deepseek/v4/decode_indexer.py | cat -n
echo
echo "== decode_indexer.py: position_ids reads =="
rg -n 'position_ids|pl\.read\(position_ids' models/deepseek/v4/decode_indexer.py
echo
echo "== decode_indexer_compressor.py: position_ids usage =="
rg -n 'position_ids|pl\.read\(position_ids' models/deepseek/v4/decode_indexer_compressor.pyRepository: hw-native-sys/pypto-lib
Length of output: 17019
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '248,312p' models/deepseek/v4/decode_indexer.py | cat -nRepository: hw-native-sys/pypto-lib
Length of output: 4453
Finish flattening the position_ids reads. position_ids is pl.Tensor[[T], pl.INT32], but models/deepseek/v4/decode_indexer.py:264 and :305 still read it as [b, s] / [batch_idx, token_s]. Use the flattened token index already in scope so these paths match the tensor rank and keep the visibility bound calculation correct.
🤖 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/decode_indexer.py` around lines 121 - 123, `position_ids`
is already flattened to a 1D token tensor, but the remaining reads in
`decode_indexer.py` still treat it as 2D with batch/token indices. Update the
code paths around the `position_ids` accesses in the decode indexing logic to
use the flattened token index already available in scope instead of `[b, s]` or
`[batch_idx, token_s]`, and keep the visibility-bound computation aligned with
the 1D rank.
| ori_block_table_req = pl.slice(ori_block_table, [B, ORI_TABLE_BLOCKS], [ridx * B, 0]) | ||
| cmp_kv_req = pl.slice(cmp_kv, [CMP_CACHE_BLOCKS, BLOCK_SIZE, 1, HEAD_DIM], | ||
| [ridx * CMP_CACHE_BLOCKS, 0, 0, 0]) | ||
| cmp_block_table_req = pl.slice(cmp_block_table, [CMP_TABLE_BLOCKS], [ridx * CMP_TABLE_BLOCKS]) | ||
| cmp_block_table_req = pl.slice(cmp_block_table, [B, CMP_TABLE_BLOCKS], [ridx * B, 0]) | ||
| idx_kv_cache_req = pl.slice(idx_kv_cache, [IDX_CACHE_BLOCKS, BLOCK_SIZE, 1, IDX_HEAD_DIM], | ||
| [ridx * IDX_CACHE_BLOCKS, 0, 0, 0]) | ||
| idx_kv_scale_req = pl.slice(idx_kv_scale, [IDX_CACHE_BLOCKS, BLOCK_SIZE, 1, 1], | ||
| [ridx * IDX_CACHE_BLOCKS, 0, 0, 0]) | ||
| idx_block_table_req = pl.slice(idx_block_table, [IDX_TABLE_BLOCKS], [ridx * IDX_TABLE_BLOCKS]) | ||
| hca_kv_state_req = pl.slice(hca_cmp_kv_state, [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, HCA_MAIN_OUT_DIM], | ||
| [ridx * HCA_STATE_BLOCK_NUM, 0, 0]) | ||
| hca_score_state_req = pl.slice(hca_cmp_score_state, [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, HCA_MAIN_OUT_DIM], | ||
| [ridx * HCA_STATE_BLOCK_NUM, 0, 0]) | ||
| hca_state_table_req = pl.slice(hca_compress_state_block_table, [HCA_STATE_MAX_BLOCKS], | ||
| [ridx * HCA_STATE_MAX_BLOCKS]) | ||
| csa_kv_state_req = pl.slice(csa_cmp_kv_state, [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, CSA_MAIN_OUT_DIM], | ||
| [ridx * CSA_STATE_BLOCK_NUM, 0, 0]) | ||
| csa_score_state_req = pl.slice(csa_cmp_score_state, [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, CSA_MAIN_OUT_DIM], | ||
| [ridx * CSA_STATE_BLOCK_NUM, 0, 0]) | ||
| csa_state_table_req = pl.slice(csa_compress_state_block_table, [CSA_STATE_MAX_BLOCKS], | ||
| [ridx * CSA_STATE_MAX_BLOCKS]) | ||
| csa_inner_kv_state_req = pl.slice(csa_inner_kv_state, [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_OUT_DIM], | ||
| [ridx * INNER_STATE_BLOCK_NUM, 0, 0]) | ||
| csa_inner_score_state_req = pl.slice(csa_inner_score_state, [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_OUT_DIM], | ||
| [ridx * INNER_STATE_BLOCK_NUM, 0, 0]) | ||
| csa_inner_state_table_req = pl.slice(csa_inner_compress_state_block_table, [INNER_STATE_MAX_BLOCKS], | ||
| [ridx * INNER_STATE_MAX_BLOCKS]) | ||
| idx_block_table_req = pl.slice(idx_block_table, [B, IDX_TABLE_BLOCKS], [ridx * B, 0]) | ||
| hca_compress_state_req = pl.slice(hca_compress_state, [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, HCA_MAIN_STATE_DIM], | ||
| [ridx * HCA_STATE_BLOCK_NUM, 0, 0]) | ||
| hca_state_table_req = pl.slice(hca_compress_state_block_table, [B, HCA_STATE_MAX_BLOCKS], [ridx * B, 0]) | ||
| csa_compress_state_req = pl.slice(csa_compress_state, [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, CSA_MAIN_STATE_DIM], | ||
| [ridx * CSA_STATE_BLOCK_NUM, 0, 0]) | ||
| csa_state_table_req = pl.slice(csa_compress_state_block_table, [B, CSA_STATE_MAX_BLOCKS], [ridx * B, 0]) | ||
| csa_inner_compress_state_req = pl.slice(csa_inner_compress_state, [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_STATE_DIM], | ||
| [ridx * INNER_STATE_BLOCK_NUM, 0, 0]) | ||
| csa_inner_state_table_req = pl.slice( | ||
| csa_inner_compress_state_block_table, | ||
| [B, INNER_STATE_MAX_BLOCKS], | ||
| [ridx * B, 0], | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm whether B / PREFILL_BATCH is imported or defined in prefill_layer.py
fd -t f 'prefill_layer.py' models/deepseek/v4 --exec sh -c '
echo "== $1 =="
rg -nP "^\s*(from|import).*\bPREFILL_BATCH\b" "$1" || echo "no PREFILL_BATCH import"
rg -nP "^\s*B\s*=" "$1" || echo "no top-level B assignment"
rg -nP "PREFILL_BATCH\s+as\s+B" "$1" || echo "no PREFILL_BATCH as B alias"
' sh {}Repository: hw-native-sys/pypto-lib
Length of output: 281
🏁 Script executed:
#!/bin/bash
sed -n '1,120p' models/deepseek/v4/prefill_layer.py
printf '\n----\n'
sed -n '700,760p' models/deepseek/v4/prefill_layer.py
printf '\n----\n'
rg -n "^\s*(from|import)\b|^\s*B\s*=|PREFILL_BATCH|\\bB\\b" models/deepseek/v4/prefill_layer.pyRepository: hw-native-sys/pypto-lib
Length of output: 8320
Import B before using it here
prefill_layer.py uses B in the request loop and _req_block_count, but this module never defines or imports it. That will raise NameError on this path; mirror prefill_fwd.py and import PREFILL_BATCH as B.
🧰 Tools
🪛 Ruff (0.15.20)
[error] 308-308: Undefined name B
(F821)
[error] 308-308: Undefined name B
(F821)
[error] 311-311: Undefined name B
(F821)
[error] 311-311: Undefined name B
(F821)
[error] 316-316: Undefined name B
(F821)
[error] 316-316: Undefined name B
(F821)
[error] 319-319: Undefined name B
(F821)
[error] 319-319: Undefined name B
(F821)
[error] 322-322: Undefined name B
(F821)
[error] 322-322: Undefined name B
(F821)
[error] 327-327: Undefined name B
(F821)
[error] 328-328: Undefined name B
(F821)
🤖 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/prefill_layer.py` around lines 308 - 329, The
request-building code in prefill_layer.py references B in the slice calls and
block-count logic, but B is not defined in this module. Fix this by importing
the same batch-size alias used in prefill_fwd.py, specifically PREFILL_BATCH as
B, so the symbols in the request loop and _req_block_count resolve correctly.
Source: Linters/SAST tools
90f56d4 to
dad3bd5
Compare
Consolidate the four compressor entry files (decode/prefill x ratio-4/128) into compressor_ratio4.py and compressor_ratio128.py with shared per-ratio cores, add compressor_schedule.py for the shared write scheduling, and move the compressor/indexer callers to token-major [T] contracts. Reconcile the decode/prefill callers with the upstream CSA task-count refactor (hw-native-sys#734) and the MoE dispatch split (hw-native-sys#729).
…io-4 compressor Rewrite decode_compressor_ratio4 to mirror main's hw-native-sys#734 decode: fuse state-scatter + softmax-pool into one task and rmsnorm + rope + cache-write into another, over a single 16-row block (pad_idx = c_idx identity, requires DECODE_B <= 16). Reads stay token-major and rope rows are gathered in-kernel from the shared freqs tables. Decode no longer shares compressor_core_ratio4 -- a single paged block and packed multi-write prefill cannot reuse the fused path -- so the core stays prefill-only; ratio-128 keeps its full shared core. Bit-exact on a2a3 (max_error_ratio=0.0) for both ratios and both modes; ratio-4 decode collapses to two fused tasks (~77.5 us, from ~87 us).
f65a463 to
ec06b4d
Compare
1 participant
Summary
compressor_ratio4.pyowns ratio-4 decode/prefill, their shared projection and pool helpers, golden references, tensor specs, and CLI mode switch.compressor_ratio128.pyowns ratio-128 decode/prefill, their shared projection and pool pipeline, golden references, tensor specs, and CLI mode switch.compressor_common.pyowns ratio-agnostic write scheduling, RoPE-row gathering, RMSNorm/RoPE computation, and cache finalization. The intermediatecompressor_schedule.pymodule is removed.compressor_core_ratio4andcompressor_core_ratio128returnnormed_kv; wrappers own scheduled cache finalization throughfinalize_compressor_writes.[T]position/slot metadata across compressor and indexer entrypoints. Compressor RoPE rows are gathered in-kernel from the shared frequency tables.compress_state[..., 2 * OUT_DIM]buffers and batch-major[B, max_blocks]compressor/indexer state tables.Latest-main conflict resolution
mainate252e07a(through perf(qwen3-14b prefill): fuse qkpv with skewed 4-token batching #746). The updated PR head is511a82a2.prefill_attention_csa.py,prefill_attention_hca.py,prefill_attention_swa.py,prefill_layer.py, andprefill_sparse_attn.pyagainst FIX: Update prefill sparse attention cache contract #739's cache-first prefill sparse-attention contract.swa_indicesaddress physical original-cache rows.cmp_indicesaddress logical compressed-cache slots.[B, ...]tensors.cmp_sparse_indices/cmp_sparse_lensinterface.A2/A3 gate fix
prefill_layer.pyandprefill_fwd.pybecause the merge mixed two block-table shape contracts.prefill_layerpacks each request's one-dimensional original/compressed sparse-cache tables contiguously. Its golden request slicer incorrectly usedCHILD_BATCH == 1, so it retained one entry instead of all 128/32 entries.prefill_fwddeclared original/compressed tables with an extra batch dimension even though their shared specs are[rank, blocks]. Its indexer table correctly remains[rank, batch, blocks], but its initializer returned[rank, blocks]until it was expanded to the declared tail shape.511a82a2restores those boundaries without changing compressor math or scheduling.Validation
git diff --check.pytest tests/golden -q: 159 passed.prefill_sparse_attn, SWA/HCA/CSA attention, fullprefill_layer, and fullprefill_fwd.git merge-tree --write-tree origin/main HEADsucceeds without conflicts.kv_cacheandx_outvalidations.prefill_layer.pypasses on two A2/A3 devices:kv_cachepassesratio_allclose(atol=1e-4, rtol=1/128, max_error_ratio=0.005)andx_nextpassesvalid_ratio_reldiff.prefill_fwd.pycompletes the full two-device runtime with exit code 0.Reproduce
Review notes
compressor_core_ratio4. Refactor: cut runtime task count in DeepSeek-V4 decode CSA flow #734's fused 16-row block path and prefill's packed multi-write schedule have different lowering and scheduling requirements.Related Issues
N/A